From ae392a28fecb4bcc25c1a98672d963c6ac3c19d5 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 11:47:19 +0800 Subject: [PATCH 01/95] platform: add multi-tenant contracts --- platform/audit.go | 51 ++++++ platform/doc.go | 15 ++ platform/errors.go | 42 +++++ platform/idempotency.go | 127 +++++++++++++ platform/identity.go | 114 ++++++++++++ platform/redaction.go | 96 ++++++++++ platform/types.go | 395 ++++++++++++++++++++++++++++++++++++++++ platform/types_test.go | 269 +++++++++++++++++++++++++++ platform/validation.go | 212 +++++++++++++++++++++ 9 files changed, 1321 insertions(+) create mode 100644 platform/audit.go create mode 100644 platform/doc.go create mode 100644 platform/errors.go create mode 100644 platform/idempotency.go create mode 100644 platform/identity.go create mode 100644 platform/redaction.go create mode 100644 platform/types.go create mode 100644 platform/types_test.go create mode 100644 platform/validation.go diff --git a/platform/audit.go b/platform/audit.go new file mode 100644 index 0000000000..deba348c82 --- /dev/null +++ b/platform/audit.go @@ -0,0 +1,51 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "context" + "sync" +) + +// AuditSink stores audit records. +type AuditSink interface { + // WriteAudit writes one audit record. + WriteAudit(ctx context.Context, record AuditRecord) error +} + +// InMemoryAuditSink is a concurrency-safe audit sink for tests and demos. +type InMemoryAuditSink struct { + mu sync.Mutex + records []AuditRecord +} + +// NewInMemoryAuditSink creates an in-memory audit sink. +func NewInMemoryAuditSink() *InMemoryAuditSink { + return &InMemoryAuditSink{} +} + +// WriteAudit writes one audit record. +func (s *InMemoryAuditSink) WriteAudit(ctx context.Context, record AuditRecord) error { + if err := ctx.Err(); err != nil { + return err + } + s.mu.Lock() + defer s.mu.Unlock() + s.records = append(s.records, record) + return nil +} + +// Records returns a snapshot of written audit records. +func (s *InMemoryAuditSink) Records() []AuditRecord { + s.mu.Lock() + defer s.mu.Unlock() + out := make([]AuditRecord, len(s.records)) + copy(out, s.records) + return out +} diff --git a/platform/doc.go b/platform/doc.go new file mode 100644 index 0000000000..e63e59b18c --- /dev/null +++ b/platform/doc.go @@ -0,0 +1,15 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +// Package platform contains reusable multi-tenant platform contracts. +// +// The package is intentionally small and dependency-light: it models tenant, +// channel, governance, idempotency, and audit data that gateway or channel +// adapter implementations can share without changing the core runner/session +// interfaces. +package platform diff --git a/platform/errors.go b/platform/errors.go new file mode 100644 index 0000000000..6778ef5533 --- /dev/null +++ b/platform/errors.go @@ -0,0 +1,42 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import "errors" + +var ( + // ErrTenantIDRequired indicates a missing tenant identifier. + ErrTenantIDRequired = errors.New("tenant_id is required") + // ErrAppIDRequired indicates a missing app identifier. + ErrAppIDRequired = errors.New("app_id is required") + // ErrBindingIDRequired indicates a missing channel binding identifier. + ErrBindingIDRequired = errors.New("binding_id is required") + // ErrChannelRequired indicates a missing channel identifier. + ErrChannelRequired = errors.New("channel is required") + // ErrAccountIDRequired indicates a missing channel account identifier. + ErrAccountIDRequired = errors.New("account_id is required") + // ErrPlatformMessageIDRequired indicates a missing platform message identifier. + ErrPlatformMessageIDRequired = errors.New("platform_message_id is required") + // ErrIdempotencyRecordNotFound indicates an unknown idempotency key. + ErrIdempotencyRecordNotFound = errors.New("idempotency record not found") + // ErrExternalUserIDRequired indicates a missing external user identifier. + ErrExternalUserIDRequired = errors.New("external_user_id is required") + // ErrExternalGroupIDRequired indicates a missing group identifier. + ErrExternalGroupIDRequired = errors.New("external_group_id is required") + // ErrConversationTypeRequired indicates a missing conversation type. + ErrConversationTypeRequired = errors.New("conversation_type is required") + // ErrInvalidConversationType indicates an unsupported conversation type. + ErrInvalidConversationType = errors.New("invalid conversation_type") + // ErrSecretReferenceRequired indicates a configuration contains inline secret material. + ErrSecretReferenceRequired = errors.New("secret reference is required") + // ErrInlineSecretRejected indicates a configuration appears to contain inline secret material. + ErrInlineSecretRejected = errors.New("inline secret values are not allowed") + // ErrWebhookPathRequired indicates a missing webhook path. + ErrWebhookPathRequired = errors.New("webhook_path is required") +) diff --git a/platform/idempotency.go b/platform/idempotency.go new file mode 100644 index 0000000000..8e9333ae58 --- /dev/null +++ b/platform/idempotency.go @@ -0,0 +1,127 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "context" + "sync" + "time" +) + +// IdempotencyStore stores inbound message processing state. +type IdempotencyStore interface { + // Start records a message as processing if it has not been seen. + Start(ctx context.Context, record IdempotencyRecord) (IdempotencyRecord, bool, error) + // Complete marks a processing record as completed. + Complete(ctx context.Context, key string, resultRef string) (IdempotencyRecord, error) + // MarkReplyFailed marks a completed record as needing outbound retry. + MarkReplyFailed(ctx context.Context, key string, resultRef string) (IdempotencyRecord, error) + // Get returns the record for key. + Get(ctx context.Context, key string) (IdempotencyRecord, bool, error) +} + +// InMemoryIdempotencyStore is a concurrency-safe idempotency store for tests and demos. +type InMemoryIdempotencyStore struct { + now func() time.Time + mu sync.Mutex + records map[string]IdempotencyRecord +} + +// NewInMemoryIdempotencyStore creates an in-memory idempotency store. +func NewInMemoryIdempotencyStore() *InMemoryIdempotencyStore { + return &InMemoryIdempotencyStore{ + now: time.Now, + records: make(map[string]IdempotencyRecord), + } +} + +// Start records a message as processing if it has not been seen. +func (s *InMemoryIdempotencyStore) Start( + ctx context.Context, + record IdempotencyRecord, +) (IdempotencyRecord, bool, error) { + if err := ctx.Err(); err != nil { + return IdempotencyRecord{}, false, err + } + key := record.IdempotencyKey + if key == "" { + key = IdempotencyKey( + record.TenantID, + record.Channel, + record.AccountID, + record.PlatformMessageID, + ) + record.IdempotencyKey = key + } + s.mu.Lock() + defer s.mu.Unlock() + if existing, ok := s.records[key]; ok { + return existing, false, nil + } + now := s.now() + record.Status = IdempotencyStatusProcessing + record.FirstSeenAt = now + record.UpdatedAt = now + s.records[key] = record + return record, true, nil +} + +// Complete marks a processing record as completed. +func (s *InMemoryIdempotencyStore) Complete( + ctx context.Context, + key string, + resultRef string, +) (IdempotencyRecord, error) { + return s.update(ctx, key, IdempotencyStatusCompleted, resultRef) +} + +// MarkReplyFailed marks a completed record as needing outbound retry. +func (s *InMemoryIdempotencyStore) MarkReplyFailed( + ctx context.Context, + key string, + resultRef string, +) (IdempotencyRecord, error) { + return s.update(ctx, key, IdempotencyStatusReplyFailed, resultRef) +} + +// Get returns the record for key. +func (s *InMemoryIdempotencyStore) Get( + ctx context.Context, + key string, +) (IdempotencyRecord, bool, error) { + if err := ctx.Err(); err != nil { + return IdempotencyRecord{}, false, err + } + s.mu.Lock() + defer s.mu.Unlock() + record, ok := s.records[key] + return record, ok, nil +} + +func (s *InMemoryIdempotencyStore) update( + ctx context.Context, + key string, + status IdempotencyStatus, + resultRef string, +) (IdempotencyRecord, error) { + if err := ctx.Err(); err != nil { + return IdempotencyRecord{}, err + } + s.mu.Lock() + defer s.mu.Unlock() + record, ok := s.records[key] + if !ok { + return IdempotencyRecord{}, ErrIdempotencyRecordNotFound + } + record.Status = status + record.ResultRef = resultRef + record.UpdatedAt = s.now() + s.records[key] = record + return record, nil +} diff --git a/platform/identity.go b/platform/identity.go new file mode 100644 index 0000000000..aab8af1728 --- /dev/null +++ b/platform/identity.go @@ -0,0 +1,114 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "crypto/sha256" + "encoding/hex" + "fmt" + "net/url" + "strings" +) + +// InternalUserID returns a stable tenant-scoped user identifier. +func InternalUserID(tenantID, channel, externalUserID string) string { + return "usr_" + shortHash(tenantID, channel, externalUserID) +} + +// UserIDHash returns a low-sensitivity hash for logs and trace attributes. +func UserIDHash(tenantID, channel, userID string) string { + return "user_hash_" + shortHash(tenantID, channel, userID) +} + +// IdempotencyKey returns the canonical duplicate-delivery key. +func IdempotencyKey(tenantID, channel, accountID, platformMessageID string) string { + return strings.Join([]string{ + "tenant", escapeKeyPart(tenantID), + "channel", escapeKeyPart(channel), + "account", escapeKeyPart(accountID), + "message", escapeKeyPart(platformMessageID), + }, ":") +} + +// SessionIDForInbound returns the stable session id for one inbound message. +func SessionIDForInbound(msg InboundMessage) (string, error) { + if err := msg.Validate(); err != nil { + return "", err + } + return SessionID( + msg.TenantID, + msg.AppID, + msg.Channel, + msg.ConversationType, + msg.ExternalUserID, + msg.ExternalGroupID, + msg.ThreadID, + ) +} + +// SessionID returns the stable tenant/app/channel-scoped session id. +func SessionID( + tenantID string, + appID string, + channel string, + conversationType ConversationType, + externalUserID string, + externalGroupID string, + threadID string, +) (string, error) { + if strings.TrimSpace(tenantID) == "" { + return "", ErrTenantIDRequired + } + if strings.TrimSpace(appID) == "" { + return "", ErrAppIDRequired + } + if strings.TrimSpace(channel) == "" { + return "", ErrChannelRequired + } + prefix := fmt.Sprintf( + "tenant:%s:app:%s:channel:%s", + escapeKeyPart(tenantID), + escapeKeyPart(appID), + escapeKeyPart(channel), + ) + switch conversationType { + case ConversationTypeDM: + if strings.TrimSpace(externalUserID) == "" { + return "", ErrExternalUserIDRequired + } + return prefix + ":dm:" + escapeKeyPart(externalUserID), nil + case ConversationTypeGroup: + if strings.TrimSpace(externalGroupID) == "" { + return "", ErrExternalGroupIDRequired + } + return prefix + ":group:" + escapeKeyPart(externalGroupID), nil + case ConversationTypeThread: + if strings.TrimSpace(externalGroupID) == "" { + return "", ErrExternalGroupIDRequired + } + if strings.TrimSpace(threadID) == "" { + return "", fmt.Errorf("thread_id is required") + } + return prefix + ":group:" + escapeKeyPart(externalGroupID) + + ":thread:" + escapeKeyPart(threadID), nil + case "": + return "", ErrConversationTypeRequired + default: + return "", ErrInvalidConversationType + } +} + +func shortHash(parts ...string) string { + sum := sha256.Sum256([]byte(strings.Join(parts, "\x00"))) + return hex.EncodeToString(sum[:])[:24] +} + +func escapeKeyPart(value string) string { + return url.PathEscape(strings.TrimSpace(value)) +} diff --git a/platform/redaction.go b/platform/redaction.go new file mode 100644 index 0000000000..ccbc327017 --- /dev/null +++ b/platform/redaction.go @@ -0,0 +1,96 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "regexp" + "strings" +) + +var defaultRedactionPatterns = []*regexp.Regexp{ + regexp.MustCompile(`(?i)(Authorization:\s*Basic\s+)[A-Za-z0-9._~+/\-]+=*`), + regexp.MustCompile(`(?i)(Bearer\s+)[A-Za-z0-9._~+/\-]+=*`), + regexp.MustCompile(`(?i)(api[_-]?key|token|secret|password|passwd|authorization|cookie)=([^&\s]+)`), + regexp.MustCompile(`(?i)(api[_-]?key|token|secret|password|passwd|authorization|cookie):\s*([^,\s]+)`), + regexp.MustCompile(`(?i)("(?:api[_-]?key|token|secret|password|passwd|authorization|cookie)"\s*:\s*")([^"]+)(")`), + regexp.MustCompile(`(?i)(sk-[A-Za-z0-9._~+/\-]{8,})`), + regexp.MustCompile(`(?i)(://[^:\s/]+:)([^@\s]+)(@)`), + regexp.MustCompile(`(?s)-----BEGIN [A-Z ]*PRIVATE KEY-----.*?-----END [A-Z ]*PRIVATE KEY-----`), +} + +// Redactor masks sensitive values before logging, tracing, or auditing. +type Redactor struct { + patterns []*regexp.Regexp +} + +// NewRedactor returns a redactor with default secret patterns and optional extras. +func NewRedactor(extraPatterns ...string) (*Redactor, error) { + patterns := append([]*regexp.Regexp(nil), defaultRedactionPatterns...) + for _, pattern := range extraPatterns { + pattern = strings.TrimSpace(pattern) + if pattern == "" { + continue + } + compiled, err := regexp.Compile(pattern) + if err != nil { + return nil, err + } + patterns = append(patterns, compiled) + } + return &Redactor{patterns: patterns}, nil +} + +// Redact returns text with known sensitive values masked. +func (r *Redactor) Redact(text string) string { + if r == nil { + r, _ = NewRedactor() + } + redacted := text + for _, pattern := range r.patterns { + redacted = pattern.ReplaceAllStringFunc(redacted, redactMatch) + } + return redacted +} + +func redactMatch(match string) string { + lower := strings.ToLower(match) + if strings.Contains(lower, "authorization:") && strings.Contains(lower, "basic ") { + return match[:strings.Index(lower, "basic ")+6] + "****" + } + if strings.Contains(lower, "bearer ") { + return match[:strings.Index(lower, "bearer ")+7] + "****" + } + if strings.Contains(match, "://") && strings.Contains(match, "@") { + start := strings.Index(match, "://") + at := strings.LastIndex(match, "@") + credential := match[start+3 : at] + colon := strings.LastIndex(credential, ":") + if colon >= 0 { + return match[:start+3+colon+1] + "****" + match[at:] + } + } + if strings.HasPrefix(match, "-----BEGIN ") { + return "-----BEGIN PRIVATE KEY-----****-----END PRIVATE KEY-----" + } + if idx := strings.Index(match, "="); idx >= 0 { + return match[:idx+1] + "****" + } + if idx := strings.Index(match, ":"); idx >= 0 { + prefix := match[:idx+1] + rest := match[idx+1:] + if strings.HasPrefix(strings.TrimLeft(rest, " \t"), "\"") && strings.HasSuffix(match, "\"") { + return prefix + " \"****\"" + } + return prefix + " ****" + } + if len(match) <= 8 { + return "****" + } + return match[:4] + "****" + match[len(match)-4:] +} diff --git a/platform/types.go b/platform/types.go new file mode 100644 index 0000000000..ca0cfba02f --- /dev/null +++ b/platform/types.go @@ -0,0 +1,395 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import "time" + +// TenantStatus is the lifecycle state of a tenant. +type TenantStatus string + +const ( + // TenantStatusActive allows normal request processing. + TenantStatusActive TenantStatus = "active" + // TenantStatusSuspended rejects new runtime requests while retaining data. + TenantStatusSuspended TenantStatus = "suspended" + // TenantStatusDeleted marks a tenant as soft-deleted. + TenantStatusDeleted TenantStatus = "deleted" +) + +// AppStatus is the lifecycle state of an agent application. +type AppStatus string + +const ( + // AppStatusActive allows the app to receive runtime traffic. + AppStatusActive AppStatus = "active" + // AppStatusSuspended rejects runtime traffic for the app. + AppStatusSuspended AppStatus = "suspended" + // AppStatusDeleted marks the app as soft-deleted. + AppStatusDeleted AppStatus = "deleted" +) + +// BindingStatus is the lifecycle state of a channel binding. +type BindingStatus string + +const ( + // BindingStatusActive allows inbound callbacks through the binding. + BindingStatusActive BindingStatus = "active" + // BindingStatusDisabled rejects inbound callbacks through the binding. + BindingStatusDisabled BindingStatus = "disabled" + // BindingStatusDeleted marks the binding as soft-deleted. + BindingStatusDeleted BindingStatus = "deleted" +) + +// ConversationType describes the IM conversation scope. +type ConversationType string + +const ( + // ConversationTypeDM is a one-to-one conversation. + ConversationTypeDM ConversationType = "dm" + // ConversationTypeGroup is a group conversation. + ConversationTypeGroup ConversationType = "group" + // ConversationTypeThread is a thread or topic inside a group conversation. + ConversationTypeThread ConversationType = "thread" +) + +// MessageType describes the normalized inbound message kind. +type MessageType string + +const ( + // MessageTypeText is a plain text message. + MessageTypeText MessageType = "text" + // MessageTypeImage is an image message. + MessageTypeImage MessageType = "image" + // MessageTypeFile is a file message. + MessageTypeFile MessageType = "file" + // MessageTypeAudio is an audio or voice message. + MessageTypeAudio MessageType = "audio" + // MessageTypeVideo is a video message. + MessageTypeVideo MessageType = "video" + // MessageTypeEvent is a non-conversational platform event. + MessageTypeEvent MessageType = "event" + // MessageTypeUnknown is an unsupported or unknown message type. + MessageTypeUnknown MessageType = "unknown" +) + +// ContentPartType describes one normalized content part. +type ContentPartType string + +const ( + // ContentPartTypeText carries text content. + ContentPartTypeText ContentPartType = "text" + // ContentPartTypeImage carries an image artifact reference. + ContentPartTypeImage ContentPartType = "image" + // ContentPartTypeFile carries a file artifact reference. + ContentPartTypeFile ContentPartType = "file" + // ContentPartTypeAudio carries an audio artifact reference. + ContentPartTypeAudio ContentPartType = "audio" + // ContentPartTypeVideo carries a video artifact reference. + ContentPartTypeVideo ContentPartType = "video" + // ContentPartTypeLocation carries a location payload. + ContentPartTypeLocation ContentPartType = "location" + // ContentPartTypeUnknown carries unsupported content metadata. + ContentPartTypeUnknown ContentPartType = "unknown" +) + +// OutboundMessageKind describes the kind of payload sent to an IM platform. +type OutboundMessageKind string + +const ( + // OutboundMessageKindText sends plain text. + OutboundMessageKindText OutboundMessageKind = "text" + // OutboundMessageKindMarkdown sends markdown when the channel supports it. + OutboundMessageKindMarkdown OutboundMessageKind = "markdown" + // OutboundMessageKindCard sends a structured card. + OutboundMessageKindCard OutboundMessageKind = "card" + // OutboundMessageKindImage sends an image. + OutboundMessageKindImage OutboundMessageKind = "image" + // OutboundMessageKindFile sends a file. + OutboundMessageKindFile OutboundMessageKind = "file" + // OutboundMessageKindStatus sends an execution status update. + OutboundMessageKindStatus OutboundMessageKind = "status" +) + +// IdempotencyStatus is the state of one inbound platform message. +type IdempotencyStatus string + +const ( + // IdempotencyStatusReceived records that the callback was accepted. + IdempotencyStatusReceived IdempotencyStatus = "received" + // IdempotencyStatusProcessing records that the runner is still executing. + IdempotencyStatusProcessing IdempotencyStatus = "processing" + // IdempotencyStatusCompleted records that the runner finished and must not rerun. + IdempotencyStatusCompleted IdempotencyStatus = "completed" + // IdempotencyStatusReplyFailed records that only outbound delivery failed. + IdempotencyStatusReplyFailed IdempotencyStatus = "reply_failed" + // IdempotencyStatusDeadLetter records an item requiring manual replay. + IdempotencyStatusDeadLetter IdempotencyStatus = "dead_letter" +) + +// OutboundStatus is the delivery state of one outbound IM message. +type OutboundStatus string + +const ( + // OutboundStatusPending is waiting for delivery. + OutboundStatusPending OutboundStatus = "pending" + // OutboundStatusSent was delivered to the platform. + OutboundStatusSent OutboundStatus = "sent" + // OutboundStatusFailed failed and may be retried. + OutboundStatusFailed OutboundStatus = "failed" + // OutboundStatusDeadLetter failed permanently or exhausted retries. + OutboundStatusDeadLetter OutboundStatus = "dead_letter" +) + +// DangerousToolAction is the default action for high-risk tools. +type DangerousToolAction string + +const ( + // DangerousToolActionDeny blocks high-risk tools. + DangerousToolActionDeny DangerousToolAction = "deny" + // DangerousToolActionAsk requires approval before high-risk tools execute. + DangerousToolActionAsk DangerousToolAction = "ask" + // DangerousToolActionAllowWithAudit allows high-risk tools with audit. + DangerousToolActionAllowWithAudit DangerousToolAction = "allow_with_audit" +) + +// Tenant is the top-level isolation boundary. +type Tenant struct { + TenantID string + Name string + Status TenantStatus + Region string + QuotaJSON string + DefaultStorageProfileID string + AuditPolicyID string + CreatedAt time.Time + UpdatedAt time.Time + DeletedAt *time.Time +} + +// AgentApp is a tenant-owned agent application configuration. +type AgentApp struct { + TenantID string + AppID string + AppName string + AgentName string + InstructionRef string + ModelProfileID string + ToolPolicyID string + StorageProfileID string + MemoryProfileID string + ReleaseVersion string + GrayPercent int + Status AppStatus + CreatedAt time.Time + UpdatedAt time.Time +} + +// ModelProfile stores model provider configuration references. +type ModelProfile struct { + TenantID string + ProfileID string + Provider string + Model string + BaseURLRef string + APIKeyRef string + TimeoutMS int + MaxTokens int + Temperature float64 + FallbackProfileID string + CostPolicyJSON string + CreatedAt time.Time + UpdatedAt time.Time +} + +// ToolPolicy stores tenant and app-level tool governance. +type ToolPolicy struct { + TenantID string + PolicyID string + AppID string + ToolWhitelist []string + ToolDenylist []string + DangerousToolAction DangerousToolAction + ApprovalChannel string + ArgumentRedactionRules []string + NetworkPolicyJSON string + FilesystemPolicyJSON string + PlatformDenylist []string + HighRiskTools []string + ToolBudgetRemainingJSON string + CreatedAt time.Time + UpdatedAt time.Time +} + +// ChannelLimits stores configurable channel capability and limit values. +type ChannelLimits struct { + MaxTextLength int + CallbackACKTimeout time.Duration + FileMaxBytes int64 + RateLimitQPS int + Burst int + SupportsAsyncReply bool + SupportsEdit bool + SupportsCardUpdate bool + RetryMaxAttempts int + RetryBackoff string +} + +// ChannelBinding maps one external IM account to one tenant app. +type ChannelBinding struct { + TenantID string + BindingID string + AppID string + Channel string + AccountID string + WebhookPath string + TokenRef string + SecretRef string + AESKeyRef string + AllowedUsers []string + AllowedGroups []string + RequiredMention bool + Status BindingStatus + ChannelLimits ChannelLimits + CreatedAt time.Time + UpdatedAt time.Time +} + +// StorageProfile stores backend choices for a tenant app. +type StorageProfile struct { + TenantID string + ProfileID string + SessionBackend string + MemoryBackend string + SummaryBackend string + ArtifactBackend string + KnowledgeBackend string + AuditBackend string + DSNRef string + Namespace string + TTLJSON string + MigrationMode string + CreatedAt time.Time + UpdatedAt time.Time +} + +// AuditPolicy stores audit retention and export settings. +type AuditPolicy struct { + TenantID string + PolicyID string + RetentionDays int + SampleRate float64 + FullAuditForRiskyTool bool + RedactionRules []string + ExportSink string + ComplianceLevel string +} + +// IMUserMapping maps an external IM identity to an internal identity. +type IMUserMapping struct { + TenantID string + Channel string + ExternalUserID string + InternalUserID string + DisplayName string + Roles []string + Status string + CreatedAt time.Time + UpdatedAt time.Time +} + +// ContentPart is one normalized part of an inbound message. +type ContentPart struct { + Type ContentPartType + Text string + FileRef string + MIMEType string + SizeBytes int64 + SHA256 string + MetadataJSON string +} + +// InboundMessage is the normalized message consumed by a gateway. +type InboundMessage struct { + TenantID string + AppID string + BindingID string + Channel string + ChannelAccountID string + PlatformMessageID string + ExternalUserID string + ExternalGroupID string + ThreadID string + ConversationType ConversationType + MessageType MessageType + ContentParts []ContentPart + RawEventType string + ReceivedAt time.Time + SignatureStatus string + TraceContext map[string]string + RequiredMentionSeen bool +} + +// OutboundMessage is the normalized payload delivered to an IM platform. +type OutboundMessage struct { + TenantID string + BindingID string + Channel string + SessionID string + ReplyToPlatformMessageID string + Kind OutboundMessageKind + Content string + FileRef string + Sequence int + DedupKey string + RetryPolicy string + TraceID string +} + +// IdempotencyRecord stores duplicate delivery state for an inbound message. +type IdempotencyRecord struct { + TenantID string + Channel string + AccountID string + PlatformMessageID string + IdempotencyKey string + RequestID string + SessionID string + Status IdempotencyStatus + FirstSeenAt time.Time + UpdatedAt time.Time + ResultRef string +} + +// AuditRecord stores governance and troubleshooting metadata. +type AuditRecord struct { + TenantID string + AuditID string + AppID string + Channel string + BindingID string + UserID string + InternalUserID string + UserIDHash string + SessionID string + MessageID string + RequestID string + AgentName string + ModelName string + ToolName string + Decision string + DecisionReason string + LatencyMS int64 + ErrorType string + Cost float64 + TokenUsageJSON string + TraceID string + RedactedDetailRef string + RedactionVersion string + CreatedAt time.Time +} diff --git a/platform/types_test.go b/platform/types_test.go new file mode 100644 index 0000000000..79c741b501 --- /dev/null +++ b/platform/types_test.go @@ -0,0 +1,269 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "context" + "errors" + "strings" + "testing" +) + +func TestSessionIDForInboundIsTenantScoped(t *testing.T) { + base := InboundMessage{ + AppID: "support", + Channel: "telegram", + ChannelAccountID: "bot-1", + PlatformMessageID: "msg-1", + ExternalUserID: "same-user", + ConversationType: ConversationTypeDM, + MessageType: MessageTypeText, + } + a := base + a.TenantID = "tenant-a" + b := base + b.TenantID = "tenant-b" + + sessionA, err := SessionIDForInbound(a) + if err != nil { + t.Fatalf("SessionIDForInbound tenant-a: %v", err) + } + sessionB, err := SessionIDForInbound(b) + if err != nil { + t.Fatalf("SessionIDForInbound tenant-b: %v", err) + } + if sessionA == sessionB { + t.Fatalf("sessions should differ across tenants: %q", sessionA) + } + if !strings.Contains(sessionA, "tenant:tenant-a:app:support:channel:telegram:dm:same-user") { + t.Fatalf("unexpected dm session id: %q", sessionA) + } +} + +func TestSessionIDForInboundSupportsGroupAndThread(t *testing.T) { + groupID, err := SessionID("tenant", "app", "wecom", ConversationTypeGroup, "user", "room 1", "") + if err != nil { + t.Fatalf("group session: %v", err) + } + if groupID != "tenant:tenant:app:app:channel:wecom:group:room%201" { + t.Fatalf("unexpected group id: %q", groupID) + } + + threadID, err := SessionID("tenant", "app", "telegram", ConversationTypeThread, "user", "chat", "topic/7") + if err != nil { + t.Fatalf("thread session: %v", err) + } + if threadID != "tenant:tenant:app:app:channel:telegram:group:chat:thread:topic%2F7" { + t.Fatalf("unexpected thread id: %q", threadID) + } +} + +func TestValidateInboundRequiresGroupForGroupConversation(t *testing.T) { + msg := InboundMessage{ + TenantID: "tenant", + AppID: "app", + Channel: "telegram", + ChannelAccountID: "bot", + PlatformMessageID: "msg", + ExternalUserID: "user", + ConversationType: ConversationTypeGroup, + } + if err := msg.Validate(); !errors.Is(err, ErrExternalGroupIDRequired) { + t.Fatalf("expected ErrExternalGroupIDRequired, got %v", err) + } +} + +func TestInternalUserIDIsStableAndTenantScoped(t *testing.T) { + a1 := InternalUserID("tenant-a", "telegram", "42") + a2 := InternalUserID("tenant-a", "telegram", "42") + b := InternalUserID("tenant-b", "telegram", "42") + if a1 != a2 { + t.Fatalf("expected stable id, got %q and %q", a1, a2) + } + if a1 == b { + t.Fatalf("expected tenant scoped ids, got %q", a1) + } +} + +func TestIdempotencyStoreDoesNotRestartCompletedMessage(t *testing.T) { + store := NewInMemoryIdempotencyStore() + record := IdempotencyRecord{ + TenantID: "tenant", + Channel: "telegram", + AccountID: "bot", + PlatformMessageID: "msg", + RequestID: "req-1", + SessionID: "session", + } + first, started, err := store.Start(context.Background(), record) + if err != nil { + t.Fatalf("start first: %v", err) + } + if !started { + t.Fatalf("first start should create the record") + } + completed, err := store.Complete(context.Background(), first.IdempotencyKey, "outbound-1") + if err != nil { + t.Fatalf("complete: %v", err) + } + if completed.Status != IdempotencyStatusCompleted { + t.Fatalf("expected completed status, got %q", completed.Status) + } + + again, started, err := store.Start(context.Background(), record) + if err != nil { + t.Fatalf("start duplicate: %v", err) + } + if started { + t.Fatalf("duplicate message should not start runner again") + } + if again.Status != IdempotencyStatusCompleted || again.ResultRef != "outbound-1" { + t.Fatalf("duplicate should return completed result, got %#v", again) + } +} + +func TestBindingRejectsInlineSecrets(t *testing.T) { + binding := ChannelBinding{ + TenantID: "tenant", + AppID: "app", + BindingID: "binding", + Channel: "telegram", + AccountID: "bot", + WebhookPath: "/channels/telegram/binding/callback", + SecretRef: "secret=plain", + } + if err := binding.Validate(); !errors.Is(err, ErrInlineSecretRejected) { + t.Fatalf("expected inline secret rejection, got %v", err) + } +} + +func TestBindingRejectsOpaqueRawSecret(t *testing.T) { + binding := ChannelBinding{ + TenantID: "tenant", + AppID: "app", + BindingID: "binding", + Channel: "telegram", + AccountID: "bot", + WebhookPath: "/channels/telegram/binding/callback", + SecretRef: "sk-1234567890abcdef", + } + if err := binding.Validate(); !errors.Is(err, ErrInlineSecretRejected) { + t.Fatalf("expected inline secret rejection, got %v", err) + } +} + +func TestBindingRejectsTelegramBotToken(t *testing.T) { + binding := ChannelBinding{ + TenantID: "tenant", + AppID: "app", + BindingID: "binding", + Channel: "telegram", + AccountID: "bot", + WebhookPath: "/channels/telegram/binding/callback", + TokenRef: "123456789:AAExampleRawTelegramToken", + } + if err := binding.Validate(); !errors.Is(err, ErrInlineSecretRejected) { + t.Fatalf("expected inline secret rejection, got %v", err) + } +} + +func TestBindingAllowsURISecretReferences(t *testing.T) { + binding := ChannelBinding{ + TenantID: "tenant", + AppID: "app", + BindingID: "binding", + Channel: "telegram", + AccountID: "bot", + WebhookPath: "/channels/telegram/binding/callback", + SecretRef: "kms://tenant/telegram-bot", + } + if err := binding.Validate(); err != nil { + t.Fatalf("expected URI reference to be accepted, got %v", err) + } +} + +func TestModelProfileRejectsInlineSecrets(t *testing.T) { + profile := ModelProfile{ + TenantID: "tenant", + ProfileID: "model", + APIKeyRef: "sk-1234567890abcdef", + } + if err := profile.Validate(); !errors.Is(err, ErrInlineSecretRejected) { + t.Fatalf("expected inline secret rejection, got %v", err) + } +} + +func TestIdempotencyUpdateUnknownKeyFails(t *testing.T) { + store := NewInMemoryIdempotencyStore() + _, err := store.Complete(context.Background(), "missing", "result") + if !errors.Is(err, ErrIdempotencyRecordNotFound) { + t.Fatalf("expected missing record error, got %v", err) + } + + record := IdempotencyRecord{ + TenantID: "tenant", + Channel: "telegram", + AccountID: "bot", + PlatformMessageID: "missing", + } + _, started, err := store.Start(context.Background(), record) + if err != nil { + t.Fatalf("start after failed complete: %v", err) + } + if !started { + t.Fatalf("failed complete must not poison future starts") + } +} + +func TestRedactorMasksSecrets(t *testing.T) { + redactor, err := NewRedactor() + if err != nil { + t.Fatalf("NewRedactor: %v", err) + } + input := `api_key=sk-1234567890abcdef Authorization=Bearer token-value Authorization: Basic abc123 token: raw "password":"json-secret" db=postgres://u:pass@example/db` + got := redactor.Redact(input) + for _, leaked := range []string{ + "sk-1234567890abcdef", + "token-value", + "abc123", + "raw", + "json-secret", + ":pass@", + } { + if strings.Contains(got, leaked) { + t.Fatalf("redacted output leaked %q: %q", leaked, got) + } + } + if !strings.Contains(got, "api_key=****") { + t.Fatalf("expected api_key mask, got %q", got) + } +} + +func TestAuditSinkStoresSnapshot(t *testing.T) { + sink := NewInMemoryAuditSink() + record := AuditRecord{ + TenantID: "tenant", + AuditID: "audit", + UserID: "internal", + InternalUserID: "usr", + UserIDHash: UserIDHash("tenant", "telegram", "external"), + TraceID: "trace", + } + if err := sink.WriteAudit(context.Background(), record); err != nil { + t.Fatalf("WriteAudit: %v", err) + } + records := sink.Records() + if len(records) != 1 { + t.Fatalf("expected one audit record, got %d", len(records)) + } + records[0].TenantID = "changed" + if sink.Records()[0].TenantID != "tenant" { + t.Fatalf("Records should return a defensive copy") + } +} diff --git a/platform/validation.go b/platform/validation.go new file mode 100644 index 0000000000..7ed5ef6440 --- /dev/null +++ b/platform/validation.go @@ -0,0 +1,212 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "fmt" + "strconv" + "strings" +) + +var rawSecretPrefixes = []string{ + "sk-", + "xoxb-", + "xoxp-", + "ya29.", + "ghp_", + "github_pat_", + "glpat-", +} + +// Validate checks that the tenant can be used as an isolation boundary. +func (t Tenant) Validate() error { + if strings.TrimSpace(t.TenantID) == "" { + return ErrTenantIDRequired + } + switch t.Status { + case "", TenantStatusActive, TenantStatusSuspended, TenantStatusDeleted: + return nil + default: + return fmt.Errorf("invalid tenant status %q", t.Status) + } +} + +// Validate checks that the app has the identifiers required for routing. +func (a AgentApp) Validate() error { + if strings.TrimSpace(a.TenantID) == "" { + return ErrTenantIDRequired + } + if strings.TrimSpace(a.AppID) == "" { + return ErrAppIDRequired + } + if a.GrayPercent < 0 || a.GrayPercent > 100 { + return fmt.Errorf("gray_percent must be between 0 and 100") + } + switch a.Status { + case "", AppStatusActive, AppStatusSuspended, AppStatusDeleted: + return nil + default: + return fmt.Errorf("invalid app status %q", a.Status) + } +} + +// Validate checks that model profile sensitive values are stored by reference. +func (p ModelProfile) Validate() error { + if strings.TrimSpace(p.TenantID) == "" { + return ErrTenantIDRequired + } + if strings.TrimSpace(p.ProfileID) == "" { + return fmt.Errorf("profile_id is required") + } + if err := validateSecretReference("base_url_ref", p.BaseURLRef); err != nil { + return err + } + if err := validateSecretReference("api_key_ref", p.APIKeyRef); err != nil { + return err + } + return nil +} + +// Validate checks that a binding has safe routing and secret references. +func (b ChannelBinding) Validate() error { + if strings.TrimSpace(b.TenantID) == "" { + return ErrTenantIDRequired + } + if strings.TrimSpace(b.AppID) == "" { + return ErrAppIDRequired + } + if strings.TrimSpace(b.BindingID) == "" { + return ErrBindingIDRequired + } + if strings.TrimSpace(b.Channel) == "" { + return ErrChannelRequired + } + if strings.TrimSpace(b.AccountID) == "" { + return ErrAccountIDRequired + } + if strings.TrimSpace(b.WebhookPath) == "" { + return ErrWebhookPathRequired + } + if err := validateSecretReference("token_ref", b.TokenRef); err != nil { + return err + } + if err := validateSecretReference("secret_ref", b.SecretRef); err != nil { + return err + } + if err := validateSecretReference("aes_key_ref", b.AESKeyRef); err != nil { + return err + } + switch b.Status { + case "", BindingStatusActive, BindingStatusDisabled, BindingStatusDeleted: + return nil + default: + return fmt.Errorf("invalid binding status %q", b.Status) + } +} + +// Validate checks that an inbound message has enough identity for routing. +func (m InboundMessage) Validate() error { + if strings.TrimSpace(m.TenantID) == "" { + return ErrTenantIDRequired + } + if strings.TrimSpace(m.AppID) == "" { + return ErrAppIDRequired + } + if strings.TrimSpace(m.Channel) == "" { + return ErrChannelRequired + } + if strings.TrimSpace(m.ChannelAccountID) == "" { + return ErrAccountIDRequired + } + if strings.TrimSpace(m.PlatformMessageID) == "" { + return ErrPlatformMessageIDRequired + } + if strings.TrimSpace(m.ExternalUserID) == "" { + return ErrExternalUserIDRequired + } + switch m.ConversationType { + case ConversationTypeDM: + return nil + case ConversationTypeGroup: + if strings.TrimSpace(m.ExternalGroupID) == "" { + return ErrExternalGroupIDRequired + } + return nil + case ConversationTypeThread: + if strings.TrimSpace(m.ExternalGroupID) == "" { + return ErrExternalGroupIDRequired + } + if strings.TrimSpace(m.ThreadID) == "" { + return fmt.Errorf("thread_id is required") + } + return nil + case "": + return ErrConversationTypeRequired + default: + return ErrInvalidConversationType + } +} + +// Validate checks that a storage profile uses references for sensitive values. +func (p StorageProfile) Validate() error { + if strings.TrimSpace(p.TenantID) == "" { + return ErrTenantIDRequired + } + if strings.TrimSpace(p.ProfileID) == "" { + return fmt.Errorf("profile_id is required") + } + if err := validateSecretReference("dsn_ref", p.DSNRef); err != nil { + return err + } + return nil +} + +func validateSecretReference(field, value string) error { + value = strings.TrimSpace(value) + if value == "" { + return nil + } + if strings.Contains(value, "=") || + hasInlineURLCredential(value) || + looksLikeRawSecret(value) { + return fmt.Errorf("%s: %w", field, ErrInlineSecretRejected) + } + return nil +} + +func hasInlineURLCredential(value string) bool { + scheme := strings.Index(value, "://") + at := strings.Index(value, "@") + if scheme < 0 || at < 0 || at < scheme { + return false + } + credential := value[scheme+3 : at] + return strings.Contains(credential, ":") +} + +func looksLikeRawSecret(value string) bool { + lower := strings.ToLower(value) + for _, prefix := range rawSecretPrefixes { + if strings.HasPrefix(lower, prefix) { + return true + } + } + if strings.HasPrefix(lower, "bot") && strings.Contains(value, ":") { + return true + } + if colon := strings.Index(value, ":"); colon > 0 { + if _, err := strconv.ParseInt(value[:colon], 10, 64); err == nil { + return true + } + } + if len(value) >= 32 && !strings.ContainsAny(value, "/:.") { + return true + } + return false +} From 0cddda9dd604ab8b57a618ed6e57d6b95df91f4e Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 12:02:59 +0800 Subject: [PATCH 02/95] platform/gateway: add text loop --- platform/gateway/doc.go | 10 + platform/gateway/errors.go | 26 ++ platform/gateway/registry.go | 133 +++++++++ platform/gateway/service.go | 342 +++++++++++++++++++++++ platform/gateway/service_test.go | 449 +++++++++++++++++++++++++++++++ platform/gateway/store.go | 64 +++++ 6 files changed, 1024 insertions(+) create mode 100644 platform/gateway/doc.go create mode 100644 platform/gateway/errors.go create mode 100644 platform/gateway/registry.go create mode 100644 platform/gateway/service.go create mode 100644 platform/gateway/service_test.go create mode 100644 platform/gateway/store.go diff --git a/platform/gateway/doc.go b/platform/gateway/doc.go new file mode 100644 index 0000000000..ebf2c84fc1 --- /dev/null +++ b/platform/gateway/doc.go @@ -0,0 +1,10 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +// Package gateway provides a small text-only IM gateway loop for platform messages. +package gateway diff --git a/platform/gateway/errors.go b/platform/gateway/errors.go new file mode 100644 index 0000000000..7908bfd139 --- /dev/null +++ b/platform/gateway/errors.go @@ -0,0 +1,26 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package gateway + +import "errors" + +var ( + // ErrRuntimeNotFound indicates that no active runtime was registered for the inbound binding. + ErrRuntimeNotFound = errors.New("gateway runtime not found") + // ErrRuntimeInactive indicates that the tenant, app, or binding rejects runtime traffic. + ErrRuntimeInactive = errors.New("gateway runtime inactive") + // ErrRuntimeMismatch indicates that a runtime's tenant, app, binding, or inbound identifiers do not match. + ErrRuntimeMismatch = errors.New("gateway runtime identifiers mismatch") + // ErrUnsupportedMessageType indicates that the gateway batch only supports text input. + ErrUnsupportedMessageType = errors.New("gateway only supports text messages") + // ErrEmptyText indicates that a text message does not contain usable text. + ErrEmptyText = errors.New("gateway text content is required") + // ErrRunnerResponseEmpty indicates that the runner completed without assistant text. + ErrRunnerResponseEmpty = errors.New("gateway runner response is empty") +) diff --git a/platform/gateway/registry.go b/platform/gateway/registry.go new file mode 100644 index 0000000000..0c15b09060 --- /dev/null +++ b/platform/gateway/registry.go @@ -0,0 +1,133 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package gateway + +import ( + "context" + "sync" + + "trpc.group/trpc-go/trpc-agent-go/platform" + "trpc.group/trpc-go/trpc-agent-go/runner" +) + +// Runtime contains the platform configuration and runner for one active binding. +type Runtime struct { + Tenant platform.Tenant + App platform.AgentApp + Binding platform.ChannelBinding + Runner runner.Runner +} + +// Validate checks that the runtime can process inbound messages. +func (r Runtime) Validate() error { + if err := r.Tenant.Validate(); err != nil { + return err + } + if err := r.App.Validate(); err != nil { + return err + } + if err := r.Binding.Validate(); err != nil { + return err + } + if r.Runner == nil { + return ErrRuntimeNotFound + } + if r.App.TenantID != r.Tenant.TenantID || + r.Binding.TenantID != r.Tenant.TenantID || + r.Binding.AppID != r.App.AppID { + return ErrRuntimeMismatch + } + if r.Tenant.Status != "" && r.Tenant.Status != platform.TenantStatusActive { + return ErrRuntimeInactive + } + if r.App.Status != "" && r.App.Status != platform.AppStatusActive { + return ErrRuntimeInactive + } + if r.Binding.Status != "" && r.Binding.Status != platform.BindingStatusActive { + return ErrRuntimeInactive + } + return nil +} + +func (r Runtime) matchesInbound(msg platform.InboundMessage) bool { + return r.Tenant.TenantID == msg.TenantID && + r.App.TenantID == msg.TenantID && + r.App.AppID == msg.AppID && + r.Binding.TenantID == msg.TenantID && + r.Binding.AppID == msg.AppID && + r.Binding.BindingID == msg.BindingID && + r.Binding.Channel == msg.Channel && + r.Binding.AccountID == msg.ChannelAccountID +} + +// Registry resolves an inbound message to an active runtime. +type Registry interface { + Lookup(ctx context.Context, msg platform.InboundMessage) (Runtime, bool, error) +} + +// InMemoryRegistry stores runtimes by tenant, app, binding, channel, and account. +type InMemoryRegistry struct { + mu sync.RWMutex + runtimes map[string]Runtime +} + +// NewInMemoryRegistry creates an in-memory runtime registry. +func NewInMemoryRegistry() *InMemoryRegistry { + return &InMemoryRegistry{ + runtimes: make(map[string]Runtime), + } +} + +// Register stores one runtime. +func (r *InMemoryRegistry) Register(runtime Runtime) error { + if err := runtime.Validate(); err != nil { + return err + } + key := runtimeKey( + runtime.Tenant.TenantID, + runtime.App.AppID, + runtime.Binding.BindingID, + runtime.Binding.Channel, + runtime.Binding.AccountID, + ) + r.mu.Lock() + defer r.mu.Unlock() + r.runtimes[key] = runtime + return nil +} + +// Lookup returns the runtime for an inbound message. +func (r *InMemoryRegistry) Lookup( + ctx context.Context, + msg platform.InboundMessage, +) (Runtime, bool, error) { + if err := ctx.Err(); err != nil { + return Runtime{}, false, err + } + key := runtimeKey( + msg.TenantID, + msg.AppID, + msg.BindingID, + msg.Channel, + msg.ChannelAccountID, + ) + r.mu.RLock() + defer r.mu.RUnlock() + runtime, ok := r.runtimes[key] + return runtime, ok, nil +} + +func runtimeKey(tenantID, appID, bindingID, channel, accountID string) string { + return platform.IdempotencyKey( + tenantID+"|"+appID, + channel, + accountID, + bindingID, + ) +} diff --git a/platform/gateway/service.go b/platform/gateway/service.go new file mode 100644 index 0000000000..cee9bc6f3f --- /dev/null +++ b/platform/gateway/service.go @@ -0,0 +1,342 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package gateway + +import ( + "context" + "fmt" + "strings" + "time" + + "trpc.group/trpc-go/trpc-agent-go/agent" + "trpc.group/trpc-go/trpc-agent-go/event" + "trpc.group/trpc-go/trpc-agent-go/model" + "trpc.group/trpc-go/trpc-agent-go/platform" +) + +// Service handles normalized inbound platform messages. +type Service struct { + registry Registry + idempotencyStore platform.IdempotencyStore + outboundStore OutboundStore + auditSink platform.AuditSink + now func() time.Time +} + +// Option configures a Service. +type Option func(*Service) + +// WithAuditSink sets the audit sink used by the service. +func WithAuditSink(sink platform.AuditSink) Option { + return func(s *Service) { + s.auditSink = sink + } +} + +// WithNow sets the clock used by the service. +func WithNow(now func() time.Time) Option { + return func(s *Service) { + if now != nil { + s.now = now + } + } +} + +// NewService creates a gateway service. +func NewService( + registry Registry, + idempotencyStore platform.IdempotencyStore, + outboundStore OutboundStore, + opts ...Option, +) *Service { + svc := &Service{ + registry: registry, + idempotencyStore: idempotencyStore, + outboundStore: outboundStore, + now: time.Now, + } + for _, opt := range opts { + if opt != nil { + opt(svc) + } + } + return svc +} + +// Result describes the outcome of handling an inbound platform message. +type Result struct { + RequestID string + SessionID string + ResultRef string + Status platform.IdempotencyStatus + Outbound platform.OutboundMessage + Duplicate bool + Processing bool + CompletedAt time.Time +} + +// HandleInbound validates, deduplicates, runs, and records a text-only inbound message. +func (s *Service) HandleInbound( + ctx context.Context, + msg platform.InboundMessage, +) (Result, error) { + start := s.now() + if err := s.validateService(); err != nil { + return Result{}, err + } + if err := msg.Validate(); err != nil { + s.writeAudit(ctx, auditFromMessage(msg, "", "", "reject", err.Error(), start, err)) + return Result{}, err + } + runtime, ok, err := s.registry.Lookup(ctx, msg) + if err != nil { + return Result{}, err + } + if !ok { + err := ErrRuntimeNotFound + s.writeAudit(ctx, auditFromMessage(msg, "", "", "reject", err.Error(), start, err)) + return Result{}, err + } + if err := runtime.Validate(); err != nil { + s.writeAudit(ctx, auditFromMessage(msg, "", "", "reject", err.Error(), start, err)) + return Result{}, err + } + if !runtime.matchesInbound(msg) { + err := ErrRuntimeMismatch + s.writeAudit(ctx, auditFromMessage(msg, "", "", "reject", err.Error(), start, err)) + return Result{}, err + } + text, err := inboundText(msg) + if err != nil { + s.writeAudit(ctx, auditFromMessage(msg, "", "", "reject", err.Error(), start, err)) + return Result{}, err + } + sessionID, err := platform.SessionIDForInbound(msg) + if err != nil { + return Result{}, err + } + internalUserID := platform.InternalUserID(msg.TenantID, msg.Channel, msg.ExternalUserID) + requestID := requestIDFor(msg) + key := platform.IdempotencyKey( + msg.TenantID, + msg.Channel, + msg.ChannelAccountID, + msg.PlatformMessageID, + ) + record, started, err := s.idempotencyStore.Start(ctx, platform.IdempotencyRecord{ + TenantID: msg.TenantID, + Channel: msg.Channel, + AccountID: msg.ChannelAccountID, + PlatformMessageID: msg.PlatformMessageID, + IdempotencyKey: key, + RequestID: requestID, + SessionID: sessionID, + }) + if err != nil { + return Result{}, err + } + if !started { + return s.duplicateResult(ctx, record) + } + + ch, err := runtime.Runner.Run( + ctx, + internalUserID, + sessionID, + model.NewUserMessage(text), + agent.WithRequestID(requestID), + ) + if err != nil { + s.writeAudit(ctx, auditFromMessage(msg, sessionID, internalUserID, "runner_error", err.Error(), start, err)) + return Result{}, err + } + content, err := collectAssistantText(ch) + if err != nil { + s.writeAudit(ctx, auditFromMessage(msg, sessionID, internalUserID, "runner_error", err.Error(), start, err)) + return Result{}, err + } + resultRef := key + ":outbound:1" + outbound := platform.OutboundMessage{ + TenantID: msg.TenantID, + BindingID: msg.BindingID, + Channel: msg.Channel, + SessionID: sessionID, + ReplyToPlatformMessageID: msg.PlatformMessageID, + Kind: platform.OutboundMessageKindText, + Content: content, + Sequence: 1, + DedupKey: resultRef, + TraceID: requestID, + } + if err := s.outboundStore.Save(ctx, resultRef, outbound); err != nil { + return Result{}, err + } + record, err = s.idempotencyStore.Complete(ctx, key, resultRef) + if err != nil { + return Result{}, err + } + s.writeAudit(ctx, auditFromMessage(msg, sessionID, internalUserID, "completed", "", start, nil)) + return Result{ + RequestID: requestID, + SessionID: sessionID, + ResultRef: resultRef, + Status: record.Status, + Outbound: outbound, + CompletedAt: s.now(), + }, nil +} + +func (s *Service) validateService() error { + if s.registry == nil { + return fmt.Errorf("gateway registry is required") + } + if s.idempotencyStore == nil { + return fmt.Errorf("gateway idempotency store is required") + } + if s.outboundStore == nil { + return fmt.Errorf("gateway outbound store is required") + } + return nil +} + +func (s *Service) duplicateResult( + ctx context.Context, + record platform.IdempotencyRecord, +) (Result, error) { + result := Result{ + RequestID: record.RequestID, + SessionID: record.SessionID, + ResultRef: record.ResultRef, + Status: record.Status, + Duplicate: true, + Processing: record.Status == platform.IdempotencyStatusProcessing, + } + if record.Status != platform.IdempotencyStatusCompleted || record.ResultRef == "" { + return result, nil + } + outbound, ok, err := s.outboundStore.Get(ctx, record.ResultRef) + if err != nil { + return Result{}, err + } + if ok { + result.Outbound = outbound + } + return result, nil +} + +func inboundText(msg platform.InboundMessage) (string, error) { + if msg.MessageType != platform.MessageTypeText { + return "", ErrUnsupportedMessageType + } + var parts []string + for _, part := range msg.ContentParts { + if part.Type != platform.ContentPartTypeText { + return "", ErrUnsupportedMessageType + } + text := strings.TrimSpace(part.Text) + if text != "" { + parts = append(parts, text) + } + } + text := strings.TrimSpace(strings.Join(parts, "\n")) + if text == "" { + return "", ErrEmptyText + } + return text, nil +} + +func collectAssistantText(ch <-chan *event.Event) (string, error) { + var parts []string + var final string + for evt := range ch { + if evt == nil || evt.Response == nil { + continue + } + if evt.IsTerminalError() { + return "", evt.Response.Error + } + if evt.IsRunnerCompletion() { + break + } + if len(evt.Choices) == 0 { + continue + } + for _, choice := range evt.Choices { + content := choice.Message.Content + if content == "" { + content = choice.Delta.Content + } + if content != "" { + if evt.Done && !evt.IsPartial && choice.Message.Content != "" { + final = content + continue + } + parts = append(parts, content) + } + } + } + if strings.TrimSpace(final) != "" { + return strings.TrimSpace(final), nil + } + content := strings.TrimSpace(strings.Join(parts, "")) + if content == "" { + return "", ErrRunnerResponseEmpty + } + return content, nil +} + +func requestIDFor(msg platform.InboundMessage) string { + if requestID := strings.TrimSpace(msg.TraceContext["request_id"]); requestID != "" { + return requestID + } + return platform.IdempotencyKey( + msg.TenantID, + msg.Channel, + msg.ChannelAccountID, + msg.PlatformMessageID, + ) +} + +func auditFromMessage( + msg platform.InboundMessage, + sessionID string, + internalUserID string, + decision string, + reason string, + start time.Time, + err error, +) platform.AuditRecord { + record := platform.AuditRecord{ + TenantID: msg.TenantID, + AppID: msg.AppID, + Channel: msg.Channel, + BindingID: msg.BindingID, + UserID: platform.UserIDHash(msg.TenantID, msg.Channel, msg.ExternalUserID), + InternalUserID: internalUserID, + UserIDHash: platform.UserIDHash(msg.TenantID, msg.Channel, msg.ExternalUserID), + SessionID: sessionID, + MessageID: msg.PlatformMessageID, + RequestID: requestIDFor(msg), + Decision: decision, + DecisionReason: reason, + LatencyMS: time.Since(start).Milliseconds(), + CreatedAt: time.Now(), + } + if err != nil { + record.ErrorType = fmt.Sprintf("%T", err) + } + return record +} + +func (s *Service) writeAudit(ctx context.Context, record platform.AuditRecord) { + if s.auditSink == nil { + return + } + _ = s.auditSink.WriteAudit(ctx, record) +} diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go new file mode 100644 index 0000000000..c133920f80 --- /dev/null +++ b/platform/gateway/service_test.go @@ -0,0 +1,449 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package gateway + +import ( + "context" + "errors" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "trpc.group/trpc-go/trpc-agent-go/agent" + "trpc.group/trpc-go/trpc-agent-go/event" + "trpc.group/trpc-go/trpc-agent-go/model" + "trpc.group/trpc-go/trpc-agent-go/platform" +) + +func TestServiceHandleInboundIsolatesTenants(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + runnerA := &recordingRunner{response: "alpha"} + runnerB := &recordingRunner{response: "beta"} + registerRuntime(t, registry, "tenant-a", runnerA) + registerRuntime(t, registry, "tenant-b", runnerB) + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + ) + + resultA, err := svc.HandleInbound(ctx, inbound("tenant-a", "msg-1", "shared-user", "hello")) + require.NoError(t, err) + resultB, err := svc.HandleInbound(ctx, inbound("tenant-b", "msg-1", "shared-user", "hello")) + require.NoError(t, err) + + assert.NotEqual(t, resultA.SessionID, resultB.SessionID) + assert.NotEqual(t, runnerA.calls[0].userID, runnerB.calls[0].userID) + assert.Equal(t, "alpha", resultA.Outbound.Content) + assert.Equal(t, "beta", resultB.Outbound.Content) +} + +func TestServiceHandleInboundDeduplicatesPlatformMessage(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "first"} + registerRuntime(t, registry, "tenant-a", r) + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + ) + msg := inbound("tenant-a", "msg-1", "user-1", "hello") + + first, err := svc.HandleInbound(ctx, msg) + require.NoError(t, err) + second, err := svc.HandleInbound(ctx, msg) + require.NoError(t, err) + + assert.False(t, first.Duplicate) + assert.True(t, second.Duplicate) + assert.False(t, second.Processing) + assert.Equal(t, first.Outbound, second.Outbound) + assert.Len(t, r.calls, 1) +} + +func TestServiceHandleInboundDuplicateProcessingDoesNotRun(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &blockingRunner{started: make(chan struct{})} + registerRuntime(t, registry, "tenant-a", r) + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + ) + msg := inbound("tenant-a", "msg-1", "user-1", "hello") + errCh := make(chan error, 1) + go func() { + _, err := svc.HandleInbound(ctx, msg) + errCh <- err + }() + <-r.started + + dup, err := svc.HandleInbound(ctx, msg) + require.NoError(t, err) + + assert.True(t, dup.Duplicate) + assert.True(t, dup.Processing) + assert.Len(t, r.calls, 1) + r.finish("done") + require.NoError(t, <-errCh) +} + +func TestServiceHandleInboundRejectsUnsupportedMessage(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "unused"} + registerRuntime(t, registry, "tenant-a", r) + audit := platform.NewInMemoryAuditSink() + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + WithAuditSink(audit), + ) + msg := inbound("tenant-a", "msg-1", "user-1", "hello") + msg.MessageType = platform.MessageTypeImage + msg.ContentParts = []platform.ContentPart{{Type: platform.ContentPartTypeImage, FileRef: "artifact://image@1"}} + + _, err := svc.HandleInbound(ctx, msg) + + require.ErrorIs(t, err, ErrUnsupportedMessageType) + assert.Empty(t, r.calls) + require.Len(t, audit.Records(), 1) + assert.Equal(t, "reject", audit.Records()[0].Decision) + assert.NotEqual(t, "user-1", audit.Records()[0].UserID) +} + +func TestServiceHandleInboundRunnerErrorDoesNotComplete(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + runnerErr := errors.New("runner failed") + r := &recordingRunner{runErr: runnerErr} + registerRuntime(t, registry, "tenant-a", r) + store := platform.NewInMemoryIdempotencyStore() + svc := NewService(registry, store, NewInMemoryOutboundStore()) + msg := inbound("tenant-a", "msg-1", "user-1", "hello") + + _, err := svc.HandleInbound(ctx, msg) + + require.ErrorIs(t, err, runnerErr) + record, ok, err := store.Get(ctx, platform.IdempotencyKey("tenant-a", "wecom", "acct", "msg-1")) + require.NoError(t, err) + require.True(t, ok) + assert.Equal(t, platform.IdempotencyStatusProcessing, record.Status) + assert.Empty(t, record.ResultRef) +} + +func TestServiceHandleInboundUsesRequestIDAndStreamsText(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{chunks: []string{"hel", "lo"}} + registerRuntime(t, registry, "tenant-a", r) + svc := NewService(registry, platform.NewInMemoryIdempotencyStore(), NewInMemoryOutboundStore()) + msg := inbound("tenant-a", "msg-1", "user-1", "hello") + msg.TraceContext = map[string]string{"request_id": "req-123"} + + result, err := svc.HandleInbound(ctx, msg) + require.NoError(t, err) + + require.Len(t, r.calls, 1) + assert.Equal(t, "req-123", r.calls[0].requestID) + assert.Equal(t, "hello", result.Outbound.Content) + assert.Equal(t, "req-123", result.Outbound.TraceID) +} + +func TestRuntimeValidateRejectsIdentifierMismatch(t *testing.T) { + runtime := validRuntime("tenant-a", &recordingRunner{response: "unused"}) + runtime.Binding.TenantID = "tenant-b" + + err := runtime.Validate() + + require.ErrorIs(t, err, ErrRuntimeMismatch) +} + +func TestServiceHandleInboundRejectsRegistryMismatch(t *testing.T) { + ctx := context.Background() + r := &recordingRunner{response: "unused"} + registry := staticRegistry{runtime: Runtime{ + Tenant: platform.Tenant{ + TenantID: "tenant-b", + Status: platform.TenantStatusActive, + }, + App: platform.AgentApp{ + TenantID: "tenant-b", + AppID: "app", + AppName: "app", + Status: platform.AppStatusActive, + }, + Binding: platform.ChannelBinding{ + TenantID: "tenant-b", + AppID: "app", + BindingID: "binding", + Channel: "wecom", + AccountID: "acct", + WebhookPath: "/webhook", + TokenRef: "secret://token", + SecretRef: "secret://secret", + Status: platform.BindingStatusActive, + }, + Runner: r, + }} + svc := NewService(registry, platform.NewInMemoryIdempotencyStore(), NewInMemoryOutboundStore()) + + _, err := svc.HandleInbound(ctx, inbound("tenant-a", "msg-1", "user-1", "hello")) + + require.ErrorIs(t, err, ErrRuntimeMismatch) + assert.Empty(t, r.calls) +} + +func TestCollectAssistantTextStopsAtRunnerCompletion(t *testing.T) { + ch := make(chan *event.Event, 2) + ch <- responseEvent("done", true) + ch <- event.NewResponseEvent( + "invocation", + "assistant", + &model.Response{ID: "rc", Object: model.ObjectTypeRunnerCompletion, Done: true}, + ) + + content, err := collectAssistantText(ch) + + require.NoError(t, err) + assert.Equal(t, "done", content) +} + +func TestCollectAssistantTextPrefersFinalFullMessage(t *testing.T) { + ch := make(chan *event.Event, 3) + ch <- chunkEvent("hel", true) + ch <- chunkEvent("lo", true) + ch <- responseEvent("hello", true) + close(ch) + + content, err := collectAssistantText(ch) + + require.NoError(t, err) + assert.Equal(t, "hello", content) +} + +func registerRuntime(t *testing.T, registry *InMemoryRegistry, tenantID string, r runnerStub) { + t.Helper() + err := registry.Register(validRuntime(tenantID, r)) + require.NoError(t, err) +} + +func validRuntime(tenantID string, r runnerStub) Runtime { + return Runtime{ + Tenant: platform.Tenant{ + TenantID: tenantID, + Status: platform.TenantStatusActive, + }, + App: platform.AgentApp{ + TenantID: tenantID, + AppID: "app", + AppName: "app", + Status: platform.AppStatusActive, + }, + Binding: platform.ChannelBinding{ + TenantID: tenantID, + AppID: "app", + BindingID: "binding", + Channel: "wecom", + AccountID: "acct", + WebhookPath: "/webhook", + TokenRef: "secret://token", + SecretRef: "secret://secret", + Status: platform.BindingStatusActive, + ChannelLimits: platform.ChannelLimits{MaxTextLength: 4096}, + }, + Runner: r, + } +} + +func inbound(tenantID, messageID, userID, text string) platform.InboundMessage { + return platform.InboundMessage{ + TenantID: tenantID, + AppID: "app", + BindingID: "binding", + Channel: "wecom", + ChannelAccountID: "acct", + PlatformMessageID: messageID, + ExternalUserID: userID, + ConversationType: platform.ConversationTypeDM, + MessageType: platform.MessageTypeText, + ContentParts: []platform.ContentPart{ + {Type: platform.ContentPartTypeText, Text: text}, + }, + ReceivedAt: time.Unix(100, 0), + } +} + +type runnerStub interface { + Run( + ctx context.Context, + userID string, + sessionID string, + message model.Message, + runOpts ...agent.RunOption, + ) (<-chan *event.Event, error) + Close() error +} + +type runnerCall struct { + userID string + sessionID string + message model.Message + requestID string +} + +type recordingRunner struct { + response string + chunks []string + runErr error + calls []runnerCall +} + +func (r *recordingRunner) Run( + ctx context.Context, + userID string, + sessionID string, + message model.Message, + runOpts ...agent.RunOption, +) (<-chan *event.Event, error) { + if r.runErr != nil { + return nil, r.runErr + } + r.calls = append(r.calls, runnerCall{ + userID: userID, + sessionID: sessionID, + message: message, + requestID: requestIDFromOptions(runOpts...), + }) + out := make(chan *event.Event, 2) + go func() { + defer close(out) + if len(r.chunks) > 0 { + for i, chunk := range r.chunks { + out <- chunkEvent(chunk, i != len(r.chunks)-1) + } + return + } + out <- responseEvent(r.response, true) + }() + return out, nil +} + +func (r *recordingRunner) Close() error { + return nil +} + +type blockingRunner struct { + mu sync.Mutex + started chan struct{} + done chan string + calls []runnerCall +} + +type staticRegistry struct { + runtime Runtime +} + +func (r staticRegistry) Lookup( + ctx context.Context, + msg platform.InboundMessage, +) (Runtime, bool, error) { + if err := ctx.Err(); err != nil { + return Runtime{}, false, err + } + return r.runtime, true, nil +} + +func (r *blockingRunner) Run( + ctx context.Context, + userID string, + sessionID string, + message model.Message, + runOpts ...agent.RunOption, +) (<-chan *event.Event, error) { + r.mu.Lock() + if r.done == nil { + r.done = make(chan string, 1) + } + r.calls = append(r.calls, runnerCall{ + userID: userID, + sessionID: sessionID, + message: message, + requestID: requestIDFromOptions(runOpts...), + }) + close(r.started) + done := r.done + r.mu.Unlock() + out := make(chan *event.Event, 1) + go func() { + defer close(out) + select { + case content := <-done: + out <- responseEvent(content, true) + case <-ctx.Done(): + } + }() + return out, nil +} + +func (r *blockingRunner) Close() error { + return nil +} + +func (r *blockingRunner) finish(content string) { + r.done <- content +} + +func responseEvent(content string, done bool) *event.Event { + return event.NewResponseEvent( + "invocation", + "assistant", + &model.Response{ + ID: content, + Object: model.ObjectTypeChatCompletion, + Done: done, + Choices: []model.Choice{ + {Index: 0, Message: model.Message{Role: model.RoleAssistant, Content: content}}, + }, + }, + ) +} + +func chunkEvent(content string, partial bool) *event.Event { + return event.NewResponseEvent( + "invocation", + "assistant", + &model.Response{ + ID: content, + Object: model.ObjectTypeChatCompletionChunk, + Done: !partial, + IsPartial: partial, + Choices: []model.Choice{ + {Index: 0, Delta: model.Message{Role: model.RoleAssistant, Content: content}}, + }, + }, + ) +} + +func requestIDFromOptions(opts ...agent.RunOption) string { + var runOptions agent.RunOptions + for _, opt := range opts { + if opt != nil { + opt(&runOptions) + } + } + return runOptions.RequestID +} diff --git a/platform/gateway/store.go b/platform/gateway/store.go new file mode 100644 index 0000000000..c3387fa89d --- /dev/null +++ b/platform/gateway/store.go @@ -0,0 +1,64 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package gateway + +import ( + "context" + "sync" + + "trpc.group/trpc-go/trpc-agent-go/platform" +) + +// OutboundStore stores gateway replies so duplicate callbacks can reuse completed results. +type OutboundStore interface { + Save(ctx context.Context, resultRef string, outbound platform.OutboundMessage) error + Get(ctx context.Context, resultRef string) (platform.OutboundMessage, bool, error) +} + +// InMemoryOutboundStore is a concurrency-safe outbound store for tests and demos. +type InMemoryOutboundStore struct { + mu sync.Mutex + messages map[string]platform.OutboundMessage +} + +// NewInMemoryOutboundStore creates an in-memory outbound store. +func NewInMemoryOutboundStore() *InMemoryOutboundStore { + return &InMemoryOutboundStore{ + messages: make(map[string]platform.OutboundMessage), + } +} + +// Save stores one outbound message under resultRef. +func (s *InMemoryOutboundStore) Save( + ctx context.Context, + resultRef string, + outbound platform.OutboundMessage, +) error { + if err := ctx.Err(); err != nil { + return err + } + s.mu.Lock() + defer s.mu.Unlock() + s.messages[resultRef] = outbound + return nil +} + +// Get returns a stored outbound message. +func (s *InMemoryOutboundStore) Get( + ctx context.Context, + resultRef string, +) (platform.OutboundMessage, bool, error) { + if err := ctx.Err(); err != nil { + return platform.OutboundMessage{}, false, err + } + s.mu.Lock() + defer s.mu.Unlock() + outbound, ok := s.messages[resultRef] + return outbound, ok, nil +} From 38636fd96ffdb57e47b518a95e108cb5c57f1895 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 12:25:44 +0800 Subject: [PATCH 03/95] platform/toolpolicy: add governance bridge --- platform/toolpolicy/doc.go | 10 + platform/toolpolicy/policy.go | 411 +++++++++++++++++++++++++++++ platform/toolpolicy/policy_test.go | 370 ++++++++++++++++++++++++++ 3 files changed, 791 insertions(+) create mode 100644 platform/toolpolicy/doc.go create mode 100644 platform/toolpolicy/policy.go create mode 100644 platform/toolpolicy/policy_test.go diff --git a/platform/toolpolicy/doc.go b/platform/toolpolicy/doc.go new file mode 100644 index 0000000000..3b57e911be --- /dev/null +++ b/platform/toolpolicy/doc.go @@ -0,0 +1,10 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +// Package toolpolicy adapts platform tool governance records to runtime policies. +package toolpolicy diff --git a/platform/toolpolicy/policy.go b/platform/toolpolicy/policy.go new file mode 100644 index 0000000000..cef6374be9 --- /dev/null +++ b/platform/toolpolicy/policy.go @@ -0,0 +1,411 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package toolpolicy + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "fmt" + "strings" + "time" + + "trpc.group/trpc-go/trpc-agent-go/platform" + "trpc.group/trpc-go/trpc-agent-go/plugin" + "trpc.group/trpc-go/trpc-agent-go/plugin/guardrail/approval" + "trpc.group/trpc-go/trpc-agent-go/plugin/guardrail/approval/review" + "trpc.group/trpc-go/trpc-agent-go/tool" +) + +// Policy adapts a platform.ToolPolicy to tool.PermissionPolicy. +type Policy struct { + name string + policy platform.ToolPolicy + audit platform.AuditSink + redactor *platform.Redactor + now func() time.Time +} + +// Option configures Policy. +type Option func(*Policy) + +// WithName sets the plugin name used when Policy is registered. +func WithName(name string) Option { + return func(p *Policy) { + name = strings.TrimSpace(name) + if name != "" { + p.name = name + } + } +} + +// WithAuditSink records each non-allow decision to audit. +func WithAuditSink(sink platform.AuditSink) Option { + return func(p *Policy) { + p.audit = sink + } +} + +// WithRedactor sets the redactor used before writing tool arguments to audit. +func WithRedactor(redactor *platform.Redactor) Option { + return func(p *Policy) { + if redactor != nil { + p.redactor = redactor + } + } +} + +// WithNow sets the clock used for audit records. +func WithNow(now func() time.Time) Option { + return func(p *Policy) { + if now != nil { + p.now = now + } + } +} + +// New creates a runtime permission policy from platform tool governance. +func New(policy platform.ToolPolicy, opts ...Option) (*Policy, error) { + if err := validate(policy); err != nil { + return nil, err + } + redactor, err := platform.NewRedactor(policy.ArgumentRedactionRules...) + if err != nil { + return nil, fmt.Errorf("newing platform tool policy: redaction rules: %w", err) + } + p := &Policy{ + name: "platform_tool_policy", + policy: policy, + redactor: redactor, + now: time.Now, + } + for _, opt := range opts { + if opt != nil { + opt(p) + } + } + return p, nil +} + +// Name implements plugin.Plugin when Policy is registered with a plugin manager. +func (p *Policy) Name() string { + if p == nil || p.name == "" { + return "platform_tool_policy" + } + return p.name +} + +// Register adds the policy as a before-tool callback for name-based governance. +// Use CheckToolPermission as a per-run tool.PermissionPolicy when decisions +// must include tool metadata such as destructive, read-only, or open-world. +func (p *Policy) Register(r *plugin.Registry) { + if p == nil || r == nil { + return + } + r.BeforeTool(p.beforeTool()) +} + +// CheckToolPermission implements tool.PermissionPolicy. +func (p *Policy) CheckToolPermission( + ctx context.Context, + req *tool.PermissionRequest, +) (tool.PermissionDecision, error) { + if req == nil { + return tool.AllowPermission(), nil + } + name := strings.TrimSpace(req.ToolName) + if name == "" && req.Declaration != nil { + name = strings.TrimSpace(req.Declaration.Name) + } + decision, reason, audit := p.decide(req, name) + if audit { + p.writeAudit(ctx, req, name, string(decision.Action), reason) + } + return decision, nil +} + +func (p *Policy) beforeTool() tool.BeforeToolCallbackStructured { + return func(ctx context.Context, args *tool.BeforeToolArgs) (*tool.BeforeToolResult, error) { + if args == nil { + return nil, nil + } + req := &tool.PermissionRequest{ + ToolName: args.ToolName, + ToolCallID: args.ToolCallID, + Declaration: args.Declaration, + Arguments: args.Arguments, + } + decision, reason, audit := p.decideNameOnly(req, req.ToolName) + if audit { + p.writeAudit(ctx, req, req.ToolName, string(decision.Action), reason) + } + var err error + if err != nil { + return nil, err + } + decision, err = tool.NormalizePermissionDecision(decision) + if err != nil { + return nil, err + } + if decision.Action == tool.PermissionActionAllow { + return nil, nil + } + return &tool.BeforeToolResult{ + CustomResult: tool.PermissionResultFor(req.ToolName, decision), + }, nil + } +} + +// ApprovalOptions maps name-based parts of the platform policy into approval +// plugin options. A non-empty whitelist remains a hard boundary. Use Policy as +// tool.PermissionPolicy when metadata-based high-risk decisions and +// allow_with_audit records are required. +func ApprovalOptions(policy platform.ToolPolicy) ([]approval.Option, error) { + if err := validate(policy); err != nil { + return nil, err + } + defaultPolicy := approval.ToolPolicySkipApproval + if len(normalizedList(policy.ToolWhitelist)) > 0 || + policy.DangerousToolAction == platform.DangerousToolActionDeny { + defaultPolicy = approval.ToolPolicyDenied + } + opts := []approval.Option{approval.WithDefaultToolPolicy(defaultPolicy)} + whitelist := normalizedList(policy.ToolWhitelist) + hasWhitelist := len(whitelist) > 0 + for _, name := range whitelist { + opts = append(opts, approval.WithToolPolicy(name, approval.ToolPolicySkipApproval)) + } + for _, name := range normalizedList(policy.HighRiskTools) { + if hasWhitelist && !contains(whitelist, name) { + continue + } + switch policy.DangerousToolAction { + case platform.DangerousToolActionDeny: + opts = append(opts, approval.WithToolPolicy(name, approval.ToolPolicyDenied)) + case platform.DangerousToolActionAsk: + opts = append(opts, approval.WithToolPolicy(name, approval.ToolPolicyRequireApproval)) + case platform.DangerousToolActionAllowWithAudit, "": + opts = append(opts, approval.WithToolPolicy(name, approval.ToolPolicySkipApproval)) + } + } + for _, name := range normalizedList(policy.ToolDenylist, policy.PlatformDenylist) { + opts = append(opts, approval.WithToolPolicy(name, approval.ToolPolicyDenied)) + } + return opts, nil +} + +// Reviewer wraps Policy as an approval reviewer for name-boundary checks. It +// rejects denylisted and non-whitelisted tools, but treats ask/approval-required +// decisions as reviewer-approved so the approval plugin can own that flow. +type Reviewer struct { + policy *Policy +} + +// NewReviewer creates an approval reviewer backed by platform tool governance. +func NewReviewer(policy platform.ToolPolicy, opts ...Option) (*Reviewer, error) { + p, err := New(policy, opts...) + if err != nil { + return nil, err + } + return &Reviewer{policy: p}, nil +} + +// Review implements approval/review.Reviewer. +func (r *Reviewer) Review(ctx context.Context, req *review.Request) (*review.Decision, error) { + if r == nil || r.policy == nil || req == nil { + return &review.Decision{Approved: true}, nil + } + permissionReq := &tool.PermissionRequest{ + ToolName: req.Action.ToolName, + Declaration: &tool.Declaration{Name: req.Action.ToolName, Description: req.Action.ToolDescription}, + Arguments: req.Action.Arguments, + } + decision, reason, audit := r.policy.decideReviewer(permissionReq, permissionReq.ToolName) + if audit { + r.policy.writeAudit(ctx, permissionReq, permissionReq.ToolName, string(decision.Action), reason) + } + var err error + decision, err = tool.NormalizePermissionDecision(decision) + if err != nil { + return nil, err + } + return &review.Decision{ + Approved: decision.Action != tool.PermissionActionDeny, + RiskLevel: string(decision.Action), + Reason: decision.Reason, + }, nil +} + +func (p *Policy) decide(req *tool.PermissionRequest, name string) (tool.PermissionDecision, string, bool) { + if contains(policyDenylist(p.policy), name) { + reason := fmt.Sprintf("tool %q is denied by platform tool policy", name) + return tool.DenyPermission(reason), reason, true + } + if len(normalizedList(p.policy.ToolWhitelist)) > 0 && + !contains(normalizedList(p.policy.ToolWhitelist), name) { + reason := fmt.Sprintf("tool %q is not in platform tool whitelist", name) + return tool.DenyPermission(reason), reason, true + } + if isHighRisk(p.policy, req, name) { + switch p.policy.DangerousToolAction { + case platform.DangerousToolActionDeny: + reason := fmt.Sprintf("high-risk tool %q is denied by platform tool policy", name) + return tool.DenyPermission(reason), reason, true + case platform.DangerousToolActionAsk: + reason := fmt.Sprintf("high-risk tool %q requires approval by platform tool policy", name) + return tool.AskPermission(reason), reason, true + case platform.DangerousToolActionAllowWithAudit, "": + reason := fmt.Sprintf("high-risk tool %q allowed with audit by platform tool policy", name) + return tool.AllowPermission(), reason, true + } + } + return tool.AllowPermission(), "", false +} + +func (p *Policy) decideNameOnly(req *tool.PermissionRequest, name string) (tool.PermissionDecision, string, bool) { + if contains(policyDenylist(p.policy), name) { + reason := fmt.Sprintf("tool %q is denied by platform tool policy", name) + return tool.DenyPermission(reason), reason, true + } + if len(normalizedList(p.policy.ToolWhitelist)) > 0 && + !contains(normalizedList(p.policy.ToolWhitelist), name) { + reason := fmt.Sprintf("tool %q is not in platform tool whitelist", name) + return tool.DenyPermission(reason), reason, true + } + if contains(normalizedList(p.policy.HighRiskTools), name) { + switch p.policy.DangerousToolAction { + case platform.DangerousToolActionDeny: + reason := fmt.Sprintf("high-risk tool %q is denied by platform tool policy", name) + return tool.DenyPermission(reason), reason, true + case platform.DangerousToolActionAsk: + reason := fmt.Sprintf("high-risk tool %q requires approval by platform tool policy", name) + return tool.AskPermission(reason), reason, true + case platform.DangerousToolActionAllowWithAudit, "": + reason := fmt.Sprintf("high-risk tool %q allowed with audit by platform tool policy", name) + return tool.AllowPermission(), reason, true + } + } + if req != nil && req.Metadata != (tool.ToolMetadata{}) { + return p.decide(req, name) + } + return tool.AllowPermission(), "", false +} + +func (p *Policy) decideReviewer(req *tool.PermissionRequest, name string) (tool.PermissionDecision, string, bool) { + if contains(policyDenylist(p.policy), name) { + reason := fmt.Sprintf("tool %q is denied by platform tool policy", name) + return tool.DenyPermission(reason), reason, true + } + if len(normalizedList(p.policy.ToolWhitelist)) > 0 && + !contains(normalizedList(p.policy.ToolWhitelist), name) { + reason := fmt.Sprintf("tool %q is not in platform tool whitelist", name) + return tool.DenyPermission(reason), reason, true + } + if contains(normalizedList(p.policy.HighRiskTools), name) { + switch p.policy.DangerousToolAction { + case platform.DangerousToolActionDeny: + reason := fmt.Sprintf("high-risk tool %q is denied by platform tool policy", name) + return tool.DenyPermission(reason), reason, true + case platform.DangerousToolActionAsk: + reason := fmt.Sprintf("high-risk tool %q approved by platform approval reviewer", name) + return tool.AllowPermission(), reason, true + case platform.DangerousToolActionAllowWithAudit, "": + reason := fmt.Sprintf("high-risk tool %q allowed with audit by platform tool policy", name) + return tool.AllowPermission(), reason, true + } + } + if req != nil && req.Metadata != (tool.ToolMetadata{}) { + return p.decide(req, name) + } + return tool.AllowPermission(), "", false +} + +func validate(policy platform.ToolPolicy) error { + switch policy.DangerousToolAction { + case "", platform.DangerousToolActionDeny, + platform.DangerousToolActionAsk, + platform.DangerousToolActionAllowWithAudit: + return nil + default: + return fmt.Errorf("invalid dangerous tool action %q", policy.DangerousToolAction) + } +} + +func isHighRisk(policy platform.ToolPolicy, req *tool.PermissionRequest, name string) bool { + if contains(normalizedList(policy.HighRiskTools), name) { + return true + } + if req == nil { + return false + } + return req.Metadata.Destructive || !req.Metadata.ReadOnly || req.Metadata.OpenWorld +} + +func policyDenylist(policy platform.ToolPolicy) []string { + return normalizedList(policy.ToolDenylist, policy.PlatformDenylist) +} + +func normalizedList(lists ...[]string) []string { + var out []string + seen := make(map[string]struct{}) + for _, list := range lists { + for _, item := range list { + item = strings.TrimSpace(item) + if item == "" { + continue + } + if _, ok := seen[item]; ok { + continue + } + seen[item] = struct{}{} + out = append(out, item) + } + } + return out +} + +func contains(items []string, target string) bool { + for _, item := range items { + if item == target { + return true + } + } + return false +} + +func (p *Policy) writeAudit( + ctx context.Context, + req *tool.PermissionRequest, + toolName string, + decision string, + reason string, +) { + if p.audit == nil { + return + } + argsSummary := argumentSummary(req.Arguments) + _ = p.audit.WriteAudit(ctx, platform.AuditRecord{ + TenantID: p.policy.TenantID, + AppID: p.policy.AppID, + ToolName: toolName, + Decision: decision, + DecisionReason: reason, + RedactedDetailRef: argsSummary, + RedactionVersion: "platform-toolpolicy-v1", + CreatedAt: p.now(), + }) +} + +func argumentSummary(args []byte) string { + if len(args) == 0 { + return "" + } + sum := sha256.Sum256(args) + return fmt.Sprintf("sha256:%s bytes:%d", hex.EncodeToString(sum[:]), len(args)) +} diff --git a/platform/toolpolicy/policy_test.go b/platform/toolpolicy/policy_test.go new file mode 100644 index 0000000000..3289a0b035 --- /dev/null +++ b/platform/toolpolicy/policy_test.go @@ -0,0 +1,370 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package toolpolicy + +import ( + "context" + "strings" + "testing" + "time" + + "trpc.group/trpc-go/trpc-agent-go/platform" + "trpc.group/trpc-go/trpc-agent-go/plugin" + "trpc.group/trpc-go/trpc-agent-go/plugin/guardrail/approval" + "trpc.group/trpc-go/trpc-agent-go/plugin/guardrail/approval/review" + "trpc.group/trpc-go/trpc-agent-go/tool" +) + +func TestPolicyDeniesToolOutsideWhitelist(t *testing.T) { + p := newPolicy(t, platform.ToolPolicy{ + TenantID: "tenant", + AppID: "app", + ToolWhitelist: []string{"knowledge_search"}, + }) + + decision, err := p.CheckToolPermission(context.Background(), request("shell", tool.ToolMetadata{})) + if err != nil { + t.Fatalf("CheckToolPermission: %v", err) + } + if decision.Action != tool.PermissionActionDeny { + t.Fatalf("expected deny, got %+v", decision) + } + if !strings.Contains(decision.Reason, "whitelist") { + t.Fatalf("expected whitelist reason, got %q", decision.Reason) + } +} + +func TestPolicyDenylistOverridesWhitelist(t *testing.T) { + p := newPolicy(t, platform.ToolPolicy{ + ToolWhitelist: []string{"shell"}, + ToolDenylist: []string{"shell"}, + }) + + decision, err := p.CheckToolPermission(context.Background(), request("shell", tool.ToolMetadata{})) + if err != nil { + t.Fatalf("CheckToolPermission: %v", err) + } + if decision.Action != tool.PermissionActionDeny { + t.Fatalf("expected deny, got %+v", decision) + } + if !strings.Contains(decision.Reason, "denied") { + t.Fatalf("expected denied reason, got %q", decision.Reason) + } +} + +func TestPolicyAsksForHighRiskTool(t *testing.T) { + p := newPolicy(t, platform.ToolPolicy{ + DangerousToolAction: platform.DangerousToolActionAsk, + HighRiskTools: []string{"workspace_write"}, + }) + + decision, err := p.CheckToolPermission(context.Background(), request("workspace_write", tool.ToolMetadata{})) + if err != nil { + t.Fatalf("CheckToolPermission: %v", err) + } + if decision.Action != tool.PermissionActionAsk { + t.Fatalf("expected ask, got %+v", decision) + } +} + +func TestPolicyAllowsHighRiskWithAuditAndRedactsArguments(t *testing.T) { + audit := platform.NewInMemoryAuditSink() + now := time.Unix(100, 0) + p := newPolicy( + t, + platform.ToolPolicy{ + TenantID: "tenant", + AppID: "app", + DangerousToolAction: platform.DangerousToolActionAllowWithAudit, + HighRiskTools: []string{"http_post"}, + }, + WithAuditSink(audit), + WithNow(func() time.Time { return now }), + ) + + decision, err := p.CheckToolPermission( + context.Background(), + request("http_post", tool.ToolMetadata{}, []byte(`{"Authorization":"Bearer raw-token","url":"https://example.com"}`)), + ) + if err != nil { + t.Fatalf("CheckToolPermission: %v", err) + } + if decision.Action != tool.PermissionActionAllow { + t.Fatalf("expected allow, got %+v", decision) + } + records := audit.Records() + if len(records) != 1 { + t.Fatalf("expected one audit record, got %d", len(records)) + } + record := records[0] + if record.Decision != string(tool.PermissionActionAllow) || + record.ToolName != "http_post" || + record.TenantID != "tenant" || + record.AppID != "app" || + !record.CreatedAt.Equal(now) { + t.Fatalf("unexpected audit record: %+v", record) + } + if strings.Contains(record.RedactedDetailRef, "raw-token") { + t.Fatalf("audit leaked raw token: %q", record.RedactedDetailRef) + } + if strings.Contains(record.RedactedDetailRef, "example.com") || + strings.Contains(record.RedactedDetailRef, "Authorization") { + t.Fatalf("audit leaked raw argument content: %q", record.RedactedDetailRef) + } + if !strings.HasPrefix(record.RedactedDetailRef, "sha256:") { + t.Fatalf("expected digest summary, got %q", record.RedactedDetailRef) + } +} + +func TestPolicyNilRedactorStillDoesNotLeakArguments(t *testing.T) { + audit := platform.NewInMemoryAuditSink() + p := newPolicy( + t, + platform.ToolPolicy{ + DangerousToolAction: platform.DangerousToolActionAllowWithAudit, + HighRiskTools: []string{"http_post"}, + }, + WithAuditSink(audit), + WithRedactor(nil), + ) + + _, err := p.CheckToolPermission( + context.Background(), + request("http_post", tool.ToolMetadata{}, []byte(`{"email":"person@example.com","path":"/private/file"}`)), + ) + if err != nil { + t.Fatalf("CheckToolPermission: %v", err) + } + records := audit.Records() + if len(records) != 1 { + t.Fatalf("expected one audit record, got %d", len(records)) + } + if strings.Contains(records[0].RedactedDetailRef, "person@example.com") || + strings.Contains(records[0].RedactedDetailRef, "/private/file") { + t.Fatalf("audit leaked raw argument content: %q", records[0].RedactedDetailRef) + } +} + +func TestPolicyDeniesDestructiveMetadata(t *testing.T) { + p := newPolicy(t, platform.ToolPolicy{ + DangerousToolAction: platform.DangerousToolActionDeny, + }) + + decision, err := p.CheckToolPermission( + context.Background(), + request("shell", tool.ToolMetadata{Destructive: true}), + ) + if err != nil { + t.Fatalf("CheckToolPermission: %v", err) + } + if decision.Action != tool.PermissionActionDeny { + t.Fatalf("expected deny, got %+v", decision) + } +} + +func TestPolicyRejectsInvalidDangerousAction(t *testing.T) { + _, err := New(platform.ToolPolicy{DangerousToolAction: platform.DangerousToolAction("bad")}) + if err == nil { + t.Fatalf("expected invalid action to fail") + } +} + +func TestApprovalOptionsMapPolicy(t *testing.T) { + opts, err := ApprovalOptions(platform.ToolPolicy{ + ToolWhitelist: []string{"search", "shell"}, + ToolDenylist: []string{"shell"}, + PlatformDenylist: []string{"admin_delete"}, + DangerousToolAction: platform.DangerousToolActionAsk, + HighRiskTools: []string{"workspace_write"}, + }) + if err != nil { + t.Fatalf("ApprovalOptions: %v", err) + } + p, err := approval.New(append(opts, approval.WithReviewer(allowReviewer{}))...) + if err != nil { + t.Fatalf("approval.New: %v", err) + } + callbacks := plugin.MustNewManager(p).ToolCallbacks() + + denied, err := callbacks.RunBeforeTool(context.Background(), &tool.BeforeToolArgs{ToolName: "shell"}) + if err != nil { + t.Fatalf("RunBeforeTool deny: %v", err) + } + if denied == nil || denied.CustomResult == nil { + t.Fatalf("expected shell to be denied") + } + approvalRequired, err := callbacks.RunBeforeTool(context.Background(), &tool.BeforeToolArgs{ToolName: "workspace_write"}) + if err != nil { + t.Fatalf("RunBeforeTool approval: %v", err) + } + if approvalRequired == nil || approvalRequired.CustomResult == nil { + t.Fatalf("expected high-risk tool outside whitelist to be denied") + } + skipped, err := callbacks.RunBeforeTool(context.Background(), &tool.BeforeToolArgs{ToolName: "search"}) + if err != nil { + t.Fatalf("RunBeforeTool skip: %v", err) + } + if skipped != nil { + t.Fatalf("expected whitelisted search to skip approval, got %+v", skipped) + } + outsideWhitelist, err := callbacks.RunBeforeTool(context.Background(), &tool.BeforeToolArgs{ToolName: "unlisted"}) + if err != nil { + t.Fatalf("RunBeforeTool outside whitelist: %v", err) + } + if outsideWhitelist == nil || outsideWhitelist.CustomResult == nil { + t.Fatalf("expected non-whitelisted tool to be denied") + } +} + +func TestPolicyRegisterAppliesNameBasedGovernance(t *testing.T) { + audit := platform.NewInMemoryAuditSink() + p := newPolicy(t, platform.ToolPolicy{ + TenantID: "tenant", + AppID: "app", + DangerousToolAction: platform.DangerousToolActionAsk, + HighRiskTools: []string{"workspace_write"}, + }, WithAuditSink(audit)) + manager := plugin.MustNewManager(p) + callbacks := manager.ToolCallbacks() + + result, err := callbacks.RunBeforeTool(context.Background(), &tool.BeforeToolArgs{ + ToolName: "workspace_write", + Arguments: []byte(`{"path":"/private/file"}`), + }) + if err != nil { + t.Fatalf("RunBeforeTool: %v", err) + } + if result == nil || result.CustomResult == nil { + t.Fatalf("expected approval-required result") + } + permissionResult, ok := result.CustomResult.(tool.PermissionResult) + if !ok { + t.Fatalf("expected tool.PermissionResult, got %T", result.CustomResult) + } + if permissionResult.Status != tool.PermissionResultStatusApprovalRequired { + t.Fatalf("expected approval_required, got %+v", permissionResult) + } + if strings.Contains(audit.Records()[0].RedactedDetailRef, "/private/file") { + t.Fatalf("audit leaked raw argument content: %q", audit.Records()[0].RedactedDetailRef) + } +} + +func TestPolicyRegisterDoesNotTreatUnknownMetadataAsHighRisk(t *testing.T) { + p := newPolicy(t, platform.ToolPolicy{ + DangerousToolAction: platform.DangerousToolActionAsk, + }) + manager := plugin.MustNewManager(p) + callbacks := manager.ToolCallbacks() + + result, err := callbacks.RunBeforeTool(context.Background(), &tool.BeforeToolArgs{ + ToolName: "read_tool", + }) + if err != nil { + t.Fatalf("RunBeforeTool: %v", err) + } + if result != nil { + t.Fatalf("expected unknown metadata in register path to continue, got %+v", result) + } +} + +func TestReviewerMapsPolicyDecisionToApprovalDecision(t *testing.T) { + reviewer, err := NewReviewer(platform.ToolPolicy{ + ToolWhitelist: []string{"search"}, + }) + if err != nil { + t.Fatalf("NewReviewer: %v", err) + } + + decision, err := reviewer.Review(context.Background(), &review.Request{ + Action: review.Action{ToolName: "shell"}, + }) + if err != nil { + t.Fatalf("Review: %v", err) + } + if decision.Approved { + t.Fatalf("expected shell outside whitelist to be rejected") + } +} + +func TestReviewerApprovesAskDecisionForApprovalPluginFlow(t *testing.T) { + reviewer, err := NewReviewer(platform.ToolPolicy{ + DangerousToolAction: platform.DangerousToolActionAsk, + HighRiskTools: []string{"workspace_write"}, + }) + if err != nil { + t.Fatalf("NewReviewer: %v", err) + } + + decision, err := reviewer.Review(context.Background(), &review.Request{ + Action: review.Action{ToolName: "workspace_write"}, + }) + if err != nil { + t.Fatalf("Review: %v", err) + } + if !decision.Approved { + t.Fatalf("expected ask decision to be approved inside approval flow, got %+v", decision) + } +} + +func TestApprovalOptionsWithReviewerAllowsWhitelistedHighRiskAsk(t *testing.T) { + policy := platform.ToolPolicy{ + ToolWhitelist: []string{"workspace_write"}, + DangerousToolAction: platform.DangerousToolActionAsk, + HighRiskTools: []string{"workspace_write"}, + } + reviewer, err := NewReviewer(policy) + if err != nil { + t.Fatalf("NewReviewer: %v", err) + } + opts, err := ApprovalOptions(policy) + if err != nil { + t.Fatalf("ApprovalOptions: %v", err) + } + p, err := approval.New(append(opts, approval.WithReviewer(reviewer))...) + if err != nil { + t.Fatalf("approval.New: %v", err) + } + callbacks := plugin.MustNewManager(p).ToolCallbacks() + + result, err := callbacks.RunBeforeTool(context.Background(), &tool.BeforeToolArgs{ToolName: "workspace_write"}) + if err != nil { + t.Fatalf("RunBeforeTool: %v", err) + } + if result != nil { + t.Fatalf("expected approved ask flow to continue, got %+v", result) + } +} + +func newPolicy(t *testing.T, policy platform.ToolPolicy, opts ...Option) *Policy { + t.Helper() + p, err := New(policy, opts...) + if err != nil { + t.Fatalf("New: %v", err) + } + return p +} + +func request(name string, metadata tool.ToolMetadata, args ...[]byte) *tool.PermissionRequest { + var payload []byte + if len(args) > 0 { + payload = args[0] + } + return &tool.PermissionRequest{ + ToolName: name, + Declaration: &tool.Declaration{Name: name}, + Arguments: payload, + Metadata: metadata, + } +} + +type allowReviewer struct{} + +func (allowReviewer) Review(context.Context, *review.Request) (*review.Decision, error) { + return &review.Decision{Approved: true}, nil +} From e39271cae1766ff8f1ec1013d924021004f5b65f Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 12:46:08 +0800 Subject: [PATCH 04/95] platform/channeladapter: add adapter outbox skeleton --- platform/channeladapter/adapter.go | 77 ++++ platform/channeladapter/dispatcher.go | 221 +++++++++++ platform/channeladapter/doc.go | 10 + platform/channeladapter/errors.go | 30 ++ platform/channeladapter/outbox.go | 404 ++++++++++++++++++++ platform/channeladapter/outbox_test.go | 501 +++++++++++++++++++++++++ 6 files changed, 1243 insertions(+) create mode 100644 platform/channeladapter/adapter.go create mode 100644 platform/channeladapter/dispatcher.go create mode 100644 platform/channeladapter/doc.go create mode 100644 platform/channeladapter/errors.go create mode 100644 platform/channeladapter/outbox.go create mode 100644 platform/channeladapter/outbox_test.go diff --git a/platform/channeladapter/adapter.go b/platform/channeladapter/adapter.go new file mode 100644 index 0000000000..e24bb7e01d --- /dev/null +++ b/platform/channeladapter/adapter.go @@ -0,0 +1,77 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package channeladapter + +import ( + "context" + "time" + + "trpc.group/trpc-go/trpc-agent-go/platform" +) + +// InboundRequest contains the platform webhook payload after routing to a binding. +type InboundRequest struct { + Binding platform.ChannelBinding + Headers map[string]string + Body []byte + ReceivedAt time.Time + TraceContext map[string]string +} + +// InboundParser converts one platform webhook request into a normalized message. +type InboundParser interface { + ParseInbound(ctx context.Context, req InboundRequest) (platform.InboundMessage, error) +} + +// OutboundProvider delivers normalized outbound messages to one IM platform. +type OutboundProvider interface { + Deliver(ctx context.Context, msg platform.OutboundMessage) (DeliveryResult, error) +} + +// Adapter is the channel boundary. It intentionally does not run agents, +// decide governance, or manage long-term memory. +type Adapter interface { + InboundParser + OutboundProvider + Name() string +} + +// DeliveryResult describes the provider response for one outbound attempt. +type DeliveryResult struct { + Status platform.OutboundStatus + ProviderMessageID string + RetryAfter time.Duration + Detail string +} + +// TextInbound builds a normalized text message from already-verified channel fields. +func TextInbound( + binding platform.ChannelBinding, + platformMessageID string, + externalUserID string, + text string, + receivedAt time.Time, +) platform.InboundMessage { + return platform.InboundMessage{ + TenantID: binding.TenantID, + AppID: binding.AppID, + BindingID: binding.BindingID, + Channel: binding.Channel, + ChannelAccountID: binding.AccountID, + PlatformMessageID: platformMessageID, + ExternalUserID: externalUserID, + ConversationType: platform.ConversationTypeDM, + MessageType: platform.MessageTypeText, + ContentParts: []platform.ContentPart{ + {Type: platform.ContentPartTypeText, Text: text}, + }, + ReceivedAt: receivedAt, + SignatureStatus: "verified", + } +} diff --git a/platform/channeladapter/dispatcher.go b/platform/channeladapter/dispatcher.go new file mode 100644 index 0000000000..66fb4a84df --- /dev/null +++ b/platform/channeladapter/dispatcher.go @@ -0,0 +1,221 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package channeladapter + +import ( + "context" + "errors" + "fmt" + "time" + + "trpc.group/trpc-go/trpc-agent-go/platform" +) + +// ProviderRegistry resolves outbound providers by channel. +type ProviderRegistry interface { + ProviderFor(channel string) (OutboundProvider, bool) +} + +// ProviderRegistryFunc adapts a function to ProviderRegistry. +type ProviderRegistryFunc func(channel string) (OutboundProvider, bool) + +// ProviderFor implements ProviderRegistry. +func (f ProviderRegistryFunc) ProviderFor(channel string) (OutboundProvider, bool) { + if f == nil { + return nil, false + } + return f(channel) +} + +// DispatchResult is the outcome of one due outbox delivery attempt. +type DispatchResult struct { + DedupKey string + Status platform.OutboundStatus + Error error +} + +// Dispatcher drains due outbound messages and sends them through channel providers. +type Dispatcher struct { + store OutboxStore + providers ProviderRegistry + policy RetryPolicy + now func() time.Time + lease time.Duration +} + +// DispatcherOption configures Dispatcher. +type DispatcherOption func(*Dispatcher) + +// WithRetryPolicy sets the dispatch retry policy. +func WithRetryPolicy(policy RetryPolicy) DispatcherOption { + return func(d *Dispatcher) { + d.policy = policy + } +} + +// WithNow sets the dispatch clock. +func WithNow(now func() time.Time) DispatcherOption { + return func(d *Dispatcher) { + if now != nil { + d.now = now + } + } +} + +// WithLeaseDuration sets how long one dispatch worker owns claimed records. +func WithLeaseDuration(lease time.Duration) DispatcherOption { + return func(d *Dispatcher) { + if lease > 0 { + d.lease = lease + } + } +} + +// NewDispatcher creates a due-outbox dispatcher. +func NewDispatcher( + store OutboxStore, + providers ProviderRegistry, + opts ...DispatcherOption, +) *Dispatcher { + d := &Dispatcher{ + store: store, + providers: providers, + policy: DefaultRetryPolicy(), + now: time.Now, + lease: 30 * time.Second, + } + for _, opt := range opts { + if opt != nil { + opt(d) + } + } + return d +} + +// DispatchDue sends due outbox messages and updates their delivery state. +func (d *Dispatcher) DispatchDue(ctx context.Context, limit int) ([]DispatchResult, error) { + if d.store == nil { + return nil, fmt.Errorf("channel adapter outbox store is required") + } + if d.providers == nil { + return nil, fmt.Errorf("channel adapter providers are required") + } + now := d.now() + due, err := d.store.ClaimDue(ctx, now, limit, d.lease) + if err != nil { + return nil, err + } + results := make([]DispatchResult, 0, len(due)) + for _, record := range due { + provider, ok := d.providers.ProviderFor(record.Message.Channel) + if !ok || provider == nil { + updated, markErr := d.store.MarkFailed( + ctx, + record.Message.DedupKey, + record.LeaseToken, + ErrNoProvider, + now, + ) + results = append(results, DispatchResult{ + DedupKey: record.Message.DedupKey, + Status: updated.Status, + Error: markErr, + }) + continue + } + delivery, deliverErr := provider.Deliver(ctx, record.Message) + if deliverErr != nil { + updated, markErr := d.store.MarkFailed( + ctx, + record.Message.DedupKey, + record.LeaseToken, + deliverErr, + now, + ) + results = append(results, DispatchResult{ + DedupKey: record.Message.DedupKey, + Status: updated.Status, + Error: firstErr(markErr, deliverErr), + }) + continue + } + switch delivery.Status { + case platform.OutboundStatusSent: + updated, markErr := d.store.MarkSent( + ctx, + record.Message.DedupKey, + record.LeaseToken, + delivery.ProviderMessageID, + now, + ) + results = append(results, DispatchResult{ + DedupKey: record.Message.DedupKey, + Status: updated.Status, + Error: markErr, + }) + case platform.OutboundStatusFailed: + updated, markErr := d.store.MarkFailedAfter( + ctx, + record.Message.DedupKey, + record.LeaseToken, + deliveryError(delivery), + now, + delivery.RetryAfter, + ) + results = append(results, DispatchResult{ + DedupKey: record.Message.DedupKey, + Status: updated.Status, + Error: markErr, + }) + case platform.OutboundStatusDeadLetter: + updated, markErr := d.store.MarkDeadLetter( + ctx, + record.Message.DedupKey, + record.LeaseToken, + deliveryError(delivery), + now, + ) + results = append(results, DispatchResult{ + DedupKey: record.Message.DedupKey, + Status: updated.Status, + Error: markErr, + }) + default: + updated, markErr := d.store.MarkFailed( + ctx, + record.Message.DedupKey, + record.LeaseToken, + fmt.Errorf("%w: %q", ErrInvalidDeliveryStatus, delivery.Status), + now, + ) + results = append(results, DispatchResult{ + DedupKey: record.Message.DedupKey, + Status: updated.Status, + Error: markErr, + }) + } + } + return results, nil +} + +func deliveryError(delivery DeliveryResult) error { + if delivery.Detail == "" { + return errors.New("provider delivery failed") + } + return errors.New(delivery.Detail) +} + +func firstErr(errs ...error) error { + for _, err := range errs { + if err != nil { + return err + } + } + return nil +} diff --git a/platform/channeladapter/doc.go b/platform/channeladapter/doc.go new file mode 100644 index 0000000000..df304d31e4 --- /dev/null +++ b/platform/channeladapter/doc.go @@ -0,0 +1,10 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +// Package channeladapter defines IM channel adapter contracts and an outbox retry skeleton. +package channeladapter diff --git a/platform/channeladapter/errors.go b/platform/channeladapter/errors.go new file mode 100644 index 0000000000..4a45fdb4be --- /dev/null +++ b/platform/channeladapter/errors.go @@ -0,0 +1,30 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package channeladapter + +import "errors" + +var ( + // ErrOutboundNotFound indicates that an outbox item does not exist. + ErrOutboundNotFound = errors.New("channel adapter outbound not found") + // ErrOutboundDuplicate indicates that the same outbound dedup key already exists. + ErrOutboundDuplicate = errors.New("channel adapter outbound duplicate") + // ErrUnsupportedOutboundKind indicates that a channel cannot deliver the message kind. + ErrUnsupportedOutboundKind = errors.New("channel adapter unsupported outbound kind") + // ErrNoProvider indicates that a dispatcher has no provider for a channel. + ErrNoProvider = errors.New("channel adapter provider not found") + // ErrOutboundNotClaimed indicates that an outbox item was not claimed for delivery. + ErrOutboundNotClaimed = errors.New("channel adapter outbound not claimed") + // ErrOutboundLeaseMismatch indicates that an outbox update used the wrong lease. + ErrOutboundLeaseMismatch = errors.New("channel adapter outbound lease mismatch") + // ErrOutboundLeaseExpired indicates that an outbox lease expired before update. + ErrOutboundLeaseExpired = errors.New("channel adapter outbound lease expired") + // ErrInvalidDeliveryStatus indicates that a provider returned an invalid status. + ErrInvalidDeliveryStatus = errors.New("channel adapter invalid delivery status") +) diff --git a/platform/channeladapter/outbox.go b/platform/channeladapter/outbox.go new file mode 100644 index 0000000000..76a26d877a --- /dev/null +++ b/platform/channeladapter/outbox.go @@ -0,0 +1,404 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package channeladapter + +import ( + "context" + "crypto/rand" + "encoding/hex" + "fmt" + "sync" + "time" + + "trpc.group/trpc-go/trpc-agent-go/platform" +) + +// RetryPolicy controls outbound delivery retry timing. +type RetryPolicy struct { + MaxAttempts int + InitialBackoff time.Duration + MaxBackoff time.Duration +} + +// DefaultRetryPolicy returns a conservative retry policy for IM delivery. +func DefaultRetryPolicy() RetryPolicy { + return RetryPolicy{ + MaxAttempts: 3, + InitialBackoff: time.Second, + MaxBackoff: time.Minute, + } +} + +// RetryPolicyForBinding builds a retry policy from channel limits when present. +func RetryPolicyForBinding(binding platform.ChannelBinding) RetryPolicy { + policy := DefaultRetryPolicy() + if binding.ChannelLimits.RetryMaxAttempts > 0 { + policy.MaxAttempts = binding.ChannelLimits.RetryMaxAttempts + } + return policy +} + +// Delay returns the retry delay for the next attempt. +func (p RetryPolicy) Delay(attempt int) time.Duration { + if p.InitialBackoff <= 0 { + p.InitialBackoff = time.Second + } + if p.MaxBackoff <= 0 { + p.MaxBackoff = time.Minute + } + if attempt <= 1 { + return p.InitialBackoff + } + delay := p.InitialBackoff + for i := 1; i < attempt; i++ { + delay *= 2 + if delay >= p.MaxBackoff { + return p.MaxBackoff + } + } + return delay +} + +// OutboxRecord stores delivery state for one outbound message. +type OutboxRecord struct { + Message platform.OutboundMessage + Status platform.OutboundStatus + Attempts int + MaxAttempts int + RetryPolicy RetryPolicy + NextAttemptAt time.Time + LeaseToken string + LeaseExpiresAt time.Time + LastError string + ProviderMessageID string + CreatedAt time.Time + UpdatedAt time.Time + SentAt *time.Time +} + +// OutboxStore stores outbound messages until they are sent or dead-lettered. +type OutboxStore interface { + Enqueue(ctx context.Context, msg platform.OutboundMessage, policy RetryPolicy) (OutboxRecord, bool, error) + Get(ctx context.Context, dedupKey string) (OutboxRecord, bool, error) + ListDue(ctx context.Context, now time.Time, limit int) ([]OutboxRecord, error) + ClaimDue(ctx context.Context, now time.Time, limit int, leaseDuration time.Duration) ([]OutboxRecord, error) + MarkSent(ctx context.Context, dedupKey string, leaseToken string, providerMessageID string, now time.Time) (OutboxRecord, error) + MarkFailed(ctx context.Context, dedupKey string, leaseToken string, err error, now time.Time) (OutboxRecord, error) + MarkFailedAfter(ctx context.Context, dedupKey string, leaseToken string, err error, now time.Time, retryAfter time.Duration) (OutboxRecord, error) + MarkDeadLetter(ctx context.Context, dedupKey string, leaseToken string, err error, now time.Time) (OutboxRecord, error) +} + +// InMemoryOutboxStore is a concurrency-safe outbox store for tests and demos. +type InMemoryOutboxStore struct { + mu sync.Mutex + records map[string]OutboxRecord +} + +// NewInMemoryOutboxStore creates an in-memory outbox store. +func NewInMemoryOutboxStore() *InMemoryOutboxStore { + return &InMemoryOutboxStore{ + records: make(map[string]OutboxRecord), + } +} + +// Enqueue stores a pending outbound message unless its dedup key already exists. +func (s *InMemoryOutboxStore) Enqueue( + ctx context.Context, + msg platform.OutboundMessage, + policy RetryPolicy, +) (OutboxRecord, bool, error) { + if err := ctx.Err(); err != nil { + return OutboxRecord{}, false, err + } + if msg.DedupKey == "" { + msg.DedupKey = platform.IdempotencyKey( + msg.TenantID, + msg.Channel, + msg.BindingID, + msg.ReplyToPlatformMessageID, + ) + fmt.Sprintf(":seq:%d", msg.Sequence) + } + s.mu.Lock() + defer s.mu.Unlock() + if existing, ok := s.records[msg.DedupKey]; ok { + if !sameOutboundIdentity(existing.Message, msg) { + return OutboxRecord{}, false, ErrOutboundDuplicate + } + return existing, false, nil + } + now := time.Now() + record := OutboxRecord{ + Message: msg, + Status: platform.OutboundStatusPending, + MaxAttempts: maxAttempts(policy), + RetryPolicy: normalizeRetryPolicy(policy), + NextAttemptAt: now, + CreatedAt: now, + UpdatedAt: now, + } + s.records[msg.DedupKey] = record + return record, true, nil +} + +// Get returns one outbox record. +func (s *InMemoryOutboxStore) Get( + ctx context.Context, + dedupKey string, +) (OutboxRecord, bool, error) { + if err := ctx.Err(); err != nil { + return OutboxRecord{}, false, err + } + s.mu.Lock() + defer s.mu.Unlock() + record, ok := s.records[dedupKey] + return record, ok, nil +} + +// ListDue returns pending or failed records due for delivery. +func (s *InMemoryOutboxStore) ListDue( + ctx context.Context, + now time.Time, + limit int, +) ([]OutboxRecord, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + s.mu.Lock() + defer s.mu.Unlock() + out := make([]OutboxRecord, 0) + for _, record := range s.records { + if record.Status != platform.OutboundStatusPending && + record.Status != platform.OutboundStatusFailed { + continue + } + if record.NextAttemptAt.After(now) { + continue + } + out = append(out, record) + if limit > 0 && len(out) >= limit { + break + } + } + return out, nil +} + +// ClaimDue atomically leases due records for one dispatcher worker. +func (s *InMemoryOutboxStore) ClaimDue( + ctx context.Context, + now time.Time, + limit int, + leaseDuration time.Duration, +) ([]OutboxRecord, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + if leaseDuration <= 0 { + leaseDuration = 30 * time.Second + } + s.mu.Lock() + defer s.mu.Unlock() + out := make([]OutboxRecord, 0) + for key, record := range s.records { + if record.Status != platform.OutboundStatusPending && + record.Status != platform.OutboundStatusFailed { + continue + } + if record.NextAttemptAt.After(now) { + continue + } + if record.LeaseToken != "" && record.LeaseExpiresAt.After(now) { + continue + } + record.LeaseToken = newLeaseToken() + record.LeaseExpiresAt = now.Add(leaseDuration) + record.UpdatedAt = now + s.records[key] = record + out = append(out, record) + if limit > 0 && len(out) >= limit { + break + } + } + return out, nil +} + +// MarkSent marks an outbox record as delivered. +func (s *InMemoryOutboxStore) MarkSent( + ctx context.Context, + dedupKey string, + leaseToken string, + providerMessageID string, + now time.Time, +) (OutboxRecord, error) { + if err := ctx.Err(); err != nil { + return OutboxRecord{}, err + } + s.mu.Lock() + defer s.mu.Unlock() + record, ok := s.records[dedupKey] + if !ok { + return OutboxRecord{}, ErrOutboundNotFound + } + if err := validateLease(record, leaseToken, now); err != nil { + return OutboxRecord{}, err + } + record.Status = platform.OutboundStatusSent + record.ProviderMessageID = providerMessageID + record.LeaseToken = "" + record.LeaseExpiresAt = time.Time{} + record.UpdatedAt = now + record.SentAt = &now + s.records[dedupKey] = record + return record, nil +} + +// MarkFailed records a failed delivery attempt and schedules retry or dead-letter. +func (s *InMemoryOutboxStore) MarkFailed( + ctx context.Context, + dedupKey string, + leaseToken string, + err error, + now time.Time, +) (OutboxRecord, error) { + return s.MarkFailedAfter(ctx, dedupKey, leaseToken, err, now, 0) +} + +// MarkFailedAfter records a failed delivery attempt and honors provider retry timing. +func (s *InMemoryOutboxStore) MarkFailedAfter( + ctx context.Context, + dedupKey string, + leaseToken string, + err error, + now time.Time, + retryAfter time.Duration, +) (OutboxRecord, error) { + if ctxErr := ctx.Err(); ctxErr != nil { + return OutboxRecord{}, ctxErr + } + s.mu.Lock() + defer s.mu.Unlock() + record, ok := s.records[dedupKey] + if !ok { + return OutboxRecord{}, ErrOutboundNotFound + } + if err := validateLease(record, leaseToken, now); err != nil { + return OutboxRecord{}, err + } + record.Attempts++ + record.LastError = errorString(err) + record.UpdatedAt = now + record.RetryPolicy = normalizeRetryPolicy(record.RetryPolicy) + record.MaxAttempts = maxAttempts(record.RetryPolicy) + record.LeaseToken = "" + record.LeaseExpiresAt = time.Time{} + if record.Attempts >= record.MaxAttempts { + record.Status = platform.OutboundStatusDeadLetter + record.NextAttemptAt = time.Time{} + } else { + record.Status = platform.OutboundStatusFailed + record.NextAttemptAt = nextAttemptAt(now, retryAfter, record.RetryPolicy.Delay(record.Attempts)) + } + s.records[dedupKey] = record + return record, nil +} + +// MarkDeadLetter records a permanent delivery failure without another retry. +func (s *InMemoryOutboxStore) MarkDeadLetter( + ctx context.Context, + dedupKey string, + leaseToken string, + err error, + now time.Time, +) (OutboxRecord, error) { + if ctxErr := ctx.Err(); ctxErr != nil { + return OutboxRecord{}, ctxErr + } + s.mu.Lock() + defer s.mu.Unlock() + record, ok := s.records[dedupKey] + if !ok { + return OutboxRecord{}, ErrOutboundNotFound + } + if err := validateLease(record, leaseToken, now); err != nil { + return OutboxRecord{}, err + } + record.Attempts++ + record.LastError = errorString(err) + record.UpdatedAt = now + record.Status = platform.OutboundStatusDeadLetter + record.NextAttemptAt = time.Time{} + record.LeaseToken = "" + record.LeaseExpiresAt = time.Time{} + s.records[dedupKey] = record + return record, nil +} + +func maxAttempts(policy RetryPolicy) int { + if policy.MaxAttempts <= 0 { + return DefaultRetryPolicy().MaxAttempts + } + return policy.MaxAttempts +} + +func normalizeRetryPolicy(policy RetryPolicy) RetryPolicy { + defaultPolicy := DefaultRetryPolicy() + if policy.MaxAttempts <= 0 { + policy.MaxAttempts = defaultPolicy.MaxAttempts + } + if policy.InitialBackoff <= 0 { + policy.InitialBackoff = defaultPolicy.InitialBackoff + } + if policy.MaxBackoff <= 0 { + policy.MaxBackoff = defaultPolicy.MaxBackoff + } + return policy +} + +func errorString(err error) string { + if err == nil { + return "" + } + return err.Error() +} + +func sameOutboundIdentity(existing platform.OutboundMessage, next platform.OutboundMessage) bool { + return existing.TenantID == next.TenantID && + existing.Channel == next.Channel && + existing.BindingID == next.BindingID && + existing.ReplyToPlatformMessageID == next.ReplyToPlatformMessageID && + existing.Sequence == next.Sequence +} + +func nextAttemptAt(now time.Time, retryAfter time.Duration, backoff time.Duration) time.Time { + delay := backoff + if retryAfter > delay { + delay = retryAfter + } + return now.Add(delay) +} + +func validateLease(record OutboxRecord, leaseToken string, now time.Time) error { + if record.LeaseToken == "" { + return ErrOutboundNotClaimed + } + if leaseToken == "" || leaseToken != record.LeaseToken { + return ErrOutboundLeaseMismatch + } + if !record.LeaseExpiresAt.IsZero() && !record.LeaseExpiresAt.After(now) { + return ErrOutboundLeaseExpired + } + return nil +} + +func newLeaseToken() string { + var buf [16]byte + if _, err := rand.Read(buf[:]); err != nil { + return fmt.Sprintf("%d", time.Now().UnixNano()) + } + return hex.EncodeToString(buf[:]) +} diff --git a/platform/channeladapter/outbox_test.go b/platform/channeladapter/outbox_test.go new file mode 100644 index 0000000000..889cd25a6f --- /dev/null +++ b/platform/channeladapter/outbox_test.go @@ -0,0 +1,501 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package channeladapter + +import ( + "context" + "errors" + "testing" + "time" + + "trpc.group/trpc-go/trpc-agent-go/platform" +) + +func TestTextInboundUsesBindingBoundary(t *testing.T) { + binding := testBinding() + msg := TextInbound(binding, "msg-1", "user-1", "hello", time.Unix(100, 0)) + + if msg.TenantID != binding.TenantID || + msg.AppID != binding.AppID || + msg.BindingID != binding.BindingID || + msg.ChannelAccountID != binding.AccountID { + t.Fatalf("message did not inherit binding boundary: %+v", msg) + } + if err := msg.Validate(); err != nil { + t.Fatalf("Validate: %v", err) + } + if len(msg.ContentParts) != 1 || msg.ContentParts[0].Text != "hello" { + t.Fatalf("unexpected content parts: %+v", msg.ContentParts) + } +} + +func TestOutboxEnqueueDeduplicates(t *testing.T) { + ctx := context.Background() + store := NewInMemoryOutboxStore() + msg := outbound("reply-1") + + first, inserted, err := store.Enqueue(ctx, msg, DefaultRetryPolicy()) + if err != nil { + t.Fatalf("enqueue first: %v", err) + } + second, insertedAgain, err := store.Enqueue(ctx, msg, DefaultRetryPolicy()) + if err != nil { + t.Fatalf("enqueue duplicate: %v", err) + } + + if !inserted || insertedAgain { + t.Fatalf("expected first insert only, got %v/%v", inserted, insertedAgain) + } + if first.Message.DedupKey != second.Message.DedupKey { + t.Fatalf("duplicate should return existing record") + } +} + +func TestOutboxEnqueueRejectsDedupKeyCollisionAcrossBoundary(t *testing.T) { + ctx := context.Background() + store := NewInMemoryOutboxStore() + msg := outbound("reply-1") + if _, _, err := store.Enqueue(ctx, msg, DefaultRetryPolicy()); err != nil { + t.Fatalf("enqueue first: %v", err) + } + colliding := msg + colliding.TenantID = "other-tenant" + + record, inserted, err := store.Enqueue(ctx, colliding, DefaultRetryPolicy()) + if !errors.Is(err, ErrOutboundDuplicate) { + t.Fatalf("expected duplicate collision, got record=%+v inserted=%v err=%v", record, inserted, err) + } + if inserted { + t.Fatal("collision must not insert") + } + existing, ok, err := store.Get(ctx, msg.DedupKey) + if err != nil || !ok { + t.Fatalf("get existing: %v ok=%v", err, ok) + } + if existing.Message.TenantID != msg.TenantID { + t.Fatalf("collision overwrote existing record: %+v", existing) + } +} + +func TestOutboxEnqueueRejectsDedupKeyCollisionAcrossBinding(t *testing.T) { + ctx := context.Background() + store := NewInMemoryOutboxStore() + msg := outbound("reply-1") + if _, _, err := store.Enqueue(ctx, msg, DefaultRetryPolicy()); err != nil { + t.Fatalf("enqueue first: %v", err) + } + colliding := msg + colliding.BindingID = "other-binding" + + record, inserted, err := store.Enqueue(ctx, colliding, DefaultRetryPolicy()) + if !errors.Is(err, ErrOutboundDuplicate) { + t.Fatalf("expected duplicate collision, got record=%+v inserted=%v err=%v", record, inserted, err) + } + if inserted { + t.Fatal("collision must not insert") + } + existing, ok, err := store.Get(ctx, msg.DedupKey) + if err != nil || !ok { + t.Fatalf("get existing: %v ok=%v", err, ok) + } + if existing.Message.BindingID != msg.BindingID { + t.Fatalf("collision overwrote existing record: %+v", existing) + } +} + +func TestOutboxFailureSchedulesRetryThenDeadLetter(t *testing.T) { + ctx := context.Background() + store := NewInMemoryOutboxStore() + msg := outbound("reply-1") + policy := RetryPolicy{MaxAttempts: 2, InitialBackoff: time.Second, MaxBackoff: time.Second} + _, _, err := store.Enqueue(ctx, msg, policy) + if err != nil { + t.Fatalf("enqueue: %v", err) + } + now := time.Now().Add(time.Hour) + + claimed, err := store.ClaimDue(ctx, now, 1, time.Minute) + if err != nil { + t.Fatalf("claim due: %v", err) + } + failed, err := store.MarkFailed( + ctx, + msg.DedupKey, + claimed[0].LeaseToken, + errors.New("rate limited"), + now, + ) + if err != nil { + t.Fatalf("mark failed: %v", err) + } + if failed.Status != platform.OutboundStatusFailed || + failed.Attempts != 1 || + !failed.NextAttemptAt.Equal(now.Add(time.Second)) { + t.Fatalf("unexpected failed record: %+v", failed) + } + claimed, err = store.ClaimDue(ctx, now.Add(time.Second), 1, time.Minute) + if err != nil { + t.Fatalf("claim retry: %v", err) + } + dead, err := store.MarkFailed( + ctx, + msg.DedupKey, + claimed[0].LeaseToken, + errors.New("still failing"), + now.Add(time.Second), + ) + if err != nil { + t.Fatalf("mark dead letter: %v", err) + } + if dead.Status != platform.OutboundStatusDeadLetter || dead.Attempts != 2 { + t.Fatalf("unexpected dead letter record: %+v", dead) + } +} + +func TestDispatcherMarksSent(t *testing.T) { + ctx := context.Background() + store := NewInMemoryOutboxStore() + msg := outbound("reply-1") + _, _, err := store.Enqueue(ctx, msg, DefaultRetryPolicy()) + if err != nil { + t.Fatalf("enqueue: %v", err) + } + provider := &fakeProvider{result: DeliveryResult{ + Status: platform.OutboundStatusSent, + ProviderMessageID: "provider-1", + }} + dispatcher := NewDispatcher( + store, + ProviderRegistryFunc(func(channel string) (OutboundProvider, bool) { + return provider, channel == "telegram" + }), + WithNow(func() time.Time { return time.Now().Add(time.Hour) }), + ) + + results, err := dispatcher.DispatchDue(ctx, 10) + if err != nil { + t.Fatalf("DispatchDue: %v", err) + } + if len(results) != 1 || results[0].Status != platform.OutboundStatusSent { + t.Fatalf("unexpected dispatch results: %+v", results) + } + record, ok, err := store.Get(ctx, msg.DedupKey) + if err != nil || !ok { + t.Fatalf("get sent record: %v ok=%v", err, ok) + } + if record.ProviderMessageID != "provider-1" || record.SentAt == nil { + t.Fatalf("unexpected sent record: %+v", record) + } + if len(provider.messages) != 1 || provider.messages[0].DedupKey != msg.DedupKey { + t.Fatalf("provider did not receive message: %+v", provider.messages) + } +} + +func TestDispatcherRetriesFailureWithoutRerunningAgent(t *testing.T) { + ctx := context.Background() + store := NewInMemoryOutboxStore() + msg := outbound("reply-1") + policy := RetryPolicy{MaxAttempts: 2, InitialBackoff: time.Second, MaxBackoff: time.Second} + _, _, err := store.Enqueue(ctx, msg, policy) + if err != nil { + t.Fatalf("enqueue: %v", err) + } + provider := &fakeProvider{err: errors.New("temporary failure")} + now := time.Now().Add(time.Hour) + dispatcher := NewDispatcher( + store, + ProviderRegistryFunc(func(channel string) (OutboundProvider, bool) { return provider, true }), + WithRetryPolicy(policy), + WithNow(func() time.Time { return now }), + ) + + results, err := dispatcher.DispatchDue(ctx, 10) + if err != nil { + t.Fatalf("DispatchDue: %v", err) + } + if len(results) != 1 || results[0].Status != platform.OutboundStatusFailed { + t.Fatalf("unexpected first dispatch: %+v", results) + } + due, err := store.ListDue(ctx, now, 10) + if err != nil { + t.Fatalf("ListDue before retry: %v", err) + } + if len(due) != 0 { + t.Fatalf("retry should not be immediately due: %+v", due) + } + + now = now.Add(time.Second) + results, err = dispatcher.DispatchDue(ctx, 10) + if err != nil { + t.Fatalf("DispatchDue retry: %v", err) + } + if len(results) != 1 || results[0].Status != platform.OutboundStatusDeadLetter { + t.Fatalf("unexpected retry dispatch: %+v", results) + } + if len(provider.messages) != 2 { + t.Fatalf("expected outbound retries only, got %d provider calls", len(provider.messages)) + } +} + +func TestDispatcherUsesRecordRetryPolicy(t *testing.T) { + ctx := context.Background() + store := NewInMemoryOutboxStore() + msg := outbound("reply-1") + _, _, err := store.Enqueue(ctx, msg, RetryPolicy{ + MaxAttempts: 1, + InitialBackoff: time.Second, + MaxBackoff: time.Second, + }) + if err != nil { + t.Fatalf("enqueue: %v", err) + } + provider := &fakeProvider{err: errors.New("temporary failure")} + dispatcher := NewDispatcher( + store, + ProviderRegistryFunc(func(channel string) (OutboundProvider, bool) { return provider, true }), + WithNow(func() time.Time { return time.Now().Add(time.Hour) }), + ) + + results, err := dispatcher.DispatchDue(ctx, 10) + if err != nil { + t.Fatalf("DispatchDue: %v", err) + } + if len(results) != 1 || results[0].Status != platform.OutboundStatusDeadLetter { + t.Fatalf("record retry policy should force dead-letter: %+v", results) + } + record, ok, err := store.Get(ctx, msg.DedupKey) + if err != nil || !ok { + t.Fatalf("get record: %v ok=%v", err, ok) + } + if record.MaxAttempts != 1 || record.Status != platform.OutboundStatusDeadLetter { + t.Fatalf("dispatcher should not override record policy: %+v", record) + } +} + +func TestDispatcherRejectsInvalidProviderStatus(t *testing.T) { + ctx := context.Background() + store := NewInMemoryOutboxStore() + msg := outbound("reply-1") + policy := RetryPolicy{MaxAttempts: 2, InitialBackoff: time.Second, MaxBackoff: time.Second} + _, _, err := store.Enqueue(ctx, msg, policy) + if err != nil { + t.Fatalf("enqueue: %v", err) + } + dispatcher := NewDispatcher( + store, + ProviderRegistryFunc(func(channel string) (OutboundProvider, bool) { + return &fakeProvider{result: DeliveryResult{Status: platform.OutboundStatusPending}}, true + }), + WithRetryPolicy(policy), + WithNow(func() time.Time { return time.Now().Add(time.Hour) }), + ) + + results, err := dispatcher.DispatchDue(ctx, 10) + if err != nil { + t.Fatalf("DispatchDue: %v", err) + } + if len(results) != 1 || results[0].Status != platform.OutboundStatusFailed { + t.Fatalf("unexpected dispatch results: %+v", results) + } + record, ok, err := store.Get(ctx, msg.DedupKey) + if err != nil || !ok { + t.Fatalf("get record: %v ok=%v", err, ok) + } + if record.Status != platform.OutboundStatusFailed || + record.LastError != `channel adapter invalid delivery status: "pending"` { + t.Fatalf("invalid status should be failed with diagnostic: %+v", record) + } +} + +func TestDispatcherHonorsProviderRetryAfter(t *testing.T) { + ctx := context.Background() + store := NewInMemoryOutboxStore() + msg := outbound("reply-1") + policy := RetryPolicy{MaxAttempts: 2, InitialBackoff: time.Second, MaxBackoff: time.Second} + _, _, err := store.Enqueue(ctx, msg, policy) + if err != nil { + t.Fatalf("enqueue: %v", err) + } + now := time.Now().Add(time.Hour) + dispatcher := NewDispatcher( + store, + ProviderRegistryFunc(func(channel string) (OutboundProvider, bool) { + return &fakeProvider{result: DeliveryResult{ + Status: platform.OutboundStatusFailed, + RetryAfter: 10 * time.Second, + Detail: "rate limited", + }}, true + }), + WithRetryPolicy(policy), + WithNow(func() time.Time { return now }), + ) + + results, err := dispatcher.DispatchDue(ctx, 10) + if err != nil { + t.Fatalf("DispatchDue: %v", err) + } + if len(results) != 1 || results[0].Status != platform.OutboundStatusFailed { + t.Fatalf("unexpected dispatch results: %+v", results) + } + record, ok, err := store.Get(ctx, msg.DedupKey) + if err != nil || !ok { + t.Fatalf("get record: %v ok=%v", err, ok) + } + if !record.NextAttemptAt.Equal(now.Add(10 * time.Second)) { + t.Fatalf("retry-after should drive next attempt, got %+v", record) + } +} + +func TestDispatcherDeadLettersPermanentProviderFailure(t *testing.T) { + ctx := context.Background() + store := NewInMemoryOutboxStore() + msg := outbound("reply-1") + _, _, err := store.Enqueue(ctx, msg, RetryPolicy{MaxAttempts: 5}) + if err != nil { + t.Fatalf("enqueue: %v", err) + } + dispatcher := NewDispatcher( + store, + ProviderRegistryFunc(func(channel string) (OutboundProvider, bool) { + return &fakeProvider{result: DeliveryResult{ + Status: platform.OutboundStatusDeadLetter, + Detail: ErrUnsupportedOutboundKind.Error(), + }}, true + }), + WithNow(func() time.Time { return time.Now().Add(time.Hour) }), + ) + + results, err := dispatcher.DispatchDue(ctx, 10) + if err != nil { + t.Fatalf("DispatchDue: %v", err) + } + if len(results) != 1 || results[0].Status != platform.OutboundStatusDeadLetter { + t.Fatalf("unexpected dispatch results: %+v", results) + } + record, ok, err := store.Get(ctx, msg.DedupKey) + if err != nil || !ok { + t.Fatalf("get record: %v ok=%v", err, ok) + } + if record.Status != platform.OutboundStatusDeadLetter || record.NextAttemptAt != (time.Time{}) { + t.Fatalf("permanent failure should not retry: %+v", record) + } +} + +func TestDispatcherClaimsDueRecordsOnce(t *testing.T) { + ctx := context.Background() + store := NewInMemoryOutboxStore() + msg := outbound("reply-1") + _, _, err := store.Enqueue(ctx, msg, DefaultRetryPolicy()) + if err != nil { + t.Fatalf("enqueue: %v", err) + } + provider := &fakeProvider{block: make(chan struct{}), result: DeliveryResult{ + Status: platform.OutboundStatusSent, + }} + dispatcher := NewDispatcher( + store, + ProviderRegistryFunc(func(channel string) (OutboundProvider, bool) { return provider, true }), + WithNow(func() time.Time { return time.Now().Add(time.Hour) }), + WithLeaseDuration(time.Minute), + ) + firstDone := make(chan error, 1) + go func() { + _, err := dispatcher.DispatchDue(ctx, 10) + firstDone <- err + }() + provider.waitForCall(t) + + results, err := dispatcher.DispatchDue(ctx, 10) + if err != nil { + t.Fatalf("second DispatchDue: %v", err) + } + if len(results) != 0 { + t.Fatalf("leased record should not be dispatched twice: %+v", results) + } + close(provider.block) + if err := <-firstDone; err != nil { + t.Fatalf("first DispatchDue: %v", err) + } + if len(provider.messages) != 1 { + t.Fatalf("provider should receive one delivery, got %d", len(provider.messages)) + } +} + +func testBinding() platform.ChannelBinding { + return platform.ChannelBinding{ + TenantID: "tenant", + AppID: "app", + BindingID: "binding", + Channel: "telegram", + AccountID: "bot", + WebhookPath: "/channels/telegram/binding/callback", + TokenRef: "secret://telegram-token", + Status: platform.BindingStatusActive, + } +} + +func outbound(dedupKey string) platform.OutboundMessage { + return platform.OutboundMessage{ + TenantID: "tenant", + BindingID: "binding", + Channel: "telegram", + SessionID: "session", + ReplyToPlatformMessageID: "msg-1", + Kind: platform.OutboundMessageKindText, + Content: "hello", + Sequence: 1, + DedupKey: dedupKey, + TraceID: "trace", + } +} + +type fakeProvider struct { + result DeliveryResult + err error + messages []platform.OutboundMessage + block chan struct{} + called chan struct{} +} + +func (p *fakeProvider) Deliver( + ctx context.Context, + msg platform.OutboundMessage, +) (DeliveryResult, error) { + if err := ctx.Err(); err != nil { + return DeliveryResult{}, err + } + if p.called == nil { + p.called = make(chan struct{}) + } + p.messages = append(p.messages, msg) + select { + case <-p.called: + default: + close(p.called) + } + if p.block != nil { + <-p.block + } + if p.err != nil { + return DeliveryResult{}, p.err + } + return p.result, nil +} + +func (p *fakeProvider) waitForCall(t *testing.T) { + t.Helper() + if p.called == nil { + p.called = make(chan struct{}) + } + select { + case <-p.called: + case <-time.After(time.Second): + t.Fatal("provider was not called") + } +} From f053160de8b05c60b65aec73686d92076324197f Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 12:58:08 +0800 Subject: [PATCH 05/95] platform/gateway: enqueue outbound handoff --- platform/gateway/service.go | 17 +++- platform/gateway/service_test.go | 133 +++++++++++++++++++++++++++++++ platform/gateway/store.go | 27 +++++++ 3 files changed, 176 insertions(+), 1 deletion(-) diff --git a/platform/gateway/service.go b/platform/gateway/service.go index cee9bc6f3f..ffe4456bc4 100644 --- a/platform/gateway/service.go +++ b/platform/gateway/service.go @@ -18,6 +18,7 @@ import ( "trpc.group/trpc-go/trpc-agent-go/event" "trpc.group/trpc-go/trpc-agent-go/model" "trpc.group/trpc-go/trpc-agent-go/platform" + "trpc.group/trpc-go/trpc-agent-go/platform/channeladapter" ) // Service handles normalized inbound platform messages. @@ -175,6 +176,18 @@ func (s *Service) HandleInbound( TraceID: requestID, } if err := s.outboundStore.Save(ctx, resultRef, outbound); err != nil { + s.writeAudit(ctx, auditFromMessage(msg, sessionID, internalUserID, "outbound_error", err.Error(), start, err)) + return Result{}, err + } + if err := s.outboundStore.Enqueue( + ctx, + outbound, + channeladapter.RetryPolicyForBinding(runtime.Binding), + ); err != nil { + if _, markErr := s.idempotencyStore.MarkReplyFailed(ctx, key, resultRef); markErr != nil { + return Result{}, markErr + } + s.writeAudit(ctx, auditFromMessage(msg, sessionID, internalUserID, "outbound_error", err.Error(), start, err)) return Result{}, err } record, err = s.idempotencyStore.Complete(ctx, key, resultRef) @@ -217,7 +230,9 @@ func (s *Service) duplicateResult( Duplicate: true, Processing: record.Status == platform.IdempotencyStatusProcessing, } - if record.Status != platform.IdempotencyStatusCompleted || record.ResultRef == "" { + if record.ResultRef == "" || + (record.Status != platform.IdempotencyStatusCompleted && + record.Status != platform.IdempotencyStatusReplyFailed) { return result, nil } outbound, ok, err := s.outboundStore.Get(ctx, record.ResultRef) diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go index c133920f80..4188b4c9dc 100644 --- a/platform/gateway/service_test.go +++ b/platform/gateway/service_test.go @@ -22,6 +22,7 @@ import ( "trpc.group/trpc-go/trpc-agent-go/event" "trpc.group/trpc-go/trpc-agent-go/model" "trpc.group/trpc-go/trpc-agent-go/platform" + "trpc.group/trpc-go/trpc-agent-go/platform/channeladapter" ) func TestServiceHandleInboundIsolatesTenants(t *testing.T) { @@ -72,6 +73,138 @@ func TestServiceHandleInboundDeduplicatesPlatformMessage(t *testing.T) { assert.Len(t, r.calls, 1) } +func TestServiceHandleInboundEnqueuesChannelOutbox(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "queued"} + runtime := validRuntime("tenant-a", r) + runtime.Binding.ChannelLimits.RetryMaxAttempts = 5 + require.NoError(t, registry.Register(runtime)) + outbox := channeladapter.NewInMemoryOutboxStore() + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewOutboxBackedOutboundStore(outbox), + ) + msg := inbound("tenant-a", "msg-1", "user-1", "hello") + + result, err := svc.HandleInbound(ctx, msg) + require.NoError(t, err) + + record, ok, err := outbox.Get(ctx, result.Outbound.DedupKey) + require.NoError(t, err) + require.True(t, ok) + assert.Equal(t, platform.OutboundStatusPending, record.Status) + assert.Equal(t, result.Outbound, record.Message) + assert.Equal(t, 5, record.MaxAttempts) +} + +func TestServiceHandleInboundDuplicateReusesOutboxBackedResult(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "queued"} + registerRuntime(t, registry, "tenant-a", r) + outbox := channeladapter.NewInMemoryOutboxStore() + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewOutboxBackedOutboundStore(outbox), + ) + msg := inbound("tenant-a", "msg-1", "user-1", "hello") + + first, err := svc.HandleInbound(ctx, msg) + require.NoError(t, err) + second, err := svc.HandleInbound(ctx, msg) + require.NoError(t, err) + + assert.True(t, second.Duplicate) + assert.Equal(t, first.Outbound, second.Outbound) + assert.Len(t, r.calls, 1) + due, err := outbox.ListDue(ctx, time.Now().Add(time.Hour), 10) + require.NoError(t, err) + assert.Len(t, due, 1) +} + +func TestServiceHandleInboundOutboxFailureDoesNotCompleteIdempotency(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "queued"} + registerRuntime(t, registry, "tenant-a", r) + idempotency := platform.NewInMemoryIdempotencyStore() + outbox := channeladapter.NewInMemoryOutboxStore() + store := NewOutboxBackedOutboundStore(outbox) + msg := inbound("tenant-a", "msg-1", "user-1", "hello") + resultRef := platform.IdempotencyKey("tenant-a", "wecom", "acct", "msg-1") + ":outbound:1" + colliding := platform.OutboundMessage{ + TenantID: "other-tenant", + BindingID: "binding", + Channel: "wecom", + SessionID: "session", + ReplyToPlatformMessageID: "msg-1", + Kind: platform.OutboundMessageKindText, + Content: "already queued", + Sequence: 1, + DedupKey: resultRef, + } + _, _, err := outbox.Enqueue(ctx, colliding, channeladapter.DefaultRetryPolicy()) + require.NoError(t, err) + audit := platform.NewInMemoryAuditSink() + svc := NewService(registry, idempotency, store, WithAuditSink(audit)) + + _, err = svc.HandleInbound(ctx, msg) + + require.ErrorIs(t, err, channeladapter.ErrOutboundDuplicate) + record, ok, err := idempotency.Get(ctx, platform.IdempotencyKey("tenant-a", "wecom", "acct", "msg-1")) + require.NoError(t, err) + require.True(t, ok) + assert.Equal(t, platform.IdempotencyStatusReplyFailed, record.Status) + assert.Equal(t, resultRef, record.ResultRef) + stored, ok, err := store.Get(ctx, resultRef) + require.NoError(t, err) + require.True(t, ok) + assert.Equal(t, "queued", stored.Content) + require.Len(t, audit.Records(), 1) + assert.Equal(t, "outbound_error", audit.Records()[0].Decision) +} + +func TestServiceHandleInboundDuplicateReplyFailedReusesStoredOutbound(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "queued"} + registerRuntime(t, registry, "tenant-a", r) + idempotency := platform.NewInMemoryIdempotencyStore() + outbox := channeladapter.NewInMemoryOutboxStore() + store := NewOutboxBackedOutboundStore(outbox) + msg := inbound("tenant-a", "msg-1", "user-1", "hello") + resultRef := platform.IdempotencyKey("tenant-a", "wecom", "acct", "msg-1") + ":outbound:1" + colliding := platform.OutboundMessage{ + TenantID: "other-tenant", + BindingID: "binding", + Channel: "wecom", + SessionID: "session", + ReplyToPlatformMessageID: "msg-1", + Kind: platform.OutboundMessageKindText, + Content: "already queued", + Sequence: 1, + DedupKey: resultRef, + } + _, _, err := outbox.Enqueue(ctx, colliding, channeladapter.DefaultRetryPolicy()) + require.NoError(t, err) + svc := NewService(registry, idempotency, store) + _, err = svc.HandleInbound(ctx, msg) + require.ErrorIs(t, err, channeladapter.ErrOutboundDuplicate) + + dup, err := svc.HandleInbound(ctx, msg) + + require.NoError(t, err) + assert.True(t, dup.Duplicate) + assert.False(t, dup.Processing) + assert.Equal(t, platform.IdempotencyStatusReplyFailed, dup.Status) + assert.Equal(t, resultRef, dup.ResultRef) + assert.Equal(t, "queued", dup.Outbound.Content) + assert.Len(t, r.calls, 1) +} + func TestServiceHandleInboundDuplicateProcessingDoesNotRun(t *testing.T) { ctx := context.Background() registry := NewInMemoryRegistry() diff --git a/platform/gateway/store.go b/platform/gateway/store.go index c3387fa89d..8965f2cb3c 100644 --- a/platform/gateway/store.go +++ b/platform/gateway/store.go @@ -13,11 +13,13 @@ import ( "sync" "trpc.group/trpc-go/trpc-agent-go/platform" + "trpc.group/trpc-go/trpc-agent-go/platform/channeladapter" ) // OutboundStore stores gateway replies so duplicate callbacks can reuse completed results. type OutboundStore interface { Save(ctx context.Context, resultRef string, outbound platform.OutboundMessage) error + Enqueue(ctx context.Context, outbound platform.OutboundMessage, policy channeladapter.RetryPolicy) error Get(ctx context.Context, resultRef string) (platform.OutboundMessage, bool, error) } @@ -25,6 +27,7 @@ type OutboundStore interface { type InMemoryOutboundStore struct { mu sync.Mutex messages map[string]platform.OutboundMessage + outbox channeladapter.OutboxStore } // NewInMemoryOutboundStore creates an in-memory outbound store. @@ -34,6 +37,14 @@ func NewInMemoryOutboundStore() *InMemoryOutboundStore { } } +// NewOutboxBackedOutboundStore creates a gateway store that also enqueues channel delivery. +func NewOutboxBackedOutboundStore(outbox channeladapter.OutboxStore) *InMemoryOutboundStore { + return &InMemoryOutboundStore{ + messages: make(map[string]platform.OutboundMessage), + outbox: outbox, + } +} + // Save stores one outbound message under resultRef. func (s *InMemoryOutboundStore) Save( ctx context.Context, @@ -49,6 +60,22 @@ func (s *InMemoryOutboundStore) Save( return nil } +// Enqueue schedules one outbound message for channel delivery when an outbox is configured. +func (s *InMemoryOutboundStore) Enqueue( + ctx context.Context, + outbound platform.OutboundMessage, + policy channeladapter.RetryPolicy, +) error { + if err := ctx.Err(); err != nil { + return err + } + if s.outbox == nil { + return nil + } + _, _, err := s.outbox.Enqueue(ctx, outbound, policy) + return err +} + // Get returns a stored outbound message. func (s *InMemoryOutboundStore) Get( ctx context.Context, From c5e55e9a6065e69793460998228563e8d569d2a8 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 13:06:03 +0800 Subject: [PATCH 06/95] platform/storagerouter: add storage routing contracts --- platform/storagerouter/doc.go | 10 + platform/storagerouter/errors.go | 22 +++ platform/storagerouter/router.go | 253 ++++++++++++++++++++++++++ platform/storagerouter/router_test.go | 147 +++++++++++++++ 4 files changed, 432 insertions(+) create mode 100644 platform/storagerouter/doc.go create mode 100644 platform/storagerouter/errors.go create mode 100644 platform/storagerouter/router.go create mode 100644 platform/storagerouter/router_test.go diff --git a/platform/storagerouter/doc.go b/platform/storagerouter/doc.go new file mode 100644 index 0000000000..26ecce2acd --- /dev/null +++ b/platform/storagerouter/doc.go @@ -0,0 +1,10 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +// Package storagerouter defines tenant-aware storage routing contracts. +package storagerouter diff --git a/platform/storagerouter/errors.go b/platform/storagerouter/errors.go new file mode 100644 index 0000000000..c5e03604d8 --- /dev/null +++ b/platform/storagerouter/errors.go @@ -0,0 +1,22 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package storagerouter + +import "errors" + +var ( + // ErrProfileNotFound indicates that a storage profile is not registered. + ErrProfileNotFound = errors.New("storage router profile not found") + // ErrTenantMismatch indicates that a profile belongs to another tenant. + ErrTenantMismatch = errors.New("storage router tenant mismatch") + // ErrBackendNotFound indicates that a requested backend is not registered. + ErrBackendNotFound = errors.New("storage router backend not found") + // ErrBackendTenantMismatch indicates that a registered backend belongs to another tenant. + ErrBackendTenantMismatch = errors.New("storage router backend tenant mismatch") +) diff --git a/platform/storagerouter/router.go b/platform/storagerouter/router.go new file mode 100644 index 0000000000..b6475bc5c4 --- /dev/null +++ b/platform/storagerouter/router.go @@ -0,0 +1,253 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package storagerouter + +import ( + "context" + "fmt" + "strings" + "sync" + + "trpc.group/trpc-go/trpc-agent-go/artifact" + "trpc.group/trpc-go/trpc-agent-go/knowledge" + "trpc.group/trpc-go/trpc-agent-go/memory" + "trpc.group/trpc-go/trpc-agent-go/platform" + "trpc.group/trpc-go/trpc-agent-go/session" +) + +// BackendSet groups concrete services for one storage backend registration. +type BackendSet struct { + TenantID string + BackendID string + Session session.Service + Memory memory.Service + Artifact artifact.Service + Knowledge knowledge.Knowledge + Audit platform.AuditSink +} + +// Router resolves tenant/app storage services from platform storage profiles. +type Router interface { + Profile(ctx context.Context, tenantID string, profileID string) (platform.StorageProfile, error) + Session(ctx context.Context, tenantID string, profileID string) (session.Service, error) + Memory(ctx context.Context, tenantID string, profileID string) (memory.Service, error) + Artifact(ctx context.Context, tenantID string, profileID string) (artifact.Service, error) + Knowledge(ctx context.Context, tenantID string, profileID string) (knowledge.Knowledge, error) + Audit(ctx context.Context, tenantID string, profileID string) (platform.AuditSink, error) +} + +// InMemoryRouter is a concurrency-safe storage router for tests and demos. +type InMemoryRouter struct { + mu sync.RWMutex + profiles map[profileKey]platform.StorageProfile + backends map[backendKey]BackendSet +} + +type profileKey struct { + tenantID string + profileID string +} + +type backendKey struct { + tenantID string + backendID string +} + +// NewInMemoryRouter creates an empty in-memory storage router. +func NewInMemoryRouter() *InMemoryRouter { + return &InMemoryRouter{ + profiles: make(map[profileKey]platform.StorageProfile), + backends: make(map[backendKey]BackendSet), + } +} + +// RegisterProfile registers or replaces one tenant storage profile. +func (r *InMemoryRouter) RegisterProfile(profile platform.StorageProfile) error { + if err := profile.Validate(); err != nil { + return err + } + r.mu.Lock() + defer r.mu.Unlock() + r.profiles[profileKey{ + tenantID: profile.TenantID, + profileID: profile.ProfileID, + }] = profile + return nil +} + +// RegisterBackend registers or replaces one concrete backend set. +func (r *InMemoryRouter) RegisterBackend(backend BackendSet) error { + if strings.TrimSpace(backend.TenantID) == "" { + return platform.ErrTenantIDRequired + } + if strings.TrimSpace(backend.BackendID) == "" { + return fmt.Errorf("backend_id is required") + } + r.mu.Lock() + defer r.mu.Unlock() + r.backends[backendKey{ + tenantID: backend.TenantID, + backendID: backend.BackendID, + }] = backend + return nil +} + +// Profile resolves one tenant storage profile. +func (r *InMemoryRouter) Profile( + ctx context.Context, + tenantID string, + profileID string, +) (platform.StorageProfile, error) { + if err := ctx.Err(); err != nil { + return platform.StorageProfile{}, err + } + r.mu.RLock() + defer r.mu.RUnlock() + profile, ok := r.profiles[profileKey{tenantID: tenantID, profileID: profileID}] + if !ok { + return platform.StorageProfile{}, ErrProfileNotFound + } + if profile.TenantID != tenantID { + return platform.StorageProfile{}, ErrTenantMismatch + } + return profile, nil +} + +// Session resolves the session service selected by a tenant storage profile. +func (r *InMemoryRouter) Session( + ctx context.Context, + tenantID string, + profileID string, +) (session.Service, error) { + backend, err := r.backend(ctx, tenantID, profileID, resourceSession) + if err != nil { + return nil, err + } + if backend.Session == nil { + return nil, ErrBackendNotFound + } + return backend.Session, nil +} + +// Memory resolves the memory service selected by a tenant storage profile. +func (r *InMemoryRouter) Memory( + ctx context.Context, + tenantID string, + profileID string, +) (memory.Service, error) { + backend, err := r.backend(ctx, tenantID, profileID, resourceMemory) + if err != nil { + return nil, err + } + if backend.Memory == nil { + return nil, ErrBackendNotFound + } + return backend.Memory, nil +} + +// Artifact resolves the artifact service selected by a tenant storage profile. +func (r *InMemoryRouter) Artifact( + ctx context.Context, + tenantID string, + profileID string, +) (artifact.Service, error) { + backend, err := r.backend(ctx, tenantID, profileID, resourceArtifact) + if err != nil { + return nil, err + } + if backend.Artifact == nil { + return nil, ErrBackendNotFound + } + return backend.Artifact, nil +} + +// Knowledge resolves the knowledge service selected by a tenant storage profile. +func (r *InMemoryRouter) Knowledge( + ctx context.Context, + tenantID string, + profileID string, +) (knowledge.Knowledge, error) { + backend, err := r.backend(ctx, tenantID, profileID, resourceKnowledge) + if err != nil { + return nil, err + } + if backend.Knowledge == nil { + return nil, ErrBackendNotFound + } + return backend.Knowledge, nil +} + +// Audit resolves the audit sink selected by a tenant storage profile. +func (r *InMemoryRouter) Audit( + ctx context.Context, + tenantID string, + profileID string, +) (platform.AuditSink, error) { + backend, err := r.backend(ctx, tenantID, profileID, resourceAudit) + if err != nil { + return nil, err + } + if backend.Audit == nil { + return nil, ErrBackendNotFound + } + return backend.Audit, nil +} + +type resourceKind string + +const ( + resourceSession resourceKind = "session" + resourceMemory resourceKind = "memory" + resourceArtifact resourceKind = "artifact" + resourceKnowledge resourceKind = "knowledge" + resourceAudit resourceKind = "audit" +) + +func (r *InMemoryRouter) backend( + ctx context.Context, + tenantID string, + profileID string, + kind resourceKind, +) (BackendSet, error) { + profile, err := r.Profile(ctx, tenantID, profileID) + if err != nil { + return BackendSet{}, err + } + backendID := backendIDFor(profile, kind) + if backendID == "" { + return BackendSet{}, ErrBackendNotFound + } + r.mu.RLock() + defer r.mu.RUnlock() + backend, ok := r.backends[backendKey{tenantID: tenantID, backendID: backendID}] + if !ok { + return BackendSet{}, ErrBackendNotFound + } + if backend.TenantID != tenantID { + return BackendSet{}, ErrBackendTenantMismatch + } + return backend, nil +} + +func backendIDFor(profile platform.StorageProfile, kind resourceKind) string { + switch kind { + case resourceSession: + return strings.TrimSpace(profile.SessionBackend) + case resourceMemory: + return strings.TrimSpace(profile.MemoryBackend) + case resourceArtifact: + return strings.TrimSpace(profile.ArtifactBackend) + case resourceKnowledge: + return strings.TrimSpace(profile.KnowledgeBackend) + case resourceAudit: + return strings.TrimSpace(profile.AuditBackend) + default: + return "" + } +} diff --git a/platform/storagerouter/router_test.go b/platform/storagerouter/router_test.go new file mode 100644 index 0000000000..830e76b210 --- /dev/null +++ b/platform/storagerouter/router_test.go @@ -0,0 +1,147 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package storagerouter + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + artifactmemory "trpc.group/trpc-go/trpc-agent-go/artifact/inmemory" + "trpc.group/trpc-go/trpc-agent-go/knowledge" + memoryinmemory "trpc.group/trpc-go/trpc-agent-go/memory/inmemory" + "trpc.group/trpc-go/trpc-agent-go/platform" + sessioninmemory "trpc.group/trpc-go/trpc-agent-go/session/inmemory" +) + +func TestRouterResolvesTenantStorageServices(t *testing.T) { + ctx := context.Background() + router := NewInMemoryRouter() + sessionSvc := sessioninmemory.NewSessionService() + memorySvc := memoryinmemory.NewMemoryService() + artifactSvc := artifactmemory.NewService() + knowledgeSvc := &stubKnowledge{} + auditSink := platform.NewInMemoryAuditSink() + require.NoError(t, router.RegisterProfile(profile("tenant-a", "profile-a", "hot"))) + require.NoError(t, router.RegisterBackend(BackendSet{ + TenantID: "tenant-a", + BackendID: "hot", + Session: sessionSvc, + Memory: memorySvc, + Artifact: artifactSvc, + Knowledge: knowledgeSvc, + Audit: auditSink, + })) + + gotSession, err := router.Session(ctx, "tenant-a", "profile-a") + require.NoError(t, err) + gotMemory, err := router.Memory(ctx, "tenant-a", "profile-a") + require.NoError(t, err) + gotArtifact, err := router.Artifact(ctx, "tenant-a", "profile-a") + require.NoError(t, err) + gotKnowledge, err := router.Knowledge(ctx, "tenant-a", "profile-a") + require.NoError(t, err) + gotAudit, err := router.Audit(ctx, "tenant-a", "profile-a") + require.NoError(t, err) + + assert.Same(t, sessionSvc, gotSession) + assert.Same(t, memorySvc, gotMemory) + assert.Same(t, artifactSvc, gotArtifact) + assert.Same(t, knowledgeSvc, gotKnowledge) + assert.Same(t, auditSink, gotAudit) +} + +func TestRouterRejectsCrossTenantLookup(t *testing.T) { + ctx := context.Background() + router := NewInMemoryRouter() + require.NoError(t, router.RegisterProfile(profile("tenant-a", "profile-a", "hot"))) + require.NoError(t, router.RegisterBackend(BackendSet{ + TenantID: "tenant-a", + BackendID: "hot", + Session: sessioninmemory.NewSessionService(), + })) + + _, err := router.Session(ctx, "tenant-b", "profile-a") + + require.ErrorIs(t, err, ErrProfileNotFound) +} + +func TestRouterRejectsMissingBackend(t *testing.T) { + ctx := context.Background() + router := NewInMemoryRouter() + require.NoError(t, router.RegisterProfile(profile("tenant-a", "profile-a", "missing"))) + + _, err := router.Session(ctx, "tenant-a", "profile-a") + + require.ErrorIs(t, err, ErrBackendNotFound) +} + +func TestRouterRejectsMissingResourceService(t *testing.T) { + ctx := context.Background() + router := NewInMemoryRouter() + require.NoError(t, router.RegisterProfile(profile("tenant-a", "profile-a", "hot"))) + require.NoError(t, router.RegisterBackend(BackendSet{ + TenantID: "tenant-a", + BackendID: "hot", + })) + + _, err := router.Memory(ctx, "tenant-a", "profile-a") + + require.ErrorIs(t, err, ErrBackendNotFound) +} + +func TestRegisterProfileValidatesSecretRefs(t *testing.T) { + router := NewInMemoryRouter() + p := profile("tenant-a", "profile-a", "hot") + p.DSNRef = "postgres://user:password@localhost/db" + + err := router.RegisterProfile(p) + + require.Error(t, err) + assert.Contains(t, err.Error(), platform.ErrInlineSecretRejected.Error()) +} + +func TestRouterHonorsContextCancellation(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + router := NewInMemoryRouter() + + _, err := router.Profile(ctx, "tenant-a", "profile-a") + + require.True(t, errors.Is(err, context.Canceled)) +} + +func profile(tenantID string, profileID string, backendID string) platform.StorageProfile { + return platform.StorageProfile{ + TenantID: tenantID, + ProfileID: profileID, + SessionBackend: backendID, + MemoryBackend: backendID, + ArtifactBackend: backendID, + KnowledgeBackend: backendID, + AuditBackend: backendID, + DSNRef: "secret://storage", + Namespace: "tenant/" + tenantID, + } +} + +type stubKnowledge struct{} + +func (s *stubKnowledge) Search( + ctx context.Context, + req *knowledge.SearchRequest, +) (*knowledge.SearchResult, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + return &knowledge.SearchResult{}, nil +} From 54881714138d1ff3618063389d7f3611ad0027a6 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 13:22:58 +0800 Subject: [PATCH 07/95] platform/gateway: add session lease --- platform/gateway/lease.go | 79 +++++++++++ platform/gateway/service.go | 57 +++++++- platform/gateway/service_test.go | 235 ++++++++++++++++++++++++++++++- 3 files changed, 361 insertions(+), 10 deletions(-) create mode 100644 platform/gateway/lease.go diff --git a/platform/gateway/lease.go b/platform/gateway/lease.go new file mode 100644 index 0000000000..f31c671e4c --- /dev/null +++ b/platform/gateway/lease.go @@ -0,0 +1,79 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package gateway + +import ( + "context" + "sync" +) + +// SessionLeaseStore serializes gateway execution for the same tenant/app/session. +type SessionLeaseStore interface { + Acquire(ctx context.Context, key SessionLeaseKey) (SessionLease, bool, error) +} + +// SessionLeaseKey identifies the gateway execution slot for one session. +type SessionLeaseKey struct { + TenantID string + AppID string + SessionID string +} + +// SessionLease releases one acquired session execution slot. +type SessionLease interface { + Release(ctx context.Context) error +} + +// InMemorySessionLeaseStore is a process-local lease store for tests and demos. +type InMemorySessionLeaseStore struct { + mu sync.Mutex + leases map[SessionLeaseKey]struct{} +} + +// NewInMemorySessionLeaseStore creates an empty process-local session lease store. +func NewInMemorySessionLeaseStore() *InMemorySessionLeaseStore { + return &InMemorySessionLeaseStore{ + leases: make(map[SessionLeaseKey]struct{}), + } +} + +// Acquire tries to acquire the session lease without waiting. +func (s *InMemorySessionLeaseStore) Acquire( + ctx context.Context, + key SessionLeaseKey, +) (SessionLease, bool, error) { + if err := ctx.Err(); err != nil { + return nil, false, err + } + s.mu.Lock() + defer s.mu.Unlock() + if _, ok := s.leases[key]; ok { + return nil, false, nil + } + s.leases[key] = struct{}{} + return &inMemorySessionLease{ + store: s, + key: key, + }, true, nil +} + +type inMemorySessionLease struct { + store *InMemorySessionLeaseStore + key SessionLeaseKey + once sync.Once +} + +func (l *inMemorySessionLease) Release(ctx context.Context) error { + l.once.Do(func() { + l.store.mu.Lock() + defer l.store.mu.Unlock() + delete(l.store.leases, l.key) + }) + return ctx.Err() +} diff --git a/platform/gateway/service.go b/platform/gateway/service.go index ffe4456bc4..60a74b7356 100644 --- a/platform/gateway/service.go +++ b/platform/gateway/service.go @@ -26,6 +26,7 @@ type Service struct { registry Registry idempotencyStore platform.IdempotencyStore outboundStore OutboundStore + leaseStore SessionLeaseStore auditSink platform.AuditSink now func() time.Time } @@ -49,6 +50,13 @@ func WithNow(now func() time.Time) Option { } } +// WithSessionLeaseStore sets the lease store used to serialize same-session runs. +func WithSessionLeaseStore(store SessionLeaseStore) Option { + return func(s *Service) { + s.leaseStore = store + } +} + // NewService creates a gateway service. func NewService( registry Registry, @@ -60,6 +68,7 @@ func NewService( registry: registry, idempotencyStore: idempotencyStore, outboundStore: outboundStore, + leaseStore: NewInMemorySessionLeaseStore(), now: time.Now, } for _, opt := range opts { @@ -130,6 +139,34 @@ func (s *Service) HandleInbound( msg.ChannelAccountID, msg.PlatformMessageID, ) + existing, ok, err := s.idempotencyStore.Get(ctx, key) + if err != nil { + return Result{}, err + } + if ok { + return s.duplicateResult(ctx, existing) + } + lease, acquired, err := s.leaseStore.Acquire(ctx, SessionLeaseKey{ + TenantID: msg.TenantID, + AppID: msg.AppID, + SessionID: sessionID, + }) + if err != nil { + return Result{}, err + } + if !acquired { + return Result{ + RequestID: requestID, + SessionID: sessionID, + Status: platform.IdempotencyStatusProcessing, + Processing: true, + }, nil + } + defer func() { + cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + defer cancel() + _ = lease.Release(cleanupCtx) + }() record, started, err := s.idempotencyStore.Start(ctx, platform.IdempotencyRecord{ TenantID: msg.TenantID, Channel: msg.Channel, @@ -157,7 +194,7 @@ func (s *Service) HandleInbound( s.writeAudit(ctx, auditFromMessage(msg, sessionID, internalUserID, "runner_error", err.Error(), start, err)) return Result{}, err } - content, err := collectAssistantText(ch) + content, err := collectAssistantText(ctx, ch) if err != nil { s.writeAudit(ctx, auditFromMessage(msg, sessionID, internalUserID, "runner_error", err.Error(), start, err)) return Result{}, err @@ -215,6 +252,9 @@ func (s *Service) validateService() error { if s.outboundStore == nil { return fmt.Errorf("gateway outbound store is required") } + if s.leaseStore == nil { + return fmt.Errorf("gateway session lease store is required") + } return nil } @@ -266,10 +306,20 @@ func inboundText(msg platform.InboundMessage) (string, error) { return text, nil } -func collectAssistantText(ch <-chan *event.Event) (string, error) { +func collectAssistantText(ctx context.Context, ch <-chan *event.Event) (string, error) { var parts []string var final string - for evt := range ch { + for { + var evt *event.Event + select { + case <-ctx.Done(): + return "", ctx.Err() + case next, ok := <-ch: + if !ok { + goto done + } + evt = next + } if evt == nil || evt.Response == nil { continue } @@ -296,6 +346,7 @@ func collectAssistantText(ch <-chan *event.Event) (string, error) { } } } +done: if strings.TrimSpace(final) != "" { return strings.TrimSpace(final), nil } diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go index 4188b4c9dc..45a9971522 100644 --- a/platform/gateway/service_test.go +++ b/platform/gateway/service_test.go @@ -233,6 +233,130 @@ func TestServiceHandleInboundDuplicateProcessingDoesNotRun(t *testing.T) { require.NoError(t, <-errCh) } +func TestServiceHandleInboundSerializesSameSession(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &blockingRunner{started: make(chan struct{})} + registerRuntime(t, registry, "tenant-a", r) + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + ) + first := inbound("tenant-a", "msg-1", "user-1", "hello") + second := inbound("tenant-a", "msg-2", "user-1", "again") + errCh := make(chan error, 1) + go func() { + _, err := svc.HandleInbound(ctx, first) + errCh <- err + }() + <-r.started + + busy, err := svc.HandleInbound(ctx, second) + require.NoError(t, err) + + assert.False(t, busy.Duplicate) + assert.True(t, busy.Processing) + assert.Equal(t, platform.IdempotencyStatusProcessing, busy.Status) + assert.Equal(t, "tenant:tenant-a:app:app:channel:wecom:dm:user-1", busy.SessionID) + assert.Len(t, r.calls, 1) + record, ok, err := svc.idempotencyStore.Get( + ctx, + platform.IdempotencyKey("tenant-a", "wecom", "acct", "msg-2"), + ) + require.NoError(t, err) + assert.False(t, ok) + assert.Empty(t, record.ResultRef) + r.finish("done") + require.NoError(t, <-errCh) + + r.finish("again") + retry, err := svc.HandleInbound(ctx, second) + require.NoError(t, err) + assert.False(t, retry.Processing) + assert.Equal(t, platform.IdempotencyStatusCompleted, retry.Status) + assert.Len(t, r.calls, 2) +} + +func TestServiceHandleInboundAllowsDifferentSessions(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "done"} + registerRuntime(t, registry, "tenant-a", r) + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + ) + + _, err := svc.HandleInbound(ctx, inbound("tenant-a", "msg-1", "user-1", "hello")) + require.NoError(t, err) + _, err = svc.HandleInbound(ctx, inbound("tenant-a", "msg-2", "user-2", "hello")) + require.NoError(t, err) + + require.Len(t, r.calls, 2) + assert.NotEqual(t, r.calls[0].sessionID, r.calls[1].sessionID) +} + +func TestServiceHandleInboundReleaseIgnoresCanceledRequestContext(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + registry := NewInMemoryRegistry() + runnerErr := errors.New("runner failed") + r := &cancelingRunner{cancel: cancel, runErr: runnerErr} + registerRuntime(t, registry, "tenant-a", r) + lease := &recordingLease{} + leaseStore := &recordingLeaseStore{lease: lease} + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + WithSessionLeaseStore(leaseStore), + ) + + _, err := svc.HandleInbound(ctx, inbound("tenant-a", "msg-1", "user-1", "hello")) + + require.ErrorIs(t, err, runnerErr) + require.True(t, lease.released) + require.NoError(t, lease.ctxErr) +} + +func TestServiceHandleInboundCancellationDuringEventCollectionReleasesSessionLease(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + registry := NewInMemoryRegistry() + r := &hangingFirstRunner{ + started: make(chan struct{}), + response: "done", + } + registerRuntime(t, registry, "tenant-a", r) + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + ) + first := inbound("tenant-a", "msg-1", "user-1", "hello") + second := inbound("tenant-a", "msg-2", "user-1", "again") + errCh := make(chan error, 1) + go func() { + _, err := svc.HandleInbound(ctx, first) + errCh <- err + }() + <-r.started + + cancel() + select { + case err := <-errCh: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(time.Second): + t.Fatal("HandleInbound did not return after context cancellation") + } + retry, err := svc.HandleInbound(context.Background(), second) + + require.NoError(t, err) + assert.False(t, retry.Processing) + assert.Equal(t, platform.IdempotencyStatusCompleted, retry.Status) + assert.Len(t, r.calls, 2) +} + func TestServiceHandleInboundRejectsUnsupportedMessage(t *testing.T) { ctx := context.Background() registry := NewInMemoryRegistry() @@ -349,7 +473,7 @@ func TestCollectAssistantTextStopsAtRunnerCompletion(t *testing.T) { &model.Response{ID: "rc", Object: model.ObjectTypeRunnerCompletion, Done: true}, ) - content, err := collectAssistantText(ch) + content, err := collectAssistantText(context.Background(), ch) require.NoError(t, err) assert.Equal(t, "done", content) @@ -362,7 +486,7 @@ func TestCollectAssistantTextPrefersFinalFullMessage(t *testing.T) { ch <- responseEvent("hello", true) close(ch) - content, err := collectAssistantText(ch) + content, err := collectAssistantText(context.Background(), ch) require.NoError(t, err) assert.Equal(t, "hello", content) @@ -480,10 +604,17 @@ func (r *recordingRunner) Close() error { } type blockingRunner struct { - mu sync.Mutex - started chan struct{} - done chan string - calls []runnerCall + mu sync.Mutex + started chan struct{} + startedOnce sync.Once + done chan string + calls []runnerCall +} + +type cancelingRunner struct { + cancel func() + runErr error + calls []runnerCall } type staticRegistry struct { @@ -517,7 +648,9 @@ func (r *blockingRunner) Run( message: message, requestID: requestIDFromOptions(runOpts...), }) - close(r.started) + r.startedOnce.Do(func() { + close(r.started) + }) done := r.done r.mu.Unlock() out := make(chan *event.Event, 1) @@ -540,6 +673,94 @@ func (r *blockingRunner) finish(content string) { r.done <- content } +func (r *cancelingRunner) Run( + ctx context.Context, + userID string, + sessionID string, + message model.Message, + runOpts ...agent.RunOption, +) (<-chan *event.Event, error) { + r.calls = append(r.calls, runnerCall{ + userID: userID, + sessionID: sessionID, + message: message, + requestID: requestIDFromOptions(runOpts...), + }) + r.cancel() + return nil, r.runErr +} + +func (r *cancelingRunner) Close() error { + return nil +} + +type hangingFirstRunner struct { + mu sync.Mutex + started chan struct{} + startedOnce sync.Once + response string + calls []runnerCall +} + +func (r *hangingFirstRunner) Run( + ctx context.Context, + userID string, + sessionID string, + message model.Message, + runOpts ...agent.RunOption, +) (<-chan *event.Event, error) { + r.mu.Lock() + callIndex := len(r.calls) + r.calls = append(r.calls, runnerCall{ + userID: userID, + sessionID: sessionID, + message: message, + requestID: requestIDFromOptions(runOpts...), + }) + if callIndex == 0 { + r.startedOnce.Do(func() { + close(r.started) + }) + } + r.mu.Unlock() + if callIndex == 0 { + return make(chan *event.Event), nil + } + out := make(chan *event.Event, 1) + out <- responseEvent(r.response, true) + close(out) + return out, nil +} + +func (r *hangingFirstRunner) Close() error { + return nil +} + +type recordingLeaseStore struct { + lease *recordingLease +} + +func (s *recordingLeaseStore) Acquire( + ctx context.Context, + key SessionLeaseKey, +) (SessionLease, bool, error) { + if err := ctx.Err(); err != nil { + return nil, false, err + } + return s.lease, true, nil +} + +type recordingLease struct { + released bool + ctxErr error +} + +func (l *recordingLease) Release(ctx context.Context) error { + l.released = true + l.ctxErr = ctx.Err() + return nil +} + func responseEvent(content string, done bool) *event.Event { return event.NewResponseEvent( "invocation", From 64d0a9a763de22048fd00ceaaef1783e35938228 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 13:30:27 +0800 Subject: [PATCH 08/95] platform: add gray routing helpers --- platform/gray.go | 50 ++++++++++++++++++++++++++++ platform/gray_test.go | 77 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 127 insertions(+) create mode 100644 platform/gray.go create mode 100644 platform/gray_test.go diff --git a/platform/gray.go b/platform/gray.go new file mode 100644 index 0000000000..41e7e1ca05 --- /dev/null +++ b/platform/gray.go @@ -0,0 +1,50 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "fmt" + "hash/fnv" + "strings" +) + +// SessionGrayBucket returns the stable 0-99 release bucket for one session. +func SessionGrayBucket(tenantID, appID, sessionID string) (int, error) { + tenantID = strings.TrimSpace(tenantID) + appID = strings.TrimSpace(appID) + sessionID = strings.TrimSpace(sessionID) + if tenantID == "" { + return 0, ErrTenantIDRequired + } + if appID == "" { + return 0, ErrAppIDRequired + } + if sessionID == "" { + return 0, fmt.Errorf("session_id is required") + } + h := fnv.New32a() + _, _ = h.Write([]byte(tenantID)) + _, _ = h.Write([]byte{0}) + _, _ = h.Write([]byte(appID)) + _, _ = h.Write([]byte{0}) + _, _ = h.Write([]byte(sessionID)) + return int(h.Sum32() % 100), nil +} + +// SessionInGrayRelease reports whether one session belongs to the app's gray release. +func SessionInGrayRelease(app AgentApp, sessionID string) (bool, int, error) { + if err := app.Validate(); err != nil { + return false, 0, err + } + bucket, err := SessionGrayBucket(app.TenantID, app.AppID, sessionID) + if err != nil { + return false, 0, err + } + return bucket < app.GrayPercent, bucket, nil +} diff --git a/platform/gray_test.go b/platform/gray_test.go new file mode 100644 index 0000000000..c4d49b8a58 --- /dev/null +++ b/platform/gray_test.go @@ -0,0 +1,77 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "errors" + "testing" +) + +func TestSessionGrayBucketIsStableByTenantAppSession(t *testing.T) { + first, err := SessionGrayBucket("tenant-a", "app-a", "session-1") + if err != nil { + t.Fatalf("bucket: %v", err) + } + second, err := SessionGrayBucket(" tenant-a ", " app-a ", " session-1 ") + if err != nil { + t.Fatalf("bucket: %v", err) + } + if first != second { + t.Fatalf("same trimmed session key should stay in same bucket: %d != %d", first, second) + } + if first < 0 || first >= 100 { + t.Fatalf("bucket should be 0-99, got %d", first) + } +} + +func TestSessionInGrayReleaseUsesConfiguredPercentBoundary(t *testing.T) { + app := AgentApp{TenantID: "tenant-a", AppID: "app-a", GrayPercent: 0} + inGray, bucket, err := SessionInGrayRelease(app, "session-1") + if err != nil { + t.Fatalf("gray decision: %v", err) + } + if inGray { + t.Fatalf("0 percent should not include bucket %d", bucket) + } + + app.GrayPercent = 100 + inGray, bucket, err = SessionInGrayRelease(app, "session-1") + if err != nil { + t.Fatalf("gray decision: %v", err) + } + if !inGray { + t.Fatalf("100 percent should include bucket %d", bucket) + } +} + +func TestSessionInGrayReleaseMatchesBucketThreshold(t *testing.T) { + app := AgentApp{TenantID: "tenant-a", AppID: "app-a", GrayPercent: 50} + inGray, bucket, err := SessionInGrayRelease(app, "session-1") + if err != nil { + t.Fatalf("gray decision: %v", err) + } + if inGray != (bucket < app.GrayPercent) { + t.Fatalf("gray decision should match bucket threshold: in_gray=%t bucket=%d percent=%d", inGray, bucket, app.GrayPercent) + } +} + +func TestSessionInGrayReleaseRejectsInvalidInputs(t *testing.T) { + _, _, err := SessionInGrayRelease(AgentApp{AppID: "app", GrayPercent: 10}, "session") + if !errors.Is(err, ErrTenantIDRequired) { + t.Fatalf("expected tenant id error, got %v", err) + } + _, _, err = SessionInGrayRelease(AgentApp{TenantID: "tenant", AppID: "app", GrayPercent: 101}, "session") + if err == nil { + t.Fatalf("expected invalid gray percent error") + } + _, _, err = SessionInGrayRelease(AgentApp{TenantID: "tenant", AppID: "app", GrayPercent: 10}, "") + if err == nil { + t.Fatalf("expected missing session id error") + } +} From cf5af61a72b8863c8db7174ba9314ae521809664 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 13:38:27 +0800 Subject: [PATCH 09/95] platform: add tenant budget helpers --- platform/budget.go | 135 ++++++++++++++++++++++++++++++++++++++ platform/budget_test.go | 139 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 274 insertions(+) create mode 100644 platform/budget.go create mode 100644 platform/budget_test.go diff --git a/platform/budget.go b/platform/budget.go new file mode 100644 index 0000000000..a3c4b6fdd1 --- /dev/null +++ b/platform/budget.go @@ -0,0 +1,135 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "encoding/json" + "fmt" + "math" + "strings" +) + +// TenantQuota captures tenant-level budget limits from Tenant.QuotaJSON. +type TenantQuota struct { + MaxPromptTokens int `json:"max_prompt_tokens,omitempty"` + MaxCompletionTokens int `json:"max_completion_tokens,omitempty"` + MaxTotalTokens int `json:"max_total_tokens,omitempty"` + MaxCost float64 `json:"max_cost,omitempty"` +} + +// UsageEstimate is the pre-run cost and token estimate checked against quota. +type UsageEstimate struct { + PromptTokens int + CompletionTokens int + TotalTokens int + Cost float64 +} + +// BudgetDecision describes whether a usage estimate is allowed by tenant quota. +type BudgetDecision struct { + Allowed bool + Reason string +} + +// ParseTenantQuota parses Tenant.QuotaJSON. Empty quota means no budget limits. +func ParseTenantQuota(tenant Tenant) (TenantQuota, error) { + if err := tenant.Validate(); err != nil { + return TenantQuota{}, err + } + quotaJSON := strings.TrimSpace(tenant.QuotaJSON) + if quotaJSON == "" { + return TenantQuota{}, nil + } + var quota TenantQuota + if err := json.Unmarshal([]byte(quotaJSON), "a); err != nil { + return TenantQuota{}, fmt.Errorf("parsing tenant quota_json: %w", err) + } + if err := quota.Validate(); err != nil { + return TenantQuota{}, err + } + return quota, nil +} + +// CheckTenantBudget checks one usage estimate against the tenant's quota_json. +func CheckTenantBudget(tenant Tenant, estimate UsageEstimate) (BudgetDecision, error) { + quota, err := ParseTenantQuota(tenant) + if err != nil { + return BudgetDecision{}, err + } + return quota.Check(estimate) +} + +// Validate checks quota limits are non-negative. +func (q TenantQuota) Validate() error { + if q.MaxPromptTokens < 0 { + return fmt.Errorf("max_prompt_tokens must be non-negative") + } + if q.MaxCompletionTokens < 0 { + return fmt.Errorf("max_completion_tokens must be non-negative") + } + if q.MaxTotalTokens < 0 { + return fmt.Errorf("max_total_tokens must be non-negative") + } + if !isFiniteNonNegative(q.MaxCost) { + return fmt.Errorf("max_cost must be finite and non-negative") + } + return nil +} + +// Check applies quota limits to one usage estimate. Zero quota fields are unlimited. +func (q TenantQuota) Check(estimate UsageEstimate) (BudgetDecision, error) { + if err := q.Validate(); err != nil { + return BudgetDecision{}, err + } + if estimate.PromptTokens < 0 || + estimate.CompletionTokens < 0 || + estimate.TotalTokens < 0 { + return BudgetDecision{}, fmt.Errorf("usage estimate values must be non-negative") + } + if !isFiniteNonNegative(estimate.Cost) { + return BudgetDecision{}, fmt.Errorf("usage estimate cost must be finite and non-negative") + } + totalTokens, err := estimate.effectiveTotalTokens() + if err != nil { + return BudgetDecision{}, err + } + if q.MaxPromptTokens > 0 && estimate.PromptTokens > q.MaxPromptTokens { + return BudgetDecision{Reason: "prompt_tokens_exceeded"}, nil + } + if q.MaxCompletionTokens > 0 && estimate.CompletionTokens > q.MaxCompletionTokens { + return BudgetDecision{Reason: "completion_tokens_exceeded"}, nil + } + if q.MaxTotalTokens > 0 && totalTokens > q.MaxTotalTokens { + return BudgetDecision{Reason: "total_tokens_exceeded"}, nil + } + if q.MaxCost > 0 && estimate.Cost > q.MaxCost { + return BudgetDecision{Reason: "cost_exceeded"}, nil + } + return BudgetDecision{Allowed: true}, nil +} + +func (e UsageEstimate) effectiveTotalTokens() (int, error) { + total := e.TotalTokens + if e.PromptTokens > maxInt()-e.CompletionTokens { + return 0, fmt.Errorf("usage estimate total tokens overflow") + } + sum := e.PromptTokens + e.CompletionTokens + if sum > total { + return sum, nil + } + return total, nil +} + +func isFiniteNonNegative(value float64) bool { + return !math.IsNaN(value) && !math.IsInf(value, 0) && value >= 0 +} + +func maxInt() int { + return int(^uint(0) >> 1) +} diff --git a/platform/budget_test.go b/platform/budget_test.go new file mode 100644 index 0000000000..def373d0df --- /dev/null +++ b/platform/budget_test.go @@ -0,0 +1,139 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "math" + "testing" +) + +func TestParseTenantQuotaAllowsEmptyQuota(t *testing.T) { + quota, err := ParseTenantQuota(Tenant{TenantID: "tenant"}) + if err != nil { + t.Fatalf("parse empty quota: %v", err) + } + if quota != (TenantQuota{}) { + t.Fatalf("empty quota should produce zero limits, got %+v", quota) + } +} + +func TestCheckTenantBudgetAllowsWithinQuota(t *testing.T) { + tenant := Tenant{ + TenantID: "tenant", + QuotaJSON: `{"max_prompt_tokens":100,"max_completion_tokens":50,"max_total_tokens":150,"max_cost":1.25}`, + } + + decision, err := CheckTenantBudget(tenant, UsageEstimate{ + PromptTokens: 100, + CompletionTokens: 50, + TotalTokens: 150, + Cost: 1.25, + }) + + if err != nil { + t.Fatalf("check budget: %v", err) + } + if !decision.Allowed || decision.Reason != "" { + t.Fatalf("expected allowed decision, got %+v", decision) + } +} + +func TestCheckTenantBudgetDeniesExceededLimits(t *testing.T) { + tests := []struct { + name string + quota TenantQuota + estimate UsageEstimate + reason string + }{ + { + name: "prompt tokens", + quota: TenantQuota{MaxPromptTokens: 10}, + estimate: UsageEstimate{PromptTokens: 11}, + reason: "prompt_tokens_exceeded", + }, + { + name: "completion tokens", + quota: TenantQuota{MaxCompletionTokens: 10}, + estimate: UsageEstimate{CompletionTokens: 11}, + reason: "completion_tokens_exceeded", + }, + { + name: "total tokens", + quota: TenantQuota{MaxTotalTokens: 10}, + estimate: UsageEstimate{TotalTokens: 11}, + reason: "total_tokens_exceeded", + }, + { + name: "derived total tokens", + quota: TenantQuota{MaxTotalTokens: 10}, + estimate: UsageEstimate{PromptTokens: 6, CompletionTokens: 5}, + reason: "total_tokens_exceeded", + }, + { + name: "cost", + quota: TenantQuota{MaxCost: 1.25}, + estimate: UsageEstimate{Cost: 1.26}, + reason: "cost_exceeded", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + decision, err := tt.quota.Check(tt.estimate) + if err != nil { + t.Fatalf("check quota: %v", err) + } + if decision.Allowed || decision.Reason != tt.reason { + t.Fatalf("expected denied %q, got %+v", tt.reason, decision) + } + }) + } +} + +func TestTenantQuotaRejectsInvalidInputs(t *testing.T) { + _, err := ParseTenantQuota(Tenant{ + TenantID: "tenant", + QuotaJSON: `{"max_total_tokens":`, + }) + if err == nil { + t.Fatalf("expected malformed quota json error") + } + _, err = ParseTenantQuota(Tenant{ + TenantID: "tenant", + QuotaJSON: `{"max_total_tokens":-1}`, + }) + if err == nil { + t.Fatalf("expected negative quota error") + } + _, err = TenantQuota{}.Check(UsageEstimate{Cost: -0.01}) + if err == nil { + t.Fatalf("expected negative usage estimate error") + } +} + +func TestTenantQuotaRejectsNonFiniteCost(t *testing.T) { + _, err := TenantQuota{MaxCost: math.NaN()}.Check(UsageEstimate{}) + if err == nil { + t.Fatalf("expected non-finite quota cost error") + } + _, err = TenantQuota{}.Check(UsageEstimate{Cost: math.Inf(1)}) + if err == nil { + t.Fatalf("expected non-finite usage cost error") + } +} + +func TestTenantQuotaRejectsTokenOverflow(t *testing.T) { + max := int(^uint(0) >> 1) + _, err := TenantQuota{MaxTotalTokens: max}.Check(UsageEstimate{ + PromptTokens: max, + CompletionTokens: 1, + }) + if err == nil { + t.Fatalf("expected total token overflow error") + } +} From 90072df13739c23d9e6aed32c50db04ebde1223b Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 13:42:42 +0800 Subject: [PATCH 10/95] platform: validate storage migration modes --- platform/migration.go | 61 +++++++++++++++++++++++++++++++++ platform/migration_test.go | 69 ++++++++++++++++++++++++++++++++++++++ platform/validation.go | 3 ++ 3 files changed, 133 insertions(+) create mode 100644 platform/migration.go create mode 100644 platform/migration_test.go diff --git a/platform/migration.go b/platform/migration.go new file mode 100644 index 0000000000..bc13dadf00 --- /dev/null +++ b/platform/migration.go @@ -0,0 +1,61 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "fmt" + "strings" +) + +// StorageMigrationMode describes a storage backend migration phase. +type StorageMigrationMode string + +const ( + // StorageMigrationModeNormal means no active migration is in progress. + StorageMigrationModeNormal StorageMigrationMode = "normal" + // StorageMigrationModeDualWrite writes new data to old and new backends. + StorageMigrationModeDualWrite StorageMigrationMode = "dual_write" + // StorageMigrationModeShadowRead compares old and new backend reads. + StorageMigrationModeShadowRead StorageMigrationMode = "shadow_read" + // StorageMigrationModeCutover routes reads and writes to the new backend. + StorageMigrationModeCutover StorageMigrationMode = "cutover" + // StorageMigrationModeRollback routes traffic back to the previous backend. + StorageMigrationModeRollback StorageMigrationMode = "rollback" +) + +// NormalizeStorageMigrationMode returns the canonical migration mode. +func NormalizeStorageMigrationMode(mode string) (StorageMigrationMode, error) { + normalized := StorageMigrationMode(strings.TrimSpace(mode)) + if normalized == "" { + return StorageMigrationModeNormal, nil + } + switch normalized { + case StorageMigrationModeNormal, + StorageMigrationModeDualWrite, + StorageMigrationModeShadowRead, + StorageMigrationModeCutover, + StorageMigrationModeRollback: + return normalized, nil + default: + return "", fmt.Errorf("invalid migration_mode %q", mode) + } +} + +// IsActiveStorageMigrationMode reports whether a mode represents active migration work. +func IsActiveStorageMigrationMode(mode StorageMigrationMode) bool { + switch mode { + case StorageMigrationModeDualWrite, + StorageMigrationModeShadowRead, + StorageMigrationModeCutover, + StorageMigrationModeRollback: + return true + default: + return false + } +} diff --git a/platform/migration_test.go b/platform/migration_test.go new file mode 100644 index 0000000000..e5558b2af4 --- /dev/null +++ b/platform/migration_test.go @@ -0,0 +1,69 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import "testing" + +func TestNormalizeStorageMigrationModeDefaultsEmptyToNormal(t *testing.T) { + mode, err := NormalizeStorageMigrationMode(" ") + if err != nil { + t.Fatalf("normalize empty migration mode: %v", err) + } + if mode != StorageMigrationModeNormal { + t.Fatalf("expected normal mode, got %q", mode) + } +} + +func TestNormalizeStorageMigrationModeAcceptsDocumentedModes(t *testing.T) { + modes := []StorageMigrationMode{ + StorageMigrationModeNormal, + StorageMigrationModeDualWrite, + StorageMigrationModeShadowRead, + StorageMigrationModeCutover, + StorageMigrationModeRollback, + } + for _, want := range modes { + t.Run(string(want), func(t *testing.T) { + got, err := NormalizeStorageMigrationMode(" " + string(want) + " ") + if err != nil { + t.Fatalf("normalize migration mode: %v", err) + } + if got != want { + t.Fatalf("expected %q, got %q", want, got) + } + }) + } +} + +func TestStorageProfileValidateRejectsInvalidMigrationMode(t *testing.T) { + profile := StorageProfile{ + TenantID: "tenant", + ProfileID: "profile", + MigrationMode: "dual-read", + } + if err := profile.Validate(); err == nil { + t.Fatalf("expected invalid migration mode error") + } +} + +func TestIsActiveStorageMigrationMode(t *testing.T) { + if IsActiveStorageMigrationMode(StorageMigrationModeNormal) { + t.Fatalf("normal mode should not be active migration") + } + for _, mode := range []StorageMigrationMode{ + StorageMigrationModeDualWrite, + StorageMigrationModeShadowRead, + StorageMigrationModeCutover, + StorageMigrationModeRollback, + } { + if !IsActiveStorageMigrationMode(mode) { + t.Fatalf("%q should be active migration", mode) + } + } +} diff --git a/platform/validation.go b/platform/validation.go index 7ed5ef6440..2aeab76472 100644 --- a/platform/validation.go +++ b/platform/validation.go @@ -164,6 +164,9 @@ func (p StorageProfile) Validate() error { if err := validateSecretReference("dsn_ref", p.DSNRef); err != nil { return err } + if _, err := NormalizeStorageMigrationMode(p.MigrationMode); err != nil { + return err + } return nil } From 2b8667910ee509591624acb744780b4425766fe3 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 13:52:35 +0800 Subject: [PATCH 11/95] platform: validate audit policies --- platform/audit_policy_test.go | 112 ++++++++++++++++++++++++++++++++++ platform/validation.go | 21 +++++++ 2 files changed, 133 insertions(+) create mode 100644 platform/audit_policy_test.go diff --git a/platform/audit_policy_test.go b/platform/audit_policy_test.go new file mode 100644 index 0000000000..086e59f60f --- /dev/null +++ b/platform/audit_policy_test.go @@ -0,0 +1,112 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "errors" + "math" + "strings" + "testing" +) + +func TestAuditPolicyValidateAcceptsValidPolicy(t *testing.T) { + policy := AuditPolicy{ + TenantID: "tenant", + PolicyID: "audit-policy", + RetentionDays: 30, + SampleRate: 0.25, + FullAuditForRiskyTool: true, + RedactionRules: []string{`(?i)(session_id=)[^\s]+`, " "}, + ExportSink: "audit://sink", + ComplianceLevel: "standard", + } + if err := policy.Validate(); err != nil { + t.Fatalf("expected valid audit policy, got %v", err) + } +} + +func TestAuditPolicyValidateRequiresTenant(t *testing.T) { + policy := validAuditPolicy() + policy.TenantID = " " + if err := policy.Validate(); !errors.Is(err, ErrTenantIDRequired) { + t.Fatalf("expected tenant requirement, got %v", err) + } +} + +func TestAuditPolicyValidateRequiresPolicyID(t *testing.T) { + policy := validAuditPolicy() + policy.PolicyID = " " + if err := policy.Validate(); err == nil || !strings.Contains(err.Error(), "policy_id is required") { + t.Fatalf("expected policy_id requirement, got %v", err) + } +} + +func TestAuditPolicyValidateRejectsNegativeRetention(t *testing.T) { + policy := validAuditPolicy() + policy.RetentionDays = -1 + if err := policy.Validate(); err == nil || !strings.Contains(err.Error(), "retention_days") { + t.Fatalf("expected retention_days validation, got %v", err) + } +} + +func TestAuditPolicyValidateRejectsInvalidSampleRate(t *testing.T) { + tests := []struct { + name string + sampleRate float64 + }{ + {name: "negative", sampleRate: -0.01}, + {name: "above_one", sampleRate: 1.01}, + {name: "nan", sampleRate: math.NaN()}, + {name: "positive_inf", sampleRate: math.Inf(1)}, + {name: "negative_inf", sampleRate: math.Inf(-1)}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + policy := validAuditPolicy() + policy.SampleRate = tt.sampleRate + if err := policy.Validate(); err == nil || !strings.Contains(err.Error(), "sample_rate") { + t.Fatalf("expected sample_rate validation, got %v", err) + } + }) + } +} + +func TestAuditPolicyValidateRejectsInvalidRedactionRule(t *testing.T) { + policy := validAuditPolicy() + policy.RedactionRules = []string{"["} + if err := policy.Validate(); err == nil || !strings.Contains(err.Error(), "redaction_rules") { + t.Fatalf("expected redaction_rules validation, got %v", err) + } +} + +func TestAuditPolicyValidateAcceptsBlankRedactionRule(t *testing.T) { + policy := validAuditPolicy() + policy.RedactionRules = []string{" "} + if err := policy.Validate(); err != nil { + t.Fatalf("expected blank redaction rule to be ignored, got %v", err) + } +} + +func TestAuditPolicyValidateAcceptsZeroSampleRate(t *testing.T) { + policy := validAuditPolicy() + policy.SampleRate = 0 + if err := policy.Validate(); err != nil { + t.Fatalf("expected zero sample rate to be accepted, got %v", err) + } +} + +func validAuditPolicy() AuditPolicy { + return AuditPolicy{ + TenantID: "tenant", + PolicyID: "audit-policy", + RetentionDays: 30, + SampleRate: 1, + } +} diff --git a/platform/validation.go b/platform/validation.go index 2aeab76472..ed98336051 100644 --- a/platform/validation.go +++ b/platform/validation.go @@ -10,6 +10,7 @@ package platform import ( "fmt" + "math" "strconv" "strings" ) @@ -170,6 +171,26 @@ func (p StorageProfile) Validate() error { return nil } +// Validate checks that audit retention and sampling policy is safe to use. +func (p AuditPolicy) Validate() error { + if strings.TrimSpace(p.TenantID) == "" { + return ErrTenantIDRequired + } + if strings.TrimSpace(p.PolicyID) == "" { + return fmt.Errorf("policy_id is required") + } + if p.RetentionDays < 0 { + return fmt.Errorf("retention_days must be greater than or equal to 0") + } + if math.IsNaN(p.SampleRate) || math.IsInf(p.SampleRate, 0) || p.SampleRate < 0 || p.SampleRate > 1 { + return fmt.Errorf("sample_rate must be between 0 and 1") + } + if _, err := NewRedactor(p.RedactionRules...); err != nil { + return fmt.Errorf("redaction_rules: %w", err) + } + return nil +} + func validateSecretReference(field, value string) error { value = strings.TrimSpace(value) if value == "" { From 08fbc12906c9d75142bfde68a4be1aaa986e0c63 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 14:02:00 +0800 Subject: [PATCH 12/95] platform: validate audit records --- platform/audit_record_test.go | 117 ++++++++++++++++++++++++++++++++++ platform/validation.go | 44 +++++++++++++ 2 files changed, 161 insertions(+) create mode 100644 platform/audit_record_test.go diff --git a/platform/audit_record_test.go b/platform/audit_record_test.go new file mode 100644 index 0000000000..f31de83dd7 --- /dev/null +++ b/platform/audit_record_test.go @@ -0,0 +1,117 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "errors" + "math" + "strings" + "testing" +) + +func TestAuditRecordValidateAcceptsSafeRecord(t *testing.T) { + record := validAuditRecord() + record.DecisionReason = "tool approved by policy" + record.TokenUsageJSON = `{"prompt_tokens":10,"completion_tokens":5}` + record.RedactedDetailRef = "sha256:0123456789abcdef bytes:128" + + if err := record.Validate(); err != nil { + t.Fatalf("expected valid audit record, got %v", err) + } +} + +func TestAuditRecordValidateRequiresTenant(t *testing.T) { + record := validAuditRecord() + record.TenantID = " " + if err := record.Validate(); !errors.Is(err, ErrTenantIDRequired) { + t.Fatalf("expected tenant requirement, got %v", err) + } +} + +func TestAuditRecordValidateRequiresAuditID(t *testing.T) { + record := validAuditRecord() + record.AuditID = " " + if err := record.Validate(); err == nil || !strings.Contains(err.Error(), "audit_id is required") { + t.Fatalf("expected audit_id requirement, got %v", err) + } +} + +func TestAuditRecordValidateRejectsNegativeLatency(t *testing.T) { + record := validAuditRecord() + record.LatencyMS = -1 + if err := record.Validate(); err == nil || !strings.Contains(err.Error(), "latency_ms") { + t.Fatalf("expected latency validation, got %v", err) + } +} + +func TestAuditRecordValidateRejectsInvalidCost(t *testing.T) { + tests := []struct { + name string + cost float64 + }{ + {name: "negative", cost: -0.01}, + {name: "nan", cost: math.NaN()}, + {name: "positive_inf", cost: math.Inf(1)}, + {name: "negative_inf", cost: math.Inf(-1)}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + record := validAuditRecord() + record.Cost = tt.cost + if err := record.Validate(); err == nil || !strings.Contains(err.Error(), "cost") { + t.Fatalf("expected cost validation, got %v", err) + } + }) + } +} + +func TestAuditRecordValidateRejectsSensitiveDecisionReason(t *testing.T) { + record := validAuditRecord() + record.DecisionReason = "Authorization: Bearer raw-token" + if err := record.Validate(); err == nil || !strings.Contains(err.Error(), "decision_reason") { + t.Fatalf("expected sensitive decision reason rejection, got %v", err) + } +} + +func TestAuditRecordValidateRejectsSensitiveErrorType(t *testing.T) { + record := validAuditRecord() + record.ErrorType = "storage_error password=plain" + if err := record.Validate(); err == nil || !strings.Contains(err.Error(), "error_type") { + t.Fatalf("expected sensitive error type rejection, got %v", err) + } +} + +func TestAuditRecordValidateRejectsSensitiveTokenUsage(t *testing.T) { + record := validAuditRecord() + record.TokenUsageJSON = `{"api_key":"sk-1234567890abcdef"}` + if err := record.Validate(); err == nil || !strings.Contains(err.Error(), "token_usage_json") { + t.Fatalf("expected sensitive token usage rejection, got %v", err) + } +} + +func TestAuditRecordValidateRejectsSensitiveDetailRef(t *testing.T) { + record := validAuditRecord() + record.RedactedDetailRef = "postgres://user:password@example.com/db" + if err := record.Validate(); err == nil || !strings.Contains(err.Error(), "redacted_detail_ref") { + t.Fatalf("expected sensitive detail rejection, got %v", err) + } +} + +func validAuditRecord() AuditRecord { + return AuditRecord{ + TenantID: "tenant", + AuditID: "audit", + UserID: "internal-user", + InternalUserID: "usr", + UserIDHash: UserIDHash("tenant", "telegram", "external"), + TraceID: "trace", + Decision: "allow", + } +} diff --git a/platform/validation.go b/platform/validation.go index ed98336051..52c0b1faf2 100644 --- a/platform/validation.go +++ b/platform/validation.go @@ -191,6 +191,50 @@ func (p AuditPolicy) Validate() error { return nil } +// Validate checks that an audit record has required identity and no raw secret detail. +func (r AuditRecord) Validate() error { + if strings.TrimSpace(r.TenantID) == "" { + return ErrTenantIDRequired + } + if strings.TrimSpace(r.AuditID) == "" { + return fmt.Errorf("audit_id is required") + } + if r.LatencyMS < 0 { + return fmt.Errorf("latency_ms must be greater than or equal to 0") + } + if math.IsNaN(r.Cost) || math.IsInf(r.Cost, 0) || r.Cost < 0 { + return fmt.Errorf("cost must be greater than or equal to 0") + } + if err := validateAuditRedactedText("decision_reason", r.DecisionReason); err != nil { + return err + } + if err := validateAuditRedactedText("error_type", r.ErrorType); err != nil { + return err + } + if err := validateAuditRedactedText("token_usage_json", r.TokenUsageJSON); err != nil { + return err + } + if err := validateAuditRedactedText("redacted_detail_ref", r.RedactedDetailRef); err != nil { + return err + } + return nil +} + +func validateAuditRedactedText(field, value string) error { + value = strings.TrimSpace(value) + if value == "" { + return nil + } + redactor, err := NewRedactor() + if err != nil { + return fmt.Errorf("%s: redactor unavailable: %w", field, err) + } + if redactor.Redact(value) != value { + return fmt.Errorf("%s contains unredacted sensitive content", field) + } + return nil +} + func validateSecretReference(field, value string) error { value = strings.TrimSpace(value) if value == "" { From b63b970e5813f4078d9fb4eab746bba3a617739b Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 14:07:39 +0800 Subject: [PATCH 13/95] platform: add capacity estimator --- platform/capacity.go | 113 +++++++++++++++++++++++++++++++++++ platform/capacity_test.go | 120 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 233 insertions(+) create mode 100644 platform/capacity.go create mode 100644 platform/capacity_test.go diff --git a/platform/capacity.go b/platform/capacity.go new file mode 100644 index 0000000000..01f3052dfb --- /dev/null +++ b/platform/capacity.go @@ -0,0 +1,113 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "fmt" + "math" +) + +// CapacityInputs captures tenant-level planning assumptions. +type CapacityInputs struct { + DAU int + MessagesPerUserPeak float64 + PeakFactor float64 + PeakWindowSeconds float64 + AverageRunnerLatencySec float64 + TargetUtilization float64 + AverageEventsPerRun float64 + RequestsPerDay int + AveragePromptTokens float64 + AverageCompletionTokens float64 + InputTokenPricePerToken float64 + OutputTokenPricePerToken float64 + AverageToolCostPerRequest float64 +} + +// CapacityEstimate summarizes rough capacity and cost signals for one tenant. +type CapacityEstimate struct { + CallbackQPS float64 + WorkerConcurrency float64 + SessionReadQPS float64 + SessionWriteQPS float64 + TokensPerDay float64 + CostPerDay float64 +} + +// EstimateCapacity applies the platform capacity formulas to tenant inputs. +func EstimateCapacity(input CapacityInputs) (CapacityEstimate, error) { + if err := input.Validate(); err != nil { + return CapacityEstimate{}, err + } + callbackQPS := float64(input.DAU) * input.MessagesPerUserPeak * input.PeakFactor / input.PeakWindowSeconds + workerConcurrency := callbackQPS * input.AverageRunnerLatencySec / input.TargetUtilization + sessionWriteQPS := callbackQPS * input.AverageEventsPerRun + requestsPerDay := float64(input.RequestsPerDay) + tokensPerRequest := input.AveragePromptTokens + input.AverageCompletionTokens + tokensPerDay := requestsPerDay * tokensPerRequest + costPerDay := requestsPerDay * ((input.AveragePromptTokens * input.InputTokenPricePerToken) + + (input.AverageCompletionTokens * input.OutputTokenPricePerToken) + + input.AverageToolCostPerRequest) + return CapacityEstimate{ + CallbackQPS: callbackQPS, + WorkerConcurrency: workerConcurrency, + SessionReadQPS: callbackQPS, + SessionWriteQPS: sessionWriteQPS, + TokensPerDay: tokensPerDay, + CostPerDay: costPerDay, + }, nil +} + +// Validate checks capacity assumptions before applying estimation formulas. +func (i CapacityInputs) Validate() error { + if i.DAU < 0 { + return fmt.Errorf("dau must be non-negative") + } + if i.RequestsPerDay < 0 { + return fmt.Errorf("requests_per_day must be non-negative") + } + if !isFiniteNonNegative(i.MessagesPerUserPeak) { + return fmt.Errorf("messages_per_user_peak must be finite and non-negative") + } + if !isFiniteNonNegative(i.PeakFactor) { + return fmt.Errorf("peak_factor must be finite and non-negative") + } + if !isFinitePositive(i.PeakWindowSeconds) { + return fmt.Errorf("peak_window_seconds must be finite and greater than 0") + } + if !isFiniteNonNegative(i.AverageRunnerLatencySec) { + return fmt.Errorf("average_runner_latency_sec must be finite and non-negative") + } + if !isFinitePositive(i.TargetUtilization) || i.TargetUtilization > 1 { + return fmt.Errorf("target_utilization must be finite and between 0 and 1") + } + if !isFiniteNonNegative(i.AverageEventsPerRun) { + return fmt.Errorf("average_events_per_run must be finite and non-negative") + } + if !isFiniteNonNegative(i.AveragePromptTokens) { + return fmt.Errorf("average_prompt_tokens must be finite and non-negative") + } + if !isFiniteNonNegative(i.AverageCompletionTokens) { + return fmt.Errorf("average_completion_tokens must be finite and non-negative") + } + if !isFiniteNonNegative(i.InputTokenPricePerToken) { + return fmt.Errorf("input_token_price_per_token must be finite and non-negative") + } + if !isFiniteNonNegative(i.OutputTokenPricePerToken) { + return fmt.Errorf("output_token_price_per_token must be finite and non-negative") + } + if !isFiniteNonNegative(i.AverageToolCostPerRequest) { + return fmt.Errorf("average_tool_cost_per_request must be finite and non-negative") + } + return nil +} + +func isFinitePositive(value float64) bool { + return !math.IsNaN(value) && !math.IsInf(value, 0) && value > 0 +} diff --git a/platform/capacity_test.go b/platform/capacity_test.go new file mode 100644 index 0000000000..0ead87a11e --- /dev/null +++ b/platform/capacity_test.go @@ -0,0 +1,120 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "math" + "strings" + "testing" +) + +func TestEstimateCapacityAppliesDesignFormulas(t *testing.T) { + estimate, err := EstimateCapacity(CapacityInputs{ + DAU: 1000, + MessagesPerUserPeak: 3, + PeakFactor: 4, + PeakWindowSeconds: 600, + AverageRunnerLatencySec: 8, + TargetUtilization: 0.8, + AverageEventsPerRun: 5, + RequestsPerDay: 20000, + AveragePromptTokens: 1000, + AverageCompletionTokens: 250, + InputTokenPricePerToken: 0.000001, + OutputTokenPricePerToken: 0.000002, + AverageToolCostPerRequest: 0.001, + }) + if err != nil { + t.Fatalf("EstimateCapacity: %v", err) + } + assertFloat(t, "CallbackQPS", estimate.CallbackQPS, 20) + assertFloat(t, "WorkerConcurrency", estimate.WorkerConcurrency, 200) + assertFloat(t, "SessionReadQPS", estimate.SessionReadQPS, 20) + assertFloat(t, "SessionWriteQPS", estimate.SessionWriteQPS, 100) + assertFloat(t, "TokensPerDay", estimate.TokensPerDay, 25_000_000) + assertFloat(t, "CostPerDay", estimate.CostPerDay, 50) +} + +func TestEstimateCapacityAllowsZeroDemand(t *testing.T) { + estimate, err := EstimateCapacity(CapacityInputs{ + PeakWindowSeconds: 1, + TargetUtilization: 1, + }) + if err != nil { + t.Fatalf("EstimateCapacity: %v", err) + } + if estimate != (CapacityEstimate{}) { + t.Fatalf("expected zero estimate, got %+v", estimate) + } +} + +func TestCapacityInputsRejectInvalidValues(t *testing.T) { + tests := []struct { + name string + input CapacityInputs + field string + }{ + { + name: "negative dau", + input: CapacityInputs{DAU: -1, PeakWindowSeconds: 1, TargetUtilization: 1}, + field: "dau", + }, + { + name: "zero peak window", + input: CapacityInputs{PeakWindowSeconds: 0, TargetUtilization: 1}, + field: "peak_window_seconds", + }, + { + name: "zero utilization", + input: CapacityInputs{PeakWindowSeconds: 1, TargetUtilization: 0}, + field: "target_utilization", + }, + { + name: "utilization above one", + input: CapacityInputs{PeakWindowSeconds: 1, TargetUtilization: 1.01}, + field: "target_utilization", + }, + { + name: "nan peak factor", + input: CapacityInputs{PeakWindowSeconds: 1, TargetUtilization: 1, PeakFactor: math.NaN()}, + field: "peak_factor", + }, + { + name: "infinite latency", + input: CapacityInputs{PeakWindowSeconds: 1, TargetUtilization: 1, AverageRunnerLatencySec: math.Inf(1)}, + field: "average_runner_latency_sec", + }, + { + name: "negative requests", + input: CapacityInputs{PeakWindowSeconds: 1, TargetUtilization: 1, RequestsPerDay: -1}, + field: "requests_per_day", + }, + { + name: "negative token price", + input: CapacityInputs{PeakWindowSeconds: 1, TargetUtilization: 1, InputTokenPricePerToken: -0.01}, + field: "input_token_price_per_token", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := EstimateCapacity(tt.input) + if err == nil || !strings.Contains(err.Error(), tt.field) { + t.Fatalf("expected %s validation error, got %v", tt.field, err) + } + }) + } +} + +func assertFloat(t *testing.T, name string, got, want float64) { + t.Helper() + if math.Abs(got-want) > 1e-9 { + t.Fatalf("%s = %v, want %v", name, got, want) + } +} From 41935692c870752bbe159288d10ec939477d3db2 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 14:16:32 +0800 Subject: [PATCH 14/95] platform/channeladapter: replay dead letters --- platform/channeladapter/errors.go | 2 + platform/channeladapter/outbox.go | 65 +++++++++++++++++ platform/channeladapter/outbox_test.go | 97 ++++++++++++++++++++++++++ 3 files changed, 164 insertions(+) diff --git a/platform/channeladapter/errors.go b/platform/channeladapter/errors.go index 4a45fdb4be..9bad4a9dd3 100644 --- a/platform/channeladapter/errors.go +++ b/platform/channeladapter/errors.go @@ -27,4 +27,6 @@ var ( ErrOutboundLeaseExpired = errors.New("channel adapter outbound lease expired") // ErrInvalidDeliveryStatus indicates that a provider returned an invalid status. ErrInvalidDeliveryStatus = errors.New("channel adapter invalid delivery status") + // ErrOutboundReplayNotDeadLetter indicates that only dead-letter records can be replayed. + ErrOutboundReplayNotDeadLetter = errors.New("channel adapter outbound replay requires dead letter") ) diff --git a/platform/channeladapter/outbox.go b/platform/channeladapter/outbox.go index 76a26d877a..0ca18a5e11 100644 --- a/platform/channeladapter/outbox.go +++ b/platform/channeladapter/outbox.go @@ -94,6 +94,12 @@ type OutboxStore interface { MarkDeadLetter(ctx context.Context, dedupKey string, leaseToken string, err error, now time.Time) (OutboxRecord, error) } +// DeadLetterOutboxStore exposes admin operations for dead-letter inspection and replay. +type DeadLetterOutboxStore interface { + ListDeadLetters(ctx context.Context, tenantID string, limit int) ([]OutboxRecord, error) + RequeueDeadLetter(ctx context.Context, dedupKey string, policy RetryPolicy, now time.Time) (OutboxRecord, error) +} + // InMemoryOutboxStore is a concurrency-safe outbox store for tests and demos. type InMemoryOutboxStore struct { mu sync.Mutex @@ -188,6 +194,33 @@ func (s *InMemoryOutboxStore) ListDue( return out, nil } +// ListDeadLetters returns dead-lettered records, optionally scoped by tenant. +func (s *InMemoryOutboxStore) ListDeadLetters( + ctx context.Context, + tenantID string, + limit int, +) ([]OutboxRecord, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + s.mu.Lock() + defer s.mu.Unlock() + out := make([]OutboxRecord, 0) + for _, record := range s.records { + if record.Status != platform.OutboundStatusDeadLetter { + continue + } + if tenantID != "" && record.Message.TenantID != tenantID { + continue + } + out = append(out, record) + if limit > 0 && len(out) >= limit { + break + } + } + return out, nil +} + // ClaimDue atomically leases due records for one dispatcher worker. func (s *InMemoryOutboxStore) ClaimDue( ctx context.Context, @@ -227,6 +260,38 @@ func (s *InMemoryOutboxStore) ClaimDue( return out, nil } +// RequeueDeadLetter moves one dead-lettered record back to pending delivery. +func (s *InMemoryOutboxStore) RequeueDeadLetter( + ctx context.Context, + dedupKey string, + policy RetryPolicy, + now time.Time, +) (OutboxRecord, error) { + if err := ctx.Err(); err != nil { + return OutboxRecord{}, err + } + s.mu.Lock() + defer s.mu.Unlock() + record, ok := s.records[dedupKey] + if !ok { + return OutboxRecord{}, ErrOutboundNotFound + } + if record.Status != platform.OutboundStatusDeadLetter { + return OutboxRecord{}, ErrOutboundReplayNotDeadLetter + } + record.Status = platform.OutboundStatusPending + record.Attempts = 0 + record.MaxAttempts = maxAttempts(policy) + record.RetryPolicy = normalizeRetryPolicy(policy) + record.NextAttemptAt = now + record.LeaseToken = "" + record.LeaseExpiresAt = time.Time{} + record.LastError = "" + record.UpdatedAt = now + s.records[dedupKey] = record + return record, nil +} + // MarkSent marks an outbox record as delivered. func (s *InMemoryOutboxStore) MarkSent( ctx context.Context, diff --git a/platform/channeladapter/outbox_test.go b/platform/channeladapter/outbox_test.go index 889cd25a6f..f9ee7c8584 100644 --- a/platform/channeladapter/outbox_test.go +++ b/platform/channeladapter/outbox_test.go @@ -158,6 +158,103 @@ func TestOutboxFailureSchedulesRetryThenDeadLetter(t *testing.T) { } } +func TestOutboxListsAndRequeuesDeadLetter(t *testing.T) { + ctx := context.Background() + store := NewInMemoryOutboxStore() + msg := outbound("reply-1") + policy := RetryPolicy{MaxAttempts: 1, InitialBackoff: time.Second, MaxBackoff: time.Second} + _, _, err := store.Enqueue(ctx, msg, policy) + if err != nil { + t.Fatalf("enqueue: %v", err) + } + now := time.Now().Add(time.Hour) + claimed, err := store.ClaimDue(ctx, now, 1, time.Minute) + if err != nil { + t.Fatalf("claim due: %v", err) + } + _, err = store.MarkDeadLetter(ctx, msg.DedupKey, claimed[0].LeaseToken, errors.New("permanent"), now) + if err != nil { + t.Fatalf("mark dead letter: %v", err) + } + + dead, err := store.ListDeadLetters(ctx, "tenant", 10) + if err != nil { + t.Fatalf("ListDeadLetters: %v", err) + } + if len(dead) != 1 || dead[0].Message.DedupKey != msg.DedupKey { + t.Fatalf("unexpected dead letters: %+v", dead) + } + otherTenant, err := store.ListDeadLetters(ctx, "other-tenant", 10) + if err != nil { + t.Fatalf("ListDeadLetters other tenant: %v", err) + } + if len(otherTenant) != 0 { + t.Fatalf("dead letter listing should respect tenant scope: %+v", otherTenant) + } + + requeued, err := store.RequeueDeadLetter(ctx, msg.DedupKey, RetryPolicy{ + MaxAttempts: 3, + InitialBackoff: 2 * time.Second, + MaxBackoff: 10 * time.Second, + }, now.Add(time.Minute)) + if err != nil { + t.Fatalf("RequeueDeadLetter: %v", err) + } + if requeued.Status != platform.OutboundStatusPending || + requeued.Attempts != 0 || + requeued.MaxAttempts != 3 || + requeued.LastError != "" || + !requeued.NextAttemptAt.Equal(now.Add(time.Minute)) { + t.Fatalf("unexpected requeued record: %+v", requeued) + } + due, err := store.ClaimDue(ctx, now.Add(time.Minute), 1, time.Minute) + if err != nil { + t.Fatalf("claim requeued: %v", err) + } + if len(due) != 1 || due[0].Message.DedupKey != msg.DedupKey { + t.Fatalf("requeued record should be due: %+v", due) + } +} + +func TestOutboxRequeueRejectsNonDeadLetter(t *testing.T) { + ctx := context.Background() + store := NewInMemoryOutboxStore() + msg := outbound("reply-1") + _, _, err := store.Enqueue(ctx, msg, DefaultRetryPolicy()) + if err != nil { + t.Fatalf("enqueue: %v", err) + } + + record, err := store.RequeueDeadLetter(ctx, msg.DedupKey, DefaultRetryPolicy(), time.Now()) + if !errors.Is(err, ErrOutboundReplayNotDeadLetter) { + t.Fatalf("expected replay rejection, got record=%+v err=%v", record, err) + } +} + +func TestOutboxRequeueRejectsSentRecord(t *testing.T) { + ctx := context.Background() + store := NewInMemoryOutboxStore() + msg := outbound("reply-1") + _, _, err := store.Enqueue(ctx, msg, DefaultRetryPolicy()) + if err != nil { + t.Fatalf("enqueue: %v", err) + } + now := time.Now().Add(time.Hour) + claimed, err := store.ClaimDue(ctx, now, 1, time.Minute) + if err != nil { + t.Fatalf("claim due: %v", err) + } + _, err = store.MarkSent(ctx, msg.DedupKey, claimed[0].LeaseToken, "provider-1", now) + if err != nil { + t.Fatalf("mark sent: %v", err) + } + + record, err := store.RequeueDeadLetter(ctx, msg.DedupKey, DefaultRetryPolicy(), now) + if !errors.Is(err, ErrOutboundReplayNotDeadLetter) { + t.Fatalf("expected replay rejection, got record=%+v err=%v", record, err) + } +} + func TestDispatcherMarksSent(t *testing.T) { ctx := context.Background() store := NewInMemoryOutboxStore() From 78675efe7e5b26697cd18600a1066a324ca92d6c Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 14:27:50 +0800 Subject: [PATCH 15/95] platform: populate audit record ids --- platform/gateway/service.go | 1 + platform/gateway/service_test.go | 2 + platform/identity.go | 5 +++ platform/toolpolicy/policy.go | 1 + platform/toolpolicy/policy_test.go | 64 ++++++++++++++++++++++++++++++ platform/types_test.go | 15 +++++++ 6 files changed, 88 insertions(+) diff --git a/platform/gateway/service.go b/platform/gateway/service.go index 60a74b7356..01e09f01b3 100644 --- a/platform/gateway/service.go +++ b/platform/gateway/service.go @@ -379,6 +379,7 @@ func auditFromMessage( err error, ) platform.AuditRecord { record := platform.AuditRecord{ + AuditID: platform.AuditID(msg.TenantID, msg.AppID, msg.Channel, msg.BindingID, msg.PlatformMessageID, sessionID, decision), TenantID: msg.TenantID, AppID: msg.AppID, Channel: msg.Channel, diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go index 45a9971522..72f5afacbc 100644 --- a/platform/gateway/service_test.go +++ b/platform/gateway/service_test.go @@ -164,6 +164,7 @@ func TestServiceHandleInboundOutboxFailureDoesNotCompleteIdempotency(t *testing. require.True(t, ok) assert.Equal(t, "queued", stored.Content) require.Len(t, audit.Records(), 1) + assert.NotEmpty(t, audit.Records()[0].AuditID) assert.Equal(t, "outbound_error", audit.Records()[0].Decision) } @@ -378,6 +379,7 @@ func TestServiceHandleInboundRejectsUnsupportedMessage(t *testing.T) { require.ErrorIs(t, err, ErrUnsupportedMessageType) assert.Empty(t, r.calls) require.Len(t, audit.Records(), 1) + assert.NotEmpty(t, audit.Records()[0].AuditID) assert.Equal(t, "reject", audit.Records()[0].Decision) assert.NotEqual(t, "user-1", audit.Records()[0].UserID) } diff --git a/platform/identity.go b/platform/identity.go index aab8af1728..9ae3916a83 100644 --- a/platform/identity.go +++ b/platform/identity.go @@ -26,6 +26,11 @@ func UserIDHash(tenantID, channel, userID string) string { return "user_hash_" + shortHash(tenantID, channel, userID) } +// AuditID returns a stable audit identifier for one audit event boundary. +func AuditID(parts ...string) string { + return "audit_" + shortHash(parts...) +} + // IdempotencyKey returns the canonical duplicate-delivery key. func IdempotencyKey(tenantID, channel, accountID, platformMessageID string) string { return strings.Join([]string{ diff --git a/platform/toolpolicy/policy.go b/platform/toolpolicy/policy.go index cef6374be9..2a9804ae8c 100644 --- a/platform/toolpolicy/policy.go +++ b/platform/toolpolicy/policy.go @@ -391,6 +391,7 @@ func (p *Policy) writeAudit( } argsSummary := argumentSummary(req.Arguments) _ = p.audit.WriteAudit(ctx, platform.AuditRecord{ + AuditID: platform.AuditID(p.policy.TenantID, p.policy.AppID, toolName, req.ToolCallID, decision, argsSummary), TenantID: p.policy.TenantID, AppID: p.policy.AppID, ToolName: toolName, diff --git a/platform/toolpolicy/policy_test.go b/platform/toolpolicy/policy_test.go index 3289a0b035..0edc55496e 100644 --- a/platform/toolpolicy/policy_test.go +++ b/platform/toolpolicy/policy_test.go @@ -104,6 +104,7 @@ func TestPolicyAllowsHighRiskWithAuditAndRedactsArguments(t *testing.T) { } record := records[0] if record.Decision != string(tool.PermissionActionAllow) || + record.AuditID == "" || record.ToolName != "http_post" || record.TenantID != "tenant" || record.AppID != "app" || @@ -122,6 +123,63 @@ func TestPolicyAllowsHighRiskWithAuditAndRedactsArguments(t *testing.T) { } } +func TestPolicyAuditIDUsesToolCallBoundary(t *testing.T) { + audit := platform.NewInMemoryAuditSink() + p := newPolicy( + t, + platform.ToolPolicy{ + TenantID: "tenant", + AppID: "app", + DangerousToolAction: platform.DangerousToolActionAllowWithAudit, + HighRiskTools: []string{"http_post"}, + }, + WithAuditSink(audit), + ) + args := []byte(`{"url":"https://example.com"}`) + req1 := request("http_post", tool.ToolMetadata{}, args) + req1.ToolCallID = "call-1" + req2 := request("http_post", tool.ToolMetadata{}, args) + req2.ToolCallID = "call-2" + + if _, err := p.CheckToolPermission(context.Background(), req1); err != nil { + t.Fatalf("CheckToolPermission req1: %v", err) + } + if _, err := p.CheckToolPermission(context.Background(), req2); err != nil { + t.Fatalf("CheckToolPermission req2: %v", err) + } + records := audit.Records() + if len(records) != 2 { + t.Fatalf("expected two audit records, got %d", len(records)) + } + if records[0].AuditID == records[1].AuditID { + t.Fatalf("expected tool-call-scoped audit ids, got %q", records[0].AuditID) + } + + retryAudit := platform.NewInMemoryAuditSink() + retryPolicy := newPolicy( + t, + platform.ToolPolicy{ + TenantID: "tenant", + AppID: "app", + DangerousToolAction: platform.DangerousToolActionAllowWithAudit, + HighRiskTools: []string{"http_post"}, + }, + WithAuditSink(retryAudit), + ) + retryReq := request("http_post", tool.ToolMetadata{}, args) + retryReq.ToolCallID = "call-1" + if _, err := retryPolicy.CheckToolPermission(context.Background(), retryReq); err != nil { + t.Fatalf("CheckToolPermission retry: %v", err) + } + retryRecords := retryAudit.Records() + if len(retryRecords) != 1 { + t.Fatalf("expected one retry audit record, got %d", len(retryRecords)) + } + if records[0].AuditID != retryRecords[0].AuditID { + t.Fatalf("expected stable audit id for same tool call, got %q and %q", records[0].AuditID, retryRecords[0].AuditID) + } +} + func TestPolicyNilRedactorStillDoesNotLeakArguments(t *testing.T) { audit := platform.NewInMemoryAuditSink() p := newPolicy( @@ -145,6 +203,9 @@ func TestPolicyNilRedactorStillDoesNotLeakArguments(t *testing.T) { if len(records) != 1 { t.Fatalf("expected one audit record, got %d", len(records)) } + if records[0].AuditID == "" { + t.Fatalf("expected audit id") + } if strings.Contains(records[0].RedactedDetailRef, "person@example.com") || strings.Contains(records[0].RedactedDetailRef, "/private/file") { t.Fatalf("audit leaked raw argument content: %q", records[0].RedactedDetailRef) @@ -243,6 +304,9 @@ func TestPolicyRegisterAppliesNameBasedGovernance(t *testing.T) { if result == nil || result.CustomResult == nil { t.Fatalf("expected approval-required result") } + if len(audit.Records()) != 1 || audit.Records()[0].AuditID == "" { + t.Fatalf("expected audit record with id, got %+v", audit.Records()) + } permissionResult, ok := result.CustomResult.(tool.PermissionResult) if !ok { t.Fatalf("expected tool.PermissionResult, got %T", result.CustomResult) diff --git a/platform/types_test.go b/platform/types_test.go index 79c741b501..e8f077fd70 100644 --- a/platform/types_test.go +++ b/platform/types_test.go @@ -91,6 +91,21 @@ func TestInternalUserIDIsStableAndTenantScoped(t *testing.T) { } } +func TestAuditIDIsStableAndScoped(t *testing.T) { + a1 := AuditID("tenant-a", "app", "trace", "decision") + a2 := AuditID("tenant-a", "app", "trace", "decision") + b := AuditID("tenant-b", "app", "trace", "decision") + if a1 != a2 { + t.Fatalf("expected stable audit id, got %q and %q", a1, a2) + } + if a1 == b { + t.Fatalf("expected scoped audit ids, got %q", a1) + } + if !strings.HasPrefix(a1, "audit_") { + t.Fatalf("expected audit id prefix, got %q", a1) + } +} + func TestIdempotencyStoreDoesNotRestartCompletedMessage(t *testing.T) { store := NewInMemoryIdempotencyStore() record := IdempotencyRecord{ From 41b6b97285d133dd46605bd6c66c824c051d30d4 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 14:36:07 +0800 Subject: [PATCH 16/95] platform: validate audit sink writes --- platform/audit.go | 3 +++ platform/toolpolicy/policy_test.go | 2 ++ platform/types_test.go | 36 ++++++++++++++++++++++++++++++ 3 files changed, 41 insertions(+) diff --git a/platform/audit.go b/platform/audit.go index deba348c82..921cd6f5b9 100644 --- a/platform/audit.go +++ b/platform/audit.go @@ -35,6 +35,9 @@ func (s *InMemoryAuditSink) WriteAudit(ctx context.Context, record AuditRecord) if err := ctx.Err(); err != nil { return err } + if err := record.Validate(); err != nil { + return err + } s.mu.Lock() defer s.mu.Unlock() s.records = append(s.records, record) diff --git a/platform/toolpolicy/policy_test.go b/platform/toolpolicy/policy_test.go index 0edc55496e..9baa26b207 100644 --- a/platform/toolpolicy/policy_test.go +++ b/platform/toolpolicy/policy_test.go @@ -185,6 +185,8 @@ func TestPolicyNilRedactorStillDoesNotLeakArguments(t *testing.T) { p := newPolicy( t, platform.ToolPolicy{ + TenantID: "tenant", + AppID: "app", DangerousToolAction: platform.DangerousToolActionAllowWithAudit, HighRiskTools: []string{"http_post"}, }, diff --git a/platform/types_test.go b/platform/types_test.go index e8f077fd70..2aee01424f 100644 --- a/platform/types_test.go +++ b/platform/types_test.go @@ -282,3 +282,39 @@ func TestAuditSinkStoresSnapshot(t *testing.T) { t.Fatalf("Records should return a defensive copy") } } + +func TestAuditSinkRejectsInvalidRecord(t *testing.T) { + sink := NewInMemoryAuditSink() + record := AuditRecord{ + TenantID: "tenant", + } + + err := sink.WriteAudit(context.Background(), record) + if err == nil || !strings.Contains(err.Error(), "audit_id is required") { + t.Fatalf("expected audit_id validation, got %v", err) + } + if got := sink.Records(); len(got) != 0 { + t.Fatalf("expected invalid record to be rejected, got %+v", got) + } +} + +func TestAuditSinkRejectsSensitiveRecord(t *testing.T) { + sink := NewInMemoryAuditSink() + record := AuditRecord{ + TenantID: "tenant", + AuditID: "audit", + UserID: "internal", + InternalUserID: "usr", + UserIDHash: UserIDHash("tenant", "telegram", "external"), + TraceID: "trace", + DecisionReason: "Authorization: Bearer raw-token", + } + + err := sink.WriteAudit(context.Background(), record) + if err == nil || !strings.Contains(err.Error(), "decision_reason") { + t.Fatalf("expected sensitive record validation, got %v", err) + } + if got := sink.Records(); len(got) != 0 { + t.Fatalf("expected sensitive record to be rejected, got %+v", got) + } +} From b72383f2df8e6fb7c46d75ee49f9e66f382cfc4a Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 14:40:56 +0800 Subject: [PATCH 17/95] platform: add usage record contracts --- platform/types.go | 20 ++++++ platform/usage_record_test.go | 111 ++++++++++++++++++++++++++++++++++ platform/validation.go | 40 ++++++++++++ 3 files changed, 171 insertions(+) create mode 100644 platform/usage_record_test.go diff --git a/platform/types.go b/platform/types.go index ca0cfba02f..17007e46be 100644 --- a/platform/types.go +++ b/platform/types.go @@ -393,3 +393,23 @@ type AuditRecord struct { RedactionVersion string CreatedAt time.Time } + +// UsageRecord stores post-run token and cost accounting dimensions. +type UsageRecord struct { + TenantID string + AppID string + UserIDHash string + SessionID string + RequestID string + ModelName string + ToolName string + PromptTokens int + CompletionTokens int + CachedTokens int + ModelUnitPrice float64 + ModelCost float64 + ToolCost float64 + TotalCost float64 + TraceID string + CreatedAt time.Time +} diff --git a/platform/usage_record_test.go b/platform/usage_record_test.go new file mode 100644 index 0000000000..977c6b904a --- /dev/null +++ b/platform/usage_record_test.go @@ -0,0 +1,111 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "errors" + "math" + "strings" + "testing" +) + +func TestUsageRecordValidateAcceptsSafeRecord(t *testing.T) { + record := validUsageRecord() + record.PromptTokens = 100 + record.CompletionTokens = 50 + record.CachedTokens = 10 + record.ModelUnitPrice = 0.00001 + record.ModelCost = 0.0015 + record.ToolCost = 0.25 + record.TotalCost = 0.2515 + + if err := record.Validate(); err != nil { + t.Fatalf("expected valid usage record, got %v", err) + } +} + +func TestUsageRecordValidateRequiresTenantAndApp(t *testing.T) { + record := validUsageRecord() + record.TenantID = " " + if err := record.Validate(); !errors.Is(err, ErrTenantIDRequired) { + t.Fatalf("expected tenant requirement, got %v", err) + } + + record = validUsageRecord() + record.AppID = " " + if err := record.Validate(); !errors.Is(err, ErrAppIDRequired) { + t.Fatalf("expected app requirement, got %v", err) + } +} + +func TestUsageRecordValidateRejectsNegativeTokens(t *testing.T) { + tests := []struct { + name string + mutate func(*UsageRecord) + }{ + {name: "prompt", mutate: func(r *UsageRecord) { r.PromptTokens = -1 }}, + {name: "completion", mutate: func(r *UsageRecord) { r.CompletionTokens = -1 }}, + {name: "cached", mutate: func(r *UsageRecord) { r.CachedTokens = -1 }}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + record := validUsageRecord() + tt.mutate(&record) + if err := record.Validate(); err == nil || !strings.Contains(err.Error(), "token") { + t.Fatalf("expected token validation, got %v", err) + } + }) + } +} + +func TestUsageRecordValidateRejectsInvalidCosts(t *testing.T) { + tests := []struct { + name string + mutate func(*UsageRecord) + field string + }{ + {name: "model unit price negative", mutate: func(r *UsageRecord) { r.ModelUnitPrice = -0.01 }, field: "model_unit_price"}, + {name: "model unit price nan", mutate: func(r *UsageRecord) { r.ModelUnitPrice = math.NaN() }, field: "model_unit_price"}, + {name: "model cost negative", mutate: func(r *UsageRecord) { r.ModelCost = -0.01 }, field: "model_cost"}, + {name: "tool cost infinite", mutate: func(r *UsageRecord) { r.ToolCost = math.Inf(1) }, field: "tool_cost"}, + {name: "total cost negative", mutate: func(r *UsageRecord) { r.TotalCost = -0.01 }, field: "total_cost"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + record := validUsageRecord() + tt.mutate(&record) + if err := record.Validate(); err == nil || !strings.Contains(err.Error(), tt.field) { + t.Fatalf("expected %s validation, got %v", tt.field, err) + } + }) + } +} + +func TestUsageRecordValidateRejectsSensitiveDimensions(t *testing.T) { + record := validUsageRecord() + record.ToolName = "http_post Authorization: Bearer raw-token" + if err := record.Validate(); err == nil || !strings.Contains(err.Error(), "tool_name") { + t.Fatalf("expected sensitive tool name rejection, got %v", err) + } +} + +func validUsageRecord() UsageRecord { + return UsageRecord{ + TenantID: "tenant", + AppID: "app", + UserIDHash: UserIDHash("tenant", "telegram", "external"), + SessionID: "session", + RequestID: "request", + ModelName: "gpt-test", + ToolName: "knowledge_search", + TraceID: "trace", + } +} diff --git a/platform/validation.go b/platform/validation.go index 52c0b1faf2..662ee6c54f 100644 --- a/platform/validation.go +++ b/platform/validation.go @@ -220,6 +220,46 @@ func (r AuditRecord) Validate() error { return nil } +// Validate checks that a usage record has required identity and safe accounting values. +func (r UsageRecord) Validate() error { + if strings.TrimSpace(r.TenantID) == "" { + return ErrTenantIDRequired + } + if strings.TrimSpace(r.AppID) == "" { + return ErrAppIDRequired + } + for field, value := range map[string]string{ + "user_id_hash": r.UserIDHash, + "session_id": r.SessionID, + "request_id": r.RequestID, + "model_name": r.ModelName, + "tool_name": r.ToolName, + "trace_id": r.TraceID, + } { + if err := validateAuditRedactedText(field, value); err != nil { + return err + } + } + if r.PromptTokens < 0 || + r.CompletionTokens < 0 || + r.CachedTokens < 0 { + return fmt.Errorf("usage token values must be non-negative") + } + if !isFiniteNonNegative(r.ModelUnitPrice) { + return fmt.Errorf("model_unit_price must be finite and non-negative") + } + if !isFiniteNonNegative(r.ModelCost) { + return fmt.Errorf("model_cost must be finite and non-negative") + } + if !isFiniteNonNegative(r.ToolCost) { + return fmt.Errorf("tool_cost must be finite and non-negative") + } + if !isFiniteNonNegative(r.TotalCost) { + return fmt.Errorf("total_cost must be finite and non-negative") + } + return nil +} + func validateAuditRedactedText(field, value string) error { value = strings.TrimSpace(value) if value == "" { From 260b3a57ed738641e4f11eb825f1f7de5d50db60 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 14:46:02 +0800 Subject: [PATCH 18/95] platform: add usage sink contracts --- platform/usage.go | 54 +++++++++++++++++++++++++++++++++++ platform/usage_record_test.go | 47 ++++++++++++++++++++++++++++++ 2 files changed, 101 insertions(+) create mode 100644 platform/usage.go diff --git a/platform/usage.go b/platform/usage.go new file mode 100644 index 0000000000..037fc7d59b --- /dev/null +++ b/platform/usage.go @@ -0,0 +1,54 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "context" + "sync" +) + +// UsageSink stores post-run usage records. +type UsageSink interface { + // WriteUsage writes one usage record. + WriteUsage(ctx context.Context, record UsageRecord) error +} + +// InMemoryUsageSink is a concurrency-safe usage sink for tests and demos. +type InMemoryUsageSink struct { + mu sync.Mutex + records []UsageRecord +} + +// NewInMemoryUsageSink creates an in-memory usage sink. +func NewInMemoryUsageSink() *InMemoryUsageSink { + return &InMemoryUsageSink{} +} + +// WriteUsage writes one usage record. +func (s *InMemoryUsageSink) WriteUsage(ctx context.Context, record UsageRecord) error { + if err := ctx.Err(); err != nil { + return err + } + if err := record.Validate(); err != nil { + return err + } + s.mu.Lock() + defer s.mu.Unlock() + s.records = append(s.records, record) + return nil +} + +// Records returns a snapshot of written usage records. +func (s *InMemoryUsageSink) Records() []UsageRecord { + s.mu.Lock() + defer s.mu.Unlock() + out := make([]UsageRecord, len(s.records)) + copy(out, s.records) + return out +} diff --git a/platform/usage_record_test.go b/platform/usage_record_test.go index 977c6b904a..d3837674df 100644 --- a/platform/usage_record_test.go +++ b/platform/usage_record_test.go @@ -9,6 +9,7 @@ package platform import ( + "context" "errors" "math" "strings" @@ -97,6 +98,52 @@ func TestUsageRecordValidateRejectsSensitiveDimensions(t *testing.T) { } } +func TestUsageSinkStoresSnapshot(t *testing.T) { + sink := NewInMemoryUsageSink() + record := validUsageRecord() + record.TotalCost = 0.25 + + if err := sink.WriteUsage(context.Background(), record); err != nil { + t.Fatalf("WriteUsage: %v", err) + } + records := sink.Records() + if len(records) != 1 { + t.Fatalf("expected one usage record, got %d", len(records)) + } + records[0].TenantID = "changed" + if sink.Records()[0].TenantID != "tenant" { + t.Fatalf("Records should return a defensive copy") + } +} + +func TestUsageSinkRejectsInvalidRecord(t *testing.T) { + sink := NewInMemoryUsageSink() + record := validUsageRecord() + record.PromptTokens = -1 + + err := sink.WriteUsage(context.Background(), record) + if err == nil || !strings.Contains(err.Error(), "token") { + t.Fatalf("expected token validation, got %v", err) + } + if got := sink.Records(); len(got) != 0 { + t.Fatalf("expected invalid record to be rejected, got %+v", got) + } +} + +func TestUsageSinkRejectsSensitiveRecord(t *testing.T) { + sink := NewInMemoryUsageSink() + record := validUsageRecord() + record.ModelName = "gpt-test api_key=sk-1234567890abcdef" + + err := sink.WriteUsage(context.Background(), record) + if err == nil || !strings.Contains(err.Error(), "model_name") { + t.Fatalf("expected sensitive record validation, got %v", err) + } + if got := sink.Records(); len(got) != 0 { + t.Fatalf("expected sensitive record to be rejected, got %+v", got) + } +} + func validUsageRecord() UsageRecord { return UsageRecord{ TenantID: "tenant", From c62e2149664e4631ab7fa463b6c5be3a73c90980 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 14:52:37 +0800 Subject: [PATCH 19/95] platform: add config version contracts --- platform/config_version_test.go | 115 ++++++++++++++++++++++++++++++++ platform/types.go | 30 +++++++++ platform/validation.go | 82 +++++++++++++++++++++++ 3 files changed, 227 insertions(+) create mode 100644 platform/config_version_test.go diff --git a/platform/config_version_test.go b/platform/config_version_test.go new file mode 100644 index 0000000000..25457511ee --- /dev/null +++ b/platform/config_version_test.go @@ -0,0 +1,115 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "errors" + "strings" + "testing" +) + +func TestAppConfigVersionValidateAcceptsValidVersion(t *testing.T) { + version := validAppConfigVersion() + version.ConfigBundleJSON = `{"model_profile_id":"model","tool_policy_id":"tools","api_key_ref":"secret://model-key"}` + + if err := version.Validate(); err != nil { + t.Fatalf("expected valid config version, got %v", err) + } +} + +func TestAppConfigVersionValidateRequiresIdentity(t *testing.T) { + version := validAppConfigVersion() + version.TenantID = " " + if err := version.Validate(); !errors.Is(err, ErrTenantIDRequired) { + t.Fatalf("expected tenant requirement, got %v", err) + } + + version = validAppConfigVersion() + version.AppID = " " + if err := version.Validate(); !errors.Is(err, ErrAppIDRequired) { + t.Fatalf("expected app requirement, got %v", err) + } + + version = validAppConfigVersion() + version.Version = " " + if err := version.Validate(); err == nil || !strings.Contains(err.Error(), "version is required") { + t.Fatalf("expected version requirement, got %v", err) + } +} + +func TestAppConfigVersionValidateRequiresBundleAndChecksum(t *testing.T) { + version := validAppConfigVersion() + version.ConfigBundleJSON = " " + if err := version.Validate(); err == nil || !strings.Contains(err.Error(), "config_bundle_json") { + t.Fatalf("expected config bundle requirement, got %v", err) + } + + version = validAppConfigVersion() + version.ConfigBundleJSON = `{"model_profile_id":` + if err := version.Validate(); err == nil || !strings.Contains(err.Error(), "valid json") { + t.Fatalf("expected config bundle json validation, got %v", err) + } + + version = validAppConfigVersion() + version.Checksum = " " + if err := version.Validate(); err == nil || !strings.Contains(err.Error(), "checksum") { + t.Fatalf("expected checksum requirement, got %v", err) + } +} + +func TestAppConfigVersionValidateRejectsUnsafeBundle(t *testing.T) { + version := validAppConfigVersion() + version.ConfigBundleJSON = `{"api_key":"sk-1234567890abcdef"}` + + if err := version.Validate(); err == nil || !strings.Contains(err.Error(), "config_bundle_json") { + t.Fatalf("expected sensitive bundle rejection, got %v", err) + } +} + +func TestAppConfigVersionValidateRejectsUnsafeRefValue(t *testing.T) { + version := validAppConfigVersion() + version.ConfigBundleJSON = `{"api_key_ref":"sk-1234567890abcdef"}` + + if err := version.Validate(); !errors.Is(err, ErrInlineSecretRejected) { + t.Fatalf("expected inline secret reference rejection, got %v", err) + } +} + +func TestAppConfigVersionValidateRejectsInvalidStatusAndGrayPercent(t *testing.T) { + version := validAppConfigVersion() + version.Status = "" + if err := version.Validate(); err == nil || !strings.Contains(err.Error(), "status is required") { + t.Fatalf("expected status requirement, got %v", err) + } + + version = validAppConfigVersion() + version.Status = AppConfigVersionStatus("paused") + if err := version.Validate(); err == nil || !strings.Contains(err.Error(), "invalid app config version status") { + t.Fatalf("expected invalid status, got %v", err) + } + + version = validAppConfigVersion() + version.GrayPercent = 101 + if err := version.Validate(); err == nil || !strings.Contains(err.Error(), "gray_percent") { + t.Fatalf("expected gray percent validation, got %v", err) + } +} + +func validAppConfigVersion() AppConfigVersion { + return AppConfigVersion{ + TenantID: "tenant", + AppID: "app", + Version: "v1", + ConfigBundleJSON: `{"model_profile_id":"model","tool_policy_id":"tools"}`, + Checksum: "sha256:0123456789abcdef", + Status: AppConfigVersionStatusDraft, + GrayPercent: 10, + CreatedBy: "operator", + } +} diff --git a/platform/types.go b/platform/types.go index 17007e46be..2a57b51f5c 100644 --- a/platform/types.go +++ b/platform/types.go @@ -34,6 +34,22 @@ const ( AppStatusDeleted AppStatus = "deleted" ) +// AppConfigVersionStatus is the lifecycle state of one app configuration version. +type AppConfigVersionStatus string + +const ( + // AppConfigVersionStatusDraft is editable and not ready for traffic. + AppConfigVersionStatusDraft AppConfigVersionStatus = "draft" + // AppConfigVersionStatusValidated passed offline validation. + AppConfigVersionStatusValidated AppConfigVersionStatus = "validated" + // AppConfigVersionStatusReleased is eligible for gray traffic. + AppConfigVersionStatusReleased AppConfigVersionStatus = "released" + // AppConfigVersionStatusActive receives normal traffic. + AppConfigVersionStatusActive AppConfigVersionStatus = "active" + // AppConfigVersionStatusRollback is retained as the rollback target. + AppConfigVersionStatusRollback AppConfigVersionStatus = "rollback" +) + // BindingStatus is the lifecycle state of a channel binding. type BindingStatus string @@ -190,6 +206,20 @@ type AgentApp struct { UpdatedAt time.Time } +// AppConfigVersion stores one deployable app configuration bundle. +type AppConfigVersion struct { + TenantID string + AppID string + Version string + ConfigBundleJSON string + Checksum string + Status AppConfigVersionStatus + GrayPercent int + CreatedBy string + CreatedAt time.Time + ActivatedAt *time.Time +} + // ModelProfile stores model provider configuration references. type ModelProfile struct { TenantID string diff --git a/platform/validation.go b/platform/validation.go index 662ee6c54f..9a9258db19 100644 --- a/platform/validation.go +++ b/platform/validation.go @@ -9,6 +9,7 @@ package platform import ( + "encoding/json" "fmt" "math" "strconv" @@ -57,6 +58,87 @@ func (a AgentApp) Validate() error { } } +// Validate checks that an app config version is safe to store and route. +func (v AppConfigVersion) Validate() error { + if strings.TrimSpace(v.TenantID) == "" { + return ErrTenantIDRequired + } + if strings.TrimSpace(v.AppID) == "" { + return ErrAppIDRequired + } + if strings.TrimSpace(v.Version) == "" { + return fmt.Errorf("version is required") + } + if strings.TrimSpace(v.ConfigBundleJSON) == "" { + return fmt.Errorf("config_bundle_json is required") + } + if !json.Valid([]byte(v.ConfigBundleJSON)) { + return fmt.Errorf("config_bundle_json must be valid json") + } + if err := validateConfigBundleJSON(v.ConfigBundleJSON); err != nil { + return err + } + if strings.TrimSpace(v.Checksum) == "" { + return fmt.Errorf("checksum is required") + } + if err := validateAuditRedactedText("checksum", v.Checksum); err != nil { + return err + } + if v.GrayPercent < 0 || v.GrayPercent > 100 { + return fmt.Errorf("gray_percent must be between 0 and 100") + } + switch v.Status { + case AppConfigVersionStatusDraft, + AppConfigVersionStatusValidated, + AppConfigVersionStatusReleased, + AppConfigVersionStatusActive, + AppConfigVersionStatusRollback: + return nil + case "": + return fmt.Errorf("status is required") + default: + return fmt.Errorf("invalid app config version status %q", v.Status) + } +} + +func validateConfigBundleJSON(bundle string) error { + var value any + if err := json.Unmarshal([]byte(bundle), &value); err != nil { + return fmt.Errorf("config_bundle_json must be valid json") + } + return validateConfigBundleValue("config_bundle_json", "", value) +} + +func validateConfigBundleValue(path, key string, value any) error { + switch typed := value.(type) { + case map[string]any: + for childKey, childValue := range typed { + childPath := path + "." + childKey + if err := validateConfigBundleValue(childPath, childKey, childValue); err != nil { + return err + } + } + case []any: + for i, childValue := range typed { + childPath := fmt.Sprintf("%s[%d]", path, i) + if err := validateConfigBundleValue(childPath, key, childValue); err != nil { + return err + } + } + case string: + if strings.HasSuffix(strings.ToLower(strings.TrimSpace(key)), "_ref") { + if err := validateSecretReference(path, typed); err != nil { + return err + } + return nil + } + if err := validateAuditRedactedText(path, typed); err != nil { + return err + } + } + return nil +} + // Validate checks that model profile sensitive values are stored by reference. func (p ModelProfile) Validate() error { if strings.TrimSpace(p.TenantID) == "" { From a74ee58758328c1bac609e2976a92fd3c6496674 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 14:58:09 +0800 Subject: [PATCH 20/95] platform: select config version by session gray bucket --- platform/gray.go | 39 ++++++++++++++ platform/gray_test.go | 120 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 159 insertions(+) diff --git a/platform/gray.go b/platform/gray.go index 41e7e1ca05..cf7cdaa1e6 100644 --- a/platform/gray.go +++ b/platform/gray.go @@ -14,6 +14,13 @@ import ( "strings" ) +// ConfigVersionSelection is the deterministic routing decision for one session. +type ConfigVersionSelection struct { + Version AppConfigVersion + Bucket int + InCandidate bool +} + // SessionGrayBucket returns the stable 0-99 release bucket for one session. func SessionGrayBucket(tenantID, appID, sessionID string) (int, error) { tenantID = strings.TrimSpace(tenantID) @@ -48,3 +55,35 @@ func SessionInGrayRelease(app AgentApp, sessionID string) (bool, int, error) { } return bucket < app.GrayPercent, bucket, nil } + +// SelectAppConfigVersionForSession chooses the active or released gray config version for one session. +func SelectAppConfigVersionForSession(active, candidate AppConfigVersion, sessionID string) (ConfigVersionSelection, error) { + var selection ConfigVersionSelection + if err := active.Validate(); err != nil { + return selection, fmt.Errorf("active config version: %w", err) + } + if active.Status != AppConfigVersionStatusActive { + return selection, fmt.Errorf("active config version status must be active") + } + if err := candidate.Validate(); err != nil { + return selection, fmt.Errorf("candidate config version: %w", err) + } + if candidate.Status != AppConfigVersionStatusReleased { + return selection, fmt.Errorf("candidate config version status must be released") + } + if strings.TrimSpace(active.TenantID) != strings.TrimSpace(candidate.TenantID) { + return selection, fmt.Errorf("candidate tenant_id must match active tenant_id") + } + if strings.TrimSpace(active.AppID) != strings.TrimSpace(candidate.AppID) { + return selection, fmt.Errorf("candidate app_id must match active app_id") + } + + bucket, err := SessionGrayBucket(active.TenantID, active.AppID, sessionID) + if err != nil { + return selection, err + } + if bucket < candidate.GrayPercent { + return ConfigVersionSelection{Version: candidate, Bucket: bucket, InCandidate: true}, nil + } + return ConfigVersionSelection{Version: active, Bucket: bucket}, nil +} diff --git a/platform/gray_test.go b/platform/gray_test.go index c4d49b8a58..425ee314da 100644 --- a/platform/gray_test.go +++ b/platform/gray_test.go @@ -75,3 +75,123 @@ func TestSessionInGrayReleaseRejectsInvalidInputs(t *testing.T) { t.Fatalf("expected missing session id error") } } + +func TestSelectAppConfigVersionForSessionIsStable(t *testing.T) { + active := validGrayActiveConfigVersion() + candidate := validGrayCandidateConfigVersion() + candidate.GrayPercent = 100 + + first, err := SelectAppConfigVersionForSession(active, candidate, "session-1") + if err != nil { + t.Fatalf("select config version: %v", err) + } + second, err := SelectAppConfigVersionForSession(active, candidate, " session-1 ") + if err != nil { + t.Fatalf("select config version: %v", err) + } + if first.Bucket != second.Bucket { + t.Fatalf("same session should use stable bucket: %d != %d", first.Bucket, second.Bucket) + } + if !first.InCandidate || first.Version.Version != candidate.Version { + t.Fatalf("expected candidate version, got %+v", first) + } +} + +func TestSelectAppConfigVersionForSessionUsesGrayPercentBoundaries(t *testing.T) { + active := validGrayActiveConfigVersion() + candidate := validGrayCandidateConfigVersion() + + candidate.GrayPercent = 0 + selection, err := SelectAppConfigVersionForSession(active, candidate, "session-1") + if err != nil { + t.Fatalf("select config version: %v", err) + } + if selection.InCandidate || selection.Version.Version != active.Version { + t.Fatalf("0 percent should choose active, got %+v", selection) + } + + candidate.GrayPercent = 100 + selection, err = SelectAppConfigVersionForSession(active, candidate, "session-1") + if err != nil { + t.Fatalf("select config version: %v", err) + } + if !selection.InCandidate || selection.Version.Version != candidate.Version { + t.Fatalf("100 percent should choose candidate, got %+v", selection) + } +} + +func TestSelectAppConfigVersionForSessionMatchesBucketThreshold(t *testing.T) { + active := validGrayActiveConfigVersion() + candidate := validGrayCandidateConfigVersion() + candidate.GrayPercent = 50 + + selection, err := SelectAppConfigVersionForSession(active, candidate, "session-1") + if err != nil { + t.Fatalf("select config version: %v", err) + } + if selection.InCandidate != (selection.Bucket < candidate.GrayPercent) { + t.Fatalf("selection should match bucket threshold: in_candidate=%t bucket=%d percent=%d", + selection.InCandidate, selection.Bucket, candidate.GrayPercent) + } + if selection.InCandidate && selection.Version.Version != candidate.Version { + t.Fatalf("candidate hit should return candidate version, got %+v", selection) + } + if !selection.InCandidate && selection.Version.Version != active.Version { + t.Fatalf("active hit should return active version, got %+v", selection) + } +} + +func TestSelectAppConfigVersionForSessionRejectsMismatchedIdentity(t *testing.T) { + active := validGrayActiveConfigVersion() + candidate := validGrayCandidateConfigVersion() + candidate.TenantID = "other-tenant" + if _, err := SelectAppConfigVersionForSession(active, candidate, "session-1"); err == nil { + t.Fatalf("expected tenant mismatch error") + } + + candidate = validGrayCandidateConfigVersion() + candidate.AppID = "other-app" + if _, err := SelectAppConfigVersionForSession(active, candidate, "session-1"); err == nil { + t.Fatalf("expected app mismatch error") + } +} + +func TestSelectAppConfigVersionForSessionRejectsInvalidStatuses(t *testing.T) { + active := validGrayActiveConfigVersion() + candidate := validGrayCandidateConfigVersion() + + active.Status = AppConfigVersionStatusReleased + if _, err := SelectAppConfigVersionForSession(active, candidate, "session-1"); err == nil { + t.Fatalf("expected active status error") + } + + active = validGrayActiveConfigVersion() + candidate.Status = AppConfigVersionStatusValidated + if _, err := SelectAppConfigVersionForSession(active, candidate, "session-1"); err == nil { + t.Fatalf("expected candidate status error") + } +} + +func TestSelectAppConfigVersionForSessionRejectsMissingSession(t *testing.T) { + active := validGrayActiveConfigVersion() + candidate := validGrayCandidateConfigVersion() + if _, err := SelectAppConfigVersionForSession(active, candidate, " "); err == nil { + t.Fatalf("expected missing session id error") + } +} + +func validGrayActiveConfigVersion() AppConfigVersion { + version := validAppConfigVersion() + version.Version = "v1" + version.Status = AppConfigVersionStatusActive + version.GrayPercent = 0 + return version +} + +func validGrayCandidateConfigVersion() AppConfigVersion { + version := validAppConfigVersion() + version.Version = "v2" + version.Status = AppConfigVersionStatusReleased + version.GrayPercent = 10 + return version +} From 75d1af703eb84cc67a9920ed872774b0433dcce6 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 15:03:11 +0800 Subject: [PATCH 21/95] platform: add config version lifecycle helpers --- platform/config_lifecycle.go | 100 +++++++++++++++++++ platform/config_lifecycle_test.go | 156 ++++++++++++++++++++++++++++++ 2 files changed, 256 insertions(+) create mode 100644 platform/config_lifecycle.go create mode 100644 platform/config_lifecycle_test.go diff --git a/platform/config_lifecycle.go b/platform/config_lifecycle.go new file mode 100644 index 0000000000..13644d6d7e --- /dev/null +++ b/platform/config_lifecycle.go @@ -0,0 +1,100 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "fmt" + "strings" + "time" +) + +// ReleaseAppConfigVersion promotes a validated app config version to gray-release candidate status. +func ReleaseAppConfigVersion(version AppConfigVersion, grayPercent int) (AppConfigVersion, error) { + if err := version.Validate(); err != nil { + return AppConfigVersion{}, err + } + if version.Status != AppConfigVersionStatusValidated { + return AppConfigVersion{}, fmt.Errorf("config version status must be validated before release") + } + if grayPercent < 0 || grayPercent > 100 { + return AppConfigVersion{}, fmt.Errorf("gray_percent must be between 0 and 100") + } + version.Status = AppConfigVersionStatusReleased + version.GrayPercent = grayPercent + version.ActivatedAt = nil + return version, nil +} + +// ActivateAppConfigVersion makes a released app config version the active version and retains the previous active version for rollback. +func ActivateAppConfigVersion(active, released AppConfigVersion, activatedAt time.Time) (AppConfigVersion, AppConfigVersion, error) { + if err := active.Validate(); err != nil { + return AppConfigVersion{}, AppConfigVersion{}, fmt.Errorf("active config version: %w", err) + } + if active.Status != AppConfigVersionStatusActive { + return AppConfigVersion{}, AppConfigVersion{}, fmt.Errorf("active config version status must be active") + } + if err := released.Validate(); err != nil { + return AppConfigVersion{}, AppConfigVersion{}, fmt.Errorf("released config version: %w", err) + } + if released.Status != AppConfigVersionStatusReleased { + return AppConfigVersion{}, AppConfigVersion{}, fmt.Errorf("released config version status must be released") + } + if err := requireSameConfigOwner(active, released); err != nil { + return AppConfigVersion{}, AppConfigVersion{}, err + } + + rollback := active + rollback.Status = AppConfigVersionStatusRollback + rollback.GrayPercent = 0 + + nextActive := released + nextActive.Status = AppConfigVersionStatusActive + nextActive.GrayPercent = 0 + nextActive.ActivatedAt = &activatedAt + return nextActive, rollback, nil +} + +// RollbackAppConfigVersion makes a rollback version active and retains the replaced version as rollback. +func RollbackAppConfigVersion(active, rollback AppConfigVersion, activatedAt time.Time) (AppConfigVersion, AppConfigVersion, error) { + if err := active.Validate(); err != nil { + return AppConfigVersion{}, AppConfigVersion{}, fmt.Errorf("active config version: %w", err) + } + if active.Status != AppConfigVersionStatusActive { + return AppConfigVersion{}, AppConfigVersion{}, fmt.Errorf("active config version status must be active") + } + if err := rollback.Validate(); err != nil { + return AppConfigVersion{}, AppConfigVersion{}, fmt.Errorf("rollback config version: %w", err) + } + if rollback.Status != AppConfigVersionStatusRollback { + return AppConfigVersion{}, AppConfigVersion{}, fmt.Errorf("rollback config version status must be rollback") + } + if err := requireSameConfigOwner(active, rollback); err != nil { + return AppConfigVersion{}, AppConfigVersion{}, err + } + + previousActive := active + previousActive.Status = AppConfigVersionStatusRollback + previousActive.GrayPercent = 0 + + nextActive := rollback + nextActive.Status = AppConfigVersionStatusActive + nextActive.GrayPercent = 0 + nextActive.ActivatedAt = &activatedAt + return nextActive, previousActive, nil +} + +func requireSameConfigOwner(left, right AppConfigVersion) error { + if strings.TrimSpace(left.TenantID) != strings.TrimSpace(right.TenantID) { + return fmt.Errorf("config version tenant_id must match") + } + if strings.TrimSpace(left.AppID) != strings.TrimSpace(right.AppID) { + return fmt.Errorf("config version app_id must match") + } + return nil +} diff --git a/platform/config_lifecycle_test.go b/platform/config_lifecycle_test.go new file mode 100644 index 0000000000..fa4f8469a9 --- /dev/null +++ b/platform/config_lifecycle_test.go @@ -0,0 +1,156 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "strings" + "testing" + "time" +) + +func TestReleaseAppConfigVersionPromotesValidatedCandidate(t *testing.T) { + version := validAppConfigVersion() + version.Status = AppConfigVersionStatusValidated + version.GrayPercent = 0 + now := time.Now() + version.ActivatedAt = &now + + released, err := ReleaseAppConfigVersion(version, 25) + if err != nil { + t.Fatalf("release config version: %v", err) + } + if released.Status != AppConfigVersionStatusReleased { + t.Fatalf("expected released status, got %q", released.Status) + } + if released.GrayPercent != 25 { + t.Fatalf("expected gray percent 25, got %d", released.GrayPercent) + } + if released.ActivatedAt != nil { + t.Fatalf("released candidate should not carry activated_at") + } +} + +func TestReleaseAppConfigVersionRejectsInvalidTransition(t *testing.T) { + version := validAppConfigVersion() + version.Status = AppConfigVersionStatusDraft + if _, err := ReleaseAppConfigVersion(version, 10); err == nil || + !strings.Contains(err.Error(), "validated") { + t.Fatalf("expected validated status requirement, got %v", err) + } + + version.Status = AppConfigVersionStatusValidated + if _, err := ReleaseAppConfigVersion(version, -1); err == nil || + !strings.Contains(err.Error(), "gray_percent") { + t.Fatalf("expected gray percent validation, got %v", err) + } + if _, err := ReleaseAppConfigVersion(version, 101); err == nil || + !strings.Contains(err.Error(), "gray_percent") { + t.Fatalf("expected gray percent validation, got %v", err) + } +} + +func TestActivateAppConfigVersionPromotesReleasedAndKeepsRollback(t *testing.T) { + active := validLifecycleConfigVersion("v1", AppConfigVersionStatusActive) + released := validLifecycleConfigVersion("v2", AppConfigVersionStatusReleased) + released.GrayPercent = 50 + activatedAt := time.Date(2026, 7, 8, 10, 0, 0, 0, time.UTC) + + nextActive, rollback, err := ActivateAppConfigVersion(active, released, activatedAt) + if err != nil { + t.Fatalf("activate config version: %v", err) + } + if nextActive.Version != "v2" || nextActive.Status != AppConfigVersionStatusActive { + t.Fatalf("expected released version to become active, got %+v", nextActive) + } + if nextActive.GrayPercent != 0 { + t.Fatalf("active version should reset gray percent, got %d", nextActive.GrayPercent) + } + if nextActive.ActivatedAt == nil || !nextActive.ActivatedAt.Equal(activatedAt) { + t.Fatalf("active version should record activation time, got %v", nextActive.ActivatedAt) + } + if rollback.Version != "v1" || rollback.Status != AppConfigVersionStatusRollback { + t.Fatalf("expected previous active to become rollback, got %+v", rollback) + } + if rollback.GrayPercent != 0 { + t.Fatalf("rollback version should not receive gray traffic, got %d", rollback.GrayPercent) + } +} + +func TestActivateAppConfigVersionRejectsInvalidTransitions(t *testing.T) { + active := validLifecycleConfigVersion("v1", AppConfigVersionStatusReleased) + released := validLifecycleConfigVersion("v2", AppConfigVersionStatusReleased) + if _, _, err := ActivateAppConfigVersion(active, released, time.Now()); err == nil || + !strings.Contains(err.Error(), "active") { + t.Fatalf("expected active status requirement, got %v", err) + } + + active = validLifecycleConfigVersion("v1", AppConfigVersionStatusActive) + released = validLifecycleConfigVersion("v2", AppConfigVersionStatusValidated) + if _, _, err := ActivateAppConfigVersion(active, released, time.Now()); err == nil || + !strings.Contains(err.Error(), "released") { + t.Fatalf("expected released status requirement, got %v", err) + } + + released = validLifecycleConfigVersion("v2", AppConfigVersionStatusReleased) + released.TenantID = "other-tenant" + if _, _, err := ActivateAppConfigVersion(active, released, time.Now()); err == nil || + !strings.Contains(err.Error(), "tenant_id") { + t.Fatalf("expected tenant mismatch error, got %v", err) + } +} + +func TestRollbackAppConfigVersionPromotesRollbackAndRetainsCurrent(t *testing.T) { + active := validLifecycleConfigVersion("v2", AppConfigVersionStatusActive) + rollback := validLifecycleConfigVersion("v1", AppConfigVersionStatusRollback) + activatedAt := time.Date(2026, 7, 8, 11, 0, 0, 0, time.UTC) + + nextActive, previousActive, err := RollbackAppConfigVersion(active, rollback, activatedAt) + if err != nil { + t.Fatalf("rollback config version: %v", err) + } + if nextActive.Version != "v1" || nextActive.Status != AppConfigVersionStatusActive { + t.Fatalf("expected rollback version to become active, got %+v", nextActive) + } + if nextActive.ActivatedAt == nil || !nextActive.ActivatedAt.Equal(activatedAt) { + t.Fatalf("rollback activation should record activation time, got %v", nextActive.ActivatedAt) + } + if previousActive.Version != "v2" || previousActive.Status != AppConfigVersionStatusRollback { + t.Fatalf("expected replaced active to become rollback, got %+v", previousActive) + } +} + +func TestRollbackAppConfigVersionRejectsInvalidTransitions(t *testing.T) { + active := validLifecycleConfigVersion("v2", AppConfigVersionStatusReleased) + rollback := validLifecycleConfigVersion("v1", AppConfigVersionStatusRollback) + if _, _, err := RollbackAppConfigVersion(active, rollback, time.Now()); err == nil || + !strings.Contains(err.Error(), "active") { + t.Fatalf("expected active status requirement, got %v", err) + } + + active = validLifecycleConfigVersion("v2", AppConfigVersionStatusActive) + rollback = validLifecycleConfigVersion("v1", AppConfigVersionStatusReleased) + if _, _, err := RollbackAppConfigVersion(active, rollback, time.Now()); err == nil || + !strings.Contains(err.Error(), "rollback") { + t.Fatalf("expected rollback status requirement, got %v", err) + } + + rollback = validLifecycleConfigVersion("v1", AppConfigVersionStatusRollback) + rollback.AppID = "other-app" + if _, _, err := RollbackAppConfigVersion(active, rollback, time.Now()); err == nil || + !strings.Contains(err.Error(), "app_id") { + t.Fatalf("expected app mismatch error, got %v", err) + } +} + +func validLifecycleConfigVersion(version string, status AppConfigVersionStatus) AppConfigVersion { + configVersion := validAppConfigVersion() + configVersion.Version = version + configVersion.Status = status + return configVersion +} From 1a12b0f9ca81f5f3d233fa1fba86c30098421909 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 15:10:14 +0800 Subject: [PATCH 22/95] platform: add usage summary contracts --- platform/usage_summary.go | 117 ++++++++++++++++++++++++++ platform/usage_summary_test.go | 149 +++++++++++++++++++++++++++++++++ 2 files changed, 266 insertions(+) create mode 100644 platform/usage_summary.go create mode 100644 platform/usage_summary_test.go diff --git a/platform/usage_summary.go b/platform/usage_summary.go new file mode 100644 index 0000000000..be29989458 --- /dev/null +++ b/platform/usage_summary.go @@ -0,0 +1,117 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "fmt" + "strings" +) + +// UsageSummaryFilter scopes usage aggregation to one tenant and optionally one app. +type UsageSummaryFilter struct { + TenantID string + AppID string +} + +// UsageSummary aggregates post-run token and cost records for dashboards and budget checks. +type UsageSummary struct { + TenantID string + AppID string + RecordCount int + PromptTokens int + CompletionTokens int + CachedTokens int + TotalTokens int + ModelCost float64 + ToolCost float64 + TotalCost float64 +} + +// SummarizeUsage aggregates usage records for one tenant and optional app. +func SummarizeUsage(records []UsageRecord, filter UsageSummaryFilter) (UsageSummary, error) { + tenantID := strings.TrimSpace(filter.TenantID) + appID := strings.TrimSpace(filter.AppID) + if tenantID == "" { + return UsageSummary{}, ErrTenantIDRequired + } + summary := UsageSummary{ + TenantID: tenantID, + AppID: appID, + } + for _, record := range records { + if strings.TrimSpace(record.TenantID) != tenantID { + continue + } + if appID != "" && strings.TrimSpace(record.AppID) != appID { + continue + } + if err := record.Validate(); err != nil { + return UsageSummary{}, err + } + if err := summary.add(record); err != nil { + return UsageSummary{}, err + } + } + return summary, nil +} + +// Summary returns an aggregate snapshot for the in-memory sink records. +func (s *InMemoryUsageSink) Summary(filter UsageSummaryFilter) (UsageSummary, error) { + return SummarizeUsage(s.Records(), filter) +} + +func (s *UsageSummary) add(record UsageRecord) error { + var err error + if s.RecordCount, err = addUsageInt("record_count", s.RecordCount, 1); err != nil { + return err + } + if s.PromptTokens, err = addUsageInt("prompt_tokens", s.PromptTokens, record.PromptTokens); err != nil { + return err + } + if s.CompletionTokens, err = addUsageInt("completion_tokens", s.CompletionTokens, record.CompletionTokens); err != nil { + return err + } + if s.CachedTokens, err = addUsageInt("cached_tokens", s.CachedTokens, record.CachedTokens); err != nil { + return err + } + if s.TotalTokens, err = addUsageInt("total_tokens", s.TotalTokens, record.PromptTokens); err != nil { + return err + } + if s.TotalTokens, err = addUsageInt("total_tokens", s.TotalTokens, record.CompletionTokens); err != nil { + return err + } + if s.ModelCost, err = addUsageCost("model_cost", s.ModelCost, record.ModelCost); err != nil { + return err + } + if s.ToolCost, err = addUsageCost("tool_cost", s.ToolCost, record.ToolCost); err != nil { + return err + } + if s.TotalCost, err = addUsageCost("total_cost", s.TotalCost, record.TotalCost); err != nil { + return err + } + return nil +} + +func addUsageInt(field string, current int, next int) (int, error) { + if next < 0 { + return 0, fmt.Errorf("%s must be non-negative", field) + } + if current > maxInt()-next { + return 0, fmt.Errorf("%s overflow", field) + } + return current + next, nil +} + +func addUsageCost(field string, current float64, next float64) (float64, error) { + total := current + next + if !isFiniteNonNegative(total) { + return 0, fmt.Errorf("%s total must be finite and non-negative", field) + } + return total, nil +} diff --git a/platform/usage_summary_test.go b/platform/usage_summary_test.go new file mode 100644 index 0000000000..ae99246eac --- /dev/null +++ b/platform/usage_summary_test.go @@ -0,0 +1,149 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "context" + "errors" + "math" + "strings" + "testing" +) + +func TestSummarizeUsageAggregatesTenantAndApp(t *testing.T) { + records := []UsageRecord{ + usageRecordForSummary("tenant-a", "app-a", 100, 50, 10, 0.15, 0.20, 0.35), + usageRecordForSummary("tenant-a", "app-a", 20, 5, 0, 0.02, 0.01, 0.03), + usageRecordForSummary("tenant-a", "app-b", 1_000, 500, 0, 10, 1, 11), + usageRecordForSummary("tenant-b", "app-a", 2_000, 600, 0, 20, 2, 22), + } + + summary, err := SummarizeUsage(records, UsageSummaryFilter{TenantID: " tenant-a ", AppID: " app-a "}) + if err != nil { + t.Fatalf("summarize usage: %v", err) + } + if summary.TenantID != "tenant-a" || summary.AppID != "app-a" { + t.Fatalf("summary should expose normalized scope, got %+v", summary) + } + if summary.RecordCount != 2 { + t.Fatalf("expected 2 records, got %d", summary.RecordCount) + } + if summary.PromptTokens != 120 || + summary.CompletionTokens != 55 || + summary.CachedTokens != 10 || + summary.TotalTokens != 175 { + t.Fatalf("unexpected token totals: %+v", summary) + } + assertFloat(t, "ModelCost", summary.ModelCost, 0.17) + assertFloat(t, "ToolCost", summary.ToolCost, 0.21) + assertFloat(t, "TotalCost", summary.TotalCost, 0.38) +} + +func TestSummarizeUsageAggregatesTenantAcrossApps(t *testing.T) { + records := []UsageRecord{ + usageRecordForSummary("tenant-a", "app-a", 10, 20, 5, 0.1, 0.2, 0.3), + usageRecordForSummary("tenant-a", "app-b", 30, 40, 0, 0.3, 0.4, 0.7), + usageRecordForSummary("tenant-b", "app-a", 100, 100, 0, 1, 1, 2), + } + + summary, err := SummarizeUsage(records, UsageSummaryFilter{TenantID: "tenant-a"}) + if err != nil { + t.Fatalf("summarize usage: %v", err) + } + if summary.RecordCount != 2 || summary.AppID != "" { + t.Fatalf("expected tenant-wide summary, got %+v", summary) + } + if summary.TotalTokens != 100 { + t.Fatalf("expected tenant token total 100, got %d", summary.TotalTokens) + } + assertFloat(t, "TotalCost", summary.TotalCost, 1.0) +} + +func TestUsageSinkSummaryUsesSnapshot(t *testing.T) { + sink := NewInMemoryUsageSink() + if err := sink.WriteUsage(context.Background(), usageRecordForSummary("tenant", "app", 10, 5, 1, 0.1, 0.2, 0.3)); err != nil { + t.Fatalf("write usage: %v", err) + } + if err := sink.WriteUsage(context.Background(), usageRecordForSummary("tenant", "app", 3, 2, 0, 0.01, 0.02, 0.03)); err != nil { + t.Fatalf("write usage: %v", err) + } + + summary, err := sink.Summary(UsageSummaryFilter{TenantID: "tenant", AppID: "app"}) + if err != nil { + t.Fatalf("sink summary: %v", err) + } + if summary.RecordCount != 2 || summary.TotalTokens != 20 { + t.Fatalf("unexpected sink summary: %+v", summary) + } + assertFloat(t, "TotalCost", summary.TotalCost, 0.33) +} + +func TestSummarizeUsageRequiresTenant(t *testing.T) { + _, err := SummarizeUsage(nil, UsageSummaryFilter{TenantID: " "}) + if !errors.Is(err, ErrTenantIDRequired) { + t.Fatalf("expected tenant requirement, got %v", err) + } +} + +func TestSummarizeUsageRejectsInvalidMatchingRecord(t *testing.T) { + record := usageRecordForSummary("tenant", "app", 10, 5, 0, 0.1, 0.2, 0.3) + record.TotalCost = math.Inf(1) + + _, err := SummarizeUsage([]UsageRecord{record}, UsageSummaryFilter{TenantID: "tenant"}) + if err == nil || !strings.Contains(err.Error(), "total_cost") { + t.Fatalf("expected invalid matching record error, got %v", err) + } +} + +func TestSummarizeUsageIgnoresInvalidNonMatchingRecord(t *testing.T) { + record := usageRecordForSummary("tenant-b", "app", 10, 5, 0, 0.1, 0.2, 0.3) + record.TotalCost = math.Inf(1) + + summary, err := SummarizeUsage([]UsageRecord{record}, UsageSummaryFilter{TenantID: "tenant-a"}) + if err != nil { + t.Fatalf("non-matching record should not be validated, got %v", err) + } + if summary.RecordCount != 0 { + t.Fatalf("expected empty summary, got %+v", summary) + } +} + +func TestSummarizeUsageRejectsTokenOverflow(t *testing.T) { + records := []UsageRecord{ + usageRecordForSummary("tenant", "app", maxInt(), 0, 0, 0, 0, 0), + usageRecordForSummary("tenant", "app", 1, 0, 0, 0, 0, 0), + } + + _, err := SummarizeUsage(records, UsageSummaryFilter{TenantID: "tenant", AppID: "app"}) + if err == nil || !strings.Contains(err.Error(), "overflow") { + t.Fatalf("expected token overflow error, got %v", err) + } +} + +func usageRecordForSummary( + tenantID string, + appID string, + promptTokens int, + completionTokens int, + cachedTokens int, + modelCost float64, + toolCost float64, + totalCost float64, +) UsageRecord { + record := validUsageRecord() + record.TenantID = tenantID + record.AppID = appID + record.PromptTokens = promptTokens + record.CompletionTokens = completionTokens + record.CachedTokens = cachedTokens + record.ModelCost = modelCost + record.ToolCost = toolCost + record.TotalCost = totalCost + return record +} From c44f477c75b3adfd0eebc03218f7e04c534352f6 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 15:16:44 +0800 Subject: [PATCH 23/95] platform: add audit query contracts --- platform/audit_query.go | 136 ++++++++++++++++++++++++++ platform/audit_query_test.go | 184 +++++++++++++++++++++++++++++++++++ 2 files changed, 320 insertions(+) create mode 100644 platform/audit_query.go create mode 100644 platform/audit_query_test.go diff --git a/platform/audit_query.go b/platform/audit_query.go new file mode 100644 index 0000000000..47e1f8dfb7 --- /dev/null +++ b/platform/audit_query.go @@ -0,0 +1,136 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "fmt" + "strings" + "time" +) + +// AuditQueryFilter scopes audit retrieval to one tenant and optional safe dimensions. +type AuditQueryFilter struct { + TenantID string + AppID string + AuditID string + Channel string + BindingID string + UserIDHash string + SessionID string + RequestID string + MessageID string + ToolName string + Decision string + TraceID string + CreatedFrom time.Time + CreatedTo time.Time + Limit int +} + +// QueryAudit returns audit records matching one tenant-scoped filter. +func QueryAudit(records []AuditRecord, filter AuditQueryFilter) ([]AuditRecord, error) { + normalized, err := filter.normalize() + if err != nil { + return nil, err + } + matches := make([]AuditRecord, 0) + for _, record := range records { + if !normalized.matchesScope(record) { + continue + } + if normalized.matches(record) { + if err := record.Validate(); err != nil { + return nil, err + } + matches = append(matches, record) + if normalized.Limit > 0 && len(matches) >= normalized.Limit { + break + } + } + } + return matches, nil +} + +// Query returns audit records matching one tenant-scoped filter. +func (s *InMemoryAuditSink) Query(filter AuditQueryFilter) ([]AuditRecord, error) { + return QueryAudit(s.Records(), filter) +} + +func (f AuditQueryFilter) normalize() (AuditQueryFilter, error) { + f.TenantID = strings.TrimSpace(f.TenantID) + if f.TenantID == "" { + return AuditQueryFilter{}, ErrTenantIDRequired + } + for field, value := range map[string]string{ + "app_id": f.AppID, + "audit_id": f.AuditID, + "channel": f.Channel, + "binding_id": f.BindingID, + "user_id_hash": f.UserIDHash, + "session_id": f.SessionID, + "request_id": f.RequestID, + "message_id": f.MessageID, + "tool_name": f.ToolName, + "decision": f.Decision, + "trace_id": f.TraceID, + } { + if err := validateAuditRedactedText(field, value); err != nil { + return AuditQueryFilter{}, err + } + } + if f.Limit < 0 { + return AuditQueryFilter{}, fmt.Errorf("limit must be non-negative") + } + if !f.CreatedFrom.IsZero() && !f.CreatedTo.IsZero() && f.CreatedFrom.After(f.CreatedTo) { + return AuditQueryFilter{}, fmt.Errorf("created_from must be before or equal to created_to") + } + f.AppID = strings.TrimSpace(f.AppID) + f.AuditID = strings.TrimSpace(f.AuditID) + f.Channel = strings.TrimSpace(f.Channel) + f.BindingID = strings.TrimSpace(f.BindingID) + f.UserIDHash = strings.TrimSpace(f.UserIDHash) + f.SessionID = strings.TrimSpace(f.SessionID) + f.RequestID = strings.TrimSpace(f.RequestID) + f.MessageID = strings.TrimSpace(f.MessageID) + f.ToolName = strings.TrimSpace(f.ToolName) + f.Decision = strings.TrimSpace(f.Decision) + f.TraceID = strings.TrimSpace(f.TraceID) + return f, nil +} + +func (f AuditQueryFilter) matchesScope(record AuditRecord) bool { + return strings.TrimSpace(record.TenantID) == f.TenantID +} + +func (f AuditQueryFilter) matches(record AuditRecord) bool { + if !matchOptional(f.AppID, record.AppID) || + !matchOptional(f.AuditID, record.AuditID) || + !matchOptional(f.Channel, record.Channel) || + !matchOptional(f.BindingID, record.BindingID) || + !matchOptional(f.UserIDHash, record.UserIDHash) || + !matchOptional(f.SessionID, record.SessionID) || + !matchOptional(f.RequestID, record.RequestID) || + !matchOptional(f.MessageID, record.MessageID) || + !matchOptional(f.ToolName, record.ToolName) || + !matchOptional(f.Decision, record.Decision) || + !matchOptional(f.TraceID, record.TraceID) { + return false + } + if !f.CreatedFrom.IsZero() && record.CreatedAt.Before(f.CreatedFrom) { + return false + } + if !f.CreatedTo.IsZero() && record.CreatedAt.After(f.CreatedTo) { + return false + } + return true +} + +func matchOptional(want, got string) bool { + return want == "" || strings.TrimSpace(got) == want +} diff --git a/platform/audit_query_test.go b/platform/audit_query_test.go new file mode 100644 index 0000000000..d0e85b5bd3 --- /dev/null +++ b/platform/audit_query_test.go @@ -0,0 +1,184 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "context" + "errors" + "strings" + "testing" + "time" +) + +func TestQueryAuditFiltersTenantAndSafeDimensions(t *testing.T) { + baseTime := time.Date(2026, 7, 8, 10, 0, 0, 0, time.UTC) + userHash := UserIDHash("tenant-a", "telegram", "external-1") + records := []AuditRecord{ + auditRecordForQuery("tenant-a", "audit-1", "app-a", "telegram", "binding-a", "session-1", "request-1", "message-1", "file_write", "deny", "trace-1", baseTime), + auditRecordForQuery("tenant-a", "audit-2", "app-a", "telegram", "binding-a", "session-2", "request-2", "message-2", "file_write", "allow", "trace-2", baseTime.Add(time.Hour)), + auditRecordForQuery("tenant-a", "audit-3", "app-b", "telegram", "binding-a", "session-1", "request-1", "message-1", "file_write", "deny", "trace-1", baseTime), + auditRecordForQuery("tenant-b", "audit-4", "app-a", "telegram", "binding-a", "session-1", "request-1", "message-1", "file_write", "deny", "trace-1", baseTime), + } + records[0].UserIDHash = userHash + records[1].UserIDHash = UserIDHash("tenant-a", "telegram", "external-2") + + matches, err := QueryAudit(records, AuditQueryFilter{ + TenantID: " tenant-a ", + AppID: " app-a ", + Channel: "telegram", + BindingID: "binding-a", + UserIDHash: userHash, + SessionID: "session-1", + RequestID: "request-1", + MessageID: "message-1", + ToolName: "file_write", + Decision: "deny", + TraceID: "trace-1", + CreatedFrom: baseTime.Add(-time.Minute), + CreatedTo: baseTime.Add(time.Minute), + }) + if err != nil { + t.Fatalf("query audit: %v", err) + } + if len(matches) != 1 || matches[0].AuditID != "audit-1" { + t.Fatalf("expected only audit-1, got %+v", matches) + } +} + +func TestQueryAuditSupportsAuditIDAndLimit(t *testing.T) { + baseTime := time.Date(2026, 7, 8, 10, 0, 0, 0, time.UTC) + records := []AuditRecord{ + auditRecordForQuery("tenant", "audit-1", "app", "telegram", "binding", "session", "request", "message", "tool", "allow", "trace", baseTime), + auditRecordForQuery("tenant", "audit-2", "app", "telegram", "binding", "session", "request", "message", "tool", "allow", "trace", baseTime), + } + + matches, err := QueryAudit(records, AuditQueryFilter{TenantID: "tenant", AuditID: "audit-2"}) + if err != nil { + t.Fatalf("query by audit id: %v", err) + } + if len(matches) != 1 || matches[0].AuditID != "audit-2" { + t.Fatalf("expected audit-2, got %+v", matches) + } + + matches, err = QueryAudit(records, AuditQueryFilter{TenantID: "tenant", Decision: "allow", Limit: 1}) + if err != nil { + t.Fatalf("query with limit: %v", err) + } + if len(matches) != 1 || matches[0].AuditID != "audit-1" { + t.Fatalf("expected first limited result, got %+v", matches) + } +} + +func TestAuditSinkQueryUsesSnapshotAndReturnsCopies(t *testing.T) { + sink := NewInMemoryAuditSink() + record := auditRecordForQuery("tenant", "audit-1", "app", "telegram", "binding", "session", "request", "message", "tool", "allow", "trace", time.Now()) + if err := sink.WriteAudit(context.Background(), record); err != nil { + t.Fatalf("write audit: %v", err) + } + + matches, err := sink.Query(AuditQueryFilter{TenantID: "tenant", AppID: "app"}) + if err != nil { + t.Fatalf("sink query: %v", err) + } + if len(matches) != 1 { + t.Fatalf("expected one match, got %d", len(matches)) + } + matches[0].TenantID = "changed" + again, err := sink.Query(AuditQueryFilter{TenantID: "tenant", AppID: "app"}) + if err != nil { + t.Fatalf("sink query again: %v", err) + } + if again[0].TenantID != "tenant" { + t.Fatalf("query should return defensive copies, got %+v", again[0]) + } +} + +func TestQueryAuditRequiresTenant(t *testing.T) { + _, err := QueryAudit(nil, AuditQueryFilter{TenantID: " "}) + if !errors.Is(err, ErrTenantIDRequired) { + t.Fatalf("expected tenant requirement, got %v", err) + } +} + +func TestQueryAuditRejectsUnsafeFilterValues(t *testing.T) { + _, err := QueryAudit(nil, AuditQueryFilter{ + TenantID: "tenant", + ToolName: "workspace_exec Authorization: Bearer raw-token", + }) + if err == nil || !strings.Contains(err.Error(), "tool_name") { + t.Fatalf("expected unsafe tool filter error, got %v", err) + } +} + +func TestQueryAuditRejectsInvalidLimitAndTimeRange(t *testing.T) { + _, err := QueryAudit(nil, AuditQueryFilter{TenantID: "tenant", Limit: -1}) + if err == nil || !strings.Contains(err.Error(), "limit") { + t.Fatalf("expected limit validation, got %v", err) + } + + from := time.Date(2026, 7, 8, 11, 0, 0, 0, time.UTC) + to := from.Add(-time.Hour) + _, err = QueryAudit(nil, AuditQueryFilter{TenantID: "tenant", CreatedFrom: from, CreatedTo: to}) + if err == nil || !strings.Contains(err.Error(), "created_from") { + t.Fatalf("expected time range validation, got %v", err) + } +} + +func TestQueryAuditRejectsInvalidMatchingRecordOnly(t *testing.T) { + matching := auditRecordForQuery("tenant-a", "audit-1", "app", "telegram", "binding", "session", "request", "message", "tool", "allow", "trace", time.Now()) + matching.LatencyMS = -1 + nonMatching := auditRecordForQuery("tenant-a", "audit-2", "other-app", "telegram", "binding", "session", "request", "message", "tool", "allow", "trace", time.Now()) + nonMatching.LatencyMS = -1 + otherTenant := auditRecordForQuery("tenant-b", "audit-2", "app", "telegram", "binding", "session", "request", "message", "tool", "allow", "trace", time.Now()) + otherTenant.LatencyMS = -1 + + _, err := QueryAudit([]AuditRecord{otherTenant}, AuditQueryFilter{TenantID: "tenant-a"}) + if err != nil { + t.Fatalf("non-matching invalid record should not be validated, got %v", err) + } + _, err = QueryAudit([]AuditRecord{nonMatching}, AuditQueryFilter{TenantID: "tenant-a", AppID: "app"}) + if err != nil { + t.Fatalf("same-tenant non-matching invalid record should not be validated, got %v", err) + } + + _, err = QueryAudit([]AuditRecord{matching}, AuditQueryFilter{TenantID: "tenant-a"}) + if err == nil || !strings.Contains(err.Error(), "latency_ms") { + t.Fatalf("expected invalid matching record error, got %v", err) + } +} + +func auditRecordForQuery( + tenantID string, + auditID string, + appID string, + channel string, + bindingID string, + sessionID string, + requestID string, + messageID string, + toolName string, + decision string, + traceID string, + createdAt time.Time, +) AuditRecord { + record := validAuditRecord() + record.TenantID = tenantID + record.AuditID = auditID + record.AppID = appID + record.Channel = channel + record.BindingID = bindingID + record.SessionID = sessionID + record.RequestID = requestID + record.MessageID = messageID + record.ToolName = toolName + record.Decision = decision + record.TraceID = traceID + record.CreatedAt = createdAt + return record +} From 8eefe55ab6bfc739f29d14c2be1af531cea4e7fa Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 15:33:56 +0800 Subject: [PATCH 24/95] platform: add config version diff contracts --- platform/config_diff.go | 202 +++++++++++++++++++++++++++++++++++ platform/config_diff_test.go | 159 +++++++++++++++++++++++++++ 2 files changed, 361 insertions(+) create mode 100644 platform/config_diff.go create mode 100644 platform/config_diff_test.go diff --git a/platform/config_diff.go b/platform/config_diff.go new file mode 100644 index 0000000000..2e8d7268d0 --- /dev/null +++ b/platform/config_diff.go @@ -0,0 +1,202 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "bytes" + "encoding/json" + "fmt" + "reflect" + "sort" + "strings" +) + +// AppConfigVersionDiffKind describes how one config version field changed. +type AppConfigVersionDiffKind string + +const ( + // AppConfigVersionDiffAdded means the field exists only in the target version. + AppConfigVersionDiffAdded AppConfigVersionDiffKind = "added" + // AppConfigVersionDiffRemoved means the field exists only in the source version. + AppConfigVersionDiffRemoved AppConfigVersionDiffKind = "removed" + // AppConfigVersionDiffChanged means the field exists in both versions with different values. + AppConfigVersionDiffChanged AppConfigVersionDiffKind = "changed" +) + +// AppConfigVersionDiffChange is one safe, displayable config version difference. +type AppConfigVersionDiffChange struct { + // Path is a metadata field name or a JSON Pointer-style config bundle path. + Path string + Kind AppConfigVersionDiffKind + Before string + After string +} + +// AppConfigVersionDiff summarizes differences between two versions owned by the same tenant app. +type AppConfigVersionDiff struct { + TenantID string + AppID string + FromVersion string + ToVersion string + Changes []AppConfigVersionDiffChange +} + +type missingConfigValue struct{} + +// DiffAppConfigVersions compares safe metadata and config bundle values for two app config versions. +func DiffAppConfigVersions(from, to AppConfigVersion) (AppConfigVersionDiff, error) { + if err := from.Validate(); err != nil { + return AppConfigVersionDiff{}, fmt.Errorf("from config version: %w", err) + } + if err := to.Validate(); err != nil { + return AppConfigVersionDiff{}, fmt.Errorf("to config version: %w", err) + } + if err := requireSameConfigOwner(from, to); err != nil { + return AppConfigVersionDiff{}, err + } + diff := AppConfigVersionDiff{ + TenantID: from.TenantID, + AppID: from.AppID, + FromVersion: from.Version, + ToVersion: to.Version, + } + addConfigScalarChange(&diff.Changes, "version", from.Version, to.Version) + addConfigScalarChange(&diff.Changes, "checksum", from.Checksum, to.Checksum) + addConfigScalarChange(&diff.Changes, "status", string(from.Status), string(to.Status)) + addConfigScalarChange(&diff.Changes, "gray_percent", from.GrayPercent, to.GrayPercent) + + fromBundle, err := decodeConfigBundle(from.ConfigBundleJSON) + if err != nil { + return AppConfigVersionDiff{}, fmt.Errorf("from config bundle: %w", err) + } + toBundle, err := decodeConfigBundle(to.ConfigBundleJSON) + if err != nil { + return AppConfigVersionDiff{}, fmt.Errorf("to config bundle: %w", err) + } + diffConfigBundleValue("/config_bundle_json", fromBundle, toBundle, &diff.Changes) + return diff, nil +} + +func decodeConfigBundle(bundle string) (any, error) { + decoder := json.NewDecoder(bytes.NewBufferString(bundle)) + decoder.UseNumber() + var value any + if err := decoder.Decode(&value); err != nil { + return nil, err + } + return value, nil +} + +func addConfigScalarChange(changes *[]AppConfigVersionDiffChange, path string, before, after any) { + if reflect.DeepEqual(before, after) { + return + } + *changes = append(*changes, AppConfigVersionDiffChange{ + Path: path, + Kind: AppConfigVersionDiffChanged, + Before: fmt.Sprint(before), + After: fmt.Sprint(after), + }) +} + +func diffConfigBundleValue(path string, before, after any, changes *[]AppConfigVersionDiffChange) { + if _, ok := before.(missingConfigValue); ok { + *changes = append(*changes, AppConfigVersionDiffChange{ + Path: path, + Kind: AppConfigVersionDiffAdded, + After: formatConfigBundleValue(after), + }) + return + } + if _, ok := after.(missingConfigValue); ok { + *changes = append(*changes, AppConfigVersionDiffChange{ + Path: path, + Kind: AppConfigVersionDiffRemoved, + Before: formatConfigBundleValue(before), + }) + return + } + + beforeMap, beforeIsMap := before.(map[string]any) + afterMap, afterIsMap := after.(map[string]any) + if beforeIsMap && afterIsMap { + for _, key := range sortedConfigKeys(beforeMap, afterMap) { + childBefore, ok := beforeMap[key] + if !ok { + childBefore = missingConfigValue{} + } + childAfter, ok := afterMap[key] + if !ok { + childAfter = missingConfigValue{} + } + diffConfigBundleValue(path+"/"+escapeConfigPathSegment(key), childBefore, childAfter, changes) + } + return + } + + beforeItems, beforeIsArray := before.([]any) + afterItems, afterIsArray := after.([]any) + if beforeIsArray && afterIsArray { + maxLen := len(beforeItems) + if len(afterItems) > maxLen { + maxLen = len(afterItems) + } + for i := 0; i < maxLen; i++ { + childBefore := any(missingConfigValue{}) + if i < len(beforeItems) { + childBefore = beforeItems[i] + } + childAfter := any(missingConfigValue{}) + if i < len(afterItems) { + childAfter = afterItems[i] + } + diffConfigBundleValue(fmt.Sprintf("%s/%d", path, i), childBefore, childAfter, changes) + } + return + } + + if reflect.DeepEqual(before, after) { + return + } + *changes = append(*changes, AppConfigVersionDiffChange{ + Path: path, + Kind: AppConfigVersionDiffChanged, + Before: formatConfigBundleValue(before), + After: formatConfigBundleValue(after), + }) +} + +func escapeConfigPathSegment(segment string) string { + segment = strings.ReplaceAll(segment, "~", "~0") + return strings.ReplaceAll(segment, "/", "~1") +} + +func sortedConfigKeys(left, right map[string]any) []string { + seen := make(map[string]struct{}, len(left)+len(right)) + for key := range left { + seen[key] = struct{}{} + } + for key := range right { + seen[key] = struct{}{} + } + keys := make([]string, 0, len(seen)) + for key := range seen { + keys = append(keys, key) + } + sort.Strings(keys) + return keys +} + +func formatConfigBundleValue(value any) string { + encoded, err := json.Marshal(value) + if err != nil { + return fmt.Sprint(value) + } + return string(encoded) +} diff --git a/platform/config_diff_test.go b/platform/config_diff_test.go new file mode 100644 index 0000000000..15fbc7c8df --- /dev/null +++ b/platform/config_diff_test.go @@ -0,0 +1,159 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "strings" + "testing" +) + +func TestDiffAppConfigVersionsReportsMetadataAndBundleChanges(t *testing.T) { + from := validAppConfigVersion() + from.Version = "v1" + from.Checksum = "sha256:1111" + from.Status = AppConfigVersionStatusActive + from.GrayPercent = 0 + from.ConfigBundleJSON = `{ + "model_profile_id":"model-a", + "tool_policy_id":"tools-a", + "limits":{"max_tokens":1000,"temperature":0.2}, + "tools":["search","ticket"], + "old_field":"removed" + }` + to := validAppConfigVersion() + to.Version = "v2" + to.Checksum = "sha256:2222" + to.Status = AppConfigVersionStatusReleased + to.GrayPercent = 25 + to.ConfigBundleJSON = `{ + "model_profile_id":"model-b", + "tool_policy_id":"tools-a", + "limits":{"max_tokens":2000,"temperature":0.2}, + "tools":["search","crm"], + "new_field":"added" + }` + + diff, err := DiffAppConfigVersions(from, to) + if err != nil { + t.Fatalf("diff config versions: %v", err) + } + if diff.TenantID != "tenant" || diff.AppID != "app" || diff.FromVersion != "v1" || diff.ToVersion != "v2" { + t.Fatalf("unexpected diff identity: %+v", diff) + } + assertConfigDiffChange(t, diff, "version", AppConfigVersionDiffChanged, "v1", "v2") + assertConfigDiffChange(t, diff, "checksum", AppConfigVersionDiffChanged, "sha256:1111", "sha256:2222") + assertConfigDiffChange(t, diff, "status", AppConfigVersionDiffChanged, "active", "released") + assertConfigDiffChange(t, diff, "gray_percent", AppConfigVersionDiffChanged, "0", "25") + assertConfigDiffChange(t, diff, "/config_bundle_json/model_profile_id", AppConfigVersionDiffChanged, `"model-a"`, `"model-b"`) + assertConfigDiffChange(t, diff, "/config_bundle_json/limits/max_tokens", AppConfigVersionDiffChanged, "1000", "2000") + assertConfigDiffChange(t, diff, "/config_bundle_json/tools/1", AppConfigVersionDiffChanged, `"ticket"`, `"crm"`) + assertConfigDiffChange(t, diff, "/config_bundle_json/old_field", AppConfigVersionDiffRemoved, `"removed"`, "") + assertConfigDiffChange(t, diff, "/config_bundle_json/new_field", AppConfigVersionDiffAdded, "", `"added"`) +} + +func TestDiffAppConfigVersionsReturnsNoChangesForEquivalentBundles(t *testing.T) { + from := validAppConfigVersion() + from.ConfigBundleJSON = `{"tool_policy_id":"tools","model_profile_id":"model"}` + to := from + to.ConfigBundleJSON = `{ + "model_profile_id":"model", + "tool_policy_id":"tools" + }` + + diff, err := DiffAppConfigVersions(from, to) + if err != nil { + t.Fatalf("diff equivalent config versions: %v", err) + } + if len(diff.Changes) != 0 { + t.Fatalf("expected no changes, got %+v", diff.Changes) + } +} + +func TestDiffAppConfigVersionsRejectsInvalidInputs(t *testing.T) { + from := validAppConfigVersion() + to := validAppConfigVersion() + to.TenantID = "other-tenant" + if _, err := DiffAppConfigVersions(from, to); err == nil || !strings.Contains(err.Error(), "tenant_id") { + t.Fatalf("expected tenant mismatch error, got %v", err) + } + + to = validAppConfigVersion() + to.AppID = "other-app" + if _, err := DiffAppConfigVersions(from, to); err == nil || !strings.Contains(err.Error(), "app_id") { + t.Fatalf("expected app mismatch error, got %v", err) + } + + to = validAppConfigVersion() + to.ConfigBundleJSON = `{"api_key":"sk-1234567890abcdef"}` + if _, err := DiffAppConfigVersions(from, to); err == nil || !strings.Contains(err.Error(), "to config version") { + t.Fatalf("expected invalid target validation error, got %v", err) + } +} + +func TestDiffAppConfigVersionsReportsArrayAddRemove(t *testing.T) { + from := validAppConfigVersion() + from.ConfigBundleJSON = `{"tools":["search","ticket"]}` + to := validAppConfigVersion() + to.Version = "v2" + to.Checksum = "sha256:2222" + to.ConfigBundleJSON = `{"tools":["search","ticket","crm"]}` + + diff, err := DiffAppConfigVersions(from, to) + if err != nil { + t.Fatalf("diff array growth: %v", err) + } + assertConfigDiffChange(t, diff, "/config_bundle_json/tools/2", AppConfigVersionDiffAdded, "", `"crm"`) + + removed, err := DiffAppConfigVersions(to, from) + if err != nil { + t.Fatalf("diff array shrink: %v", err) + } + assertConfigDiffChange(t, removed, "/config_bundle_json/tools/2", AppConfigVersionDiffRemoved, `"crm"`, "") +} + +func TestDiffAppConfigVersionsEscapesAmbiguousObjectKeys(t *testing.T) { + from := validAppConfigVersion() + from.ConfigBundleJSON = `{"a.b":1,"a":{"b":1},"slash/key":"old","tilde~key":"old","tools[1]":"old","tools":["search","ticket"]}` + to := validAppConfigVersion() + to.Version = "v2" + to.Checksum = "sha256:2222" + to.ConfigBundleJSON = `{"a.b":2,"a":{"b":3},"slash/key":"new","tilde~key":"new","tools[1]":"new","tools":["search","crm"]}` + + diff, err := DiffAppConfigVersions(from, to) + if err != nil { + t.Fatalf("diff escaped keys: %v", err) + } + assertConfigDiffChange(t, diff, "/config_bundle_json/a.b", AppConfigVersionDiffChanged, "1", "2") + assertConfigDiffChange(t, diff, "/config_bundle_json/a/b", AppConfigVersionDiffChanged, "1", "3") + assertConfigDiffChange(t, diff, "/config_bundle_json/slash~1key", AppConfigVersionDiffChanged, `"old"`, `"new"`) + assertConfigDiffChange(t, diff, "/config_bundle_json/tilde~0key", AppConfigVersionDiffChanged, `"old"`, `"new"`) + assertConfigDiffChange(t, diff, "/config_bundle_json/tools[1]", AppConfigVersionDiffChanged, `"old"`, `"new"`) + assertConfigDiffChange(t, diff, "/config_bundle_json/tools/1", AppConfigVersionDiffChanged, `"ticket"`, `"crm"`) +} + +func assertConfigDiffChange( + t *testing.T, + diff AppConfigVersionDiff, + path string, + kind AppConfigVersionDiffKind, + before string, + after string, +) { + t.Helper() + for _, change := range diff.Changes { + if change.Path != path { + continue + } + if change.Kind != kind || change.Before != before || change.After != after { + t.Fatalf("unexpected change for %s: %+v", path, change) + } + return + } + t.Fatalf("missing change %s in %+v", path, diff.Changes) +} From 6b8b00c2d7d236977fb7d6a622722ec804418010 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 15:39:50 +0800 Subject: [PATCH 25/95] platform: add config gray status summary --- platform/gray_status.go | 104 ++++++++++++++++++++++++++ platform/gray_status_test.go | 141 +++++++++++++++++++++++++++++++++++ 2 files changed, 245 insertions(+) create mode 100644 platform/gray_status.go create mode 100644 platform/gray_status_test.go diff --git a/platform/gray_status.go b/platform/gray_status.go new file mode 100644 index 0000000000..3604fce40c --- /dev/null +++ b/platform/gray_status.go @@ -0,0 +1,104 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import "fmt" + +// ConfigGrayStatusSummary is an operations-facing view of config gray rollout state. +type ConfigGrayStatusSummary struct { + TenantID string + AppID string + ActiveVersion string + ActiveChecksum string + // ActiveTrafficPercent is the configured routing share, not observed live traffic. + ActiveTrafficPercent int + HasCandidate bool + CandidateVersion string + CandidateChecksum string + CandidateGrayPercent int + // CandidateTrafficPercent is the configured routing share, not observed live traffic. + CandidateTrafficPercent int + HasRollback bool + RollbackVersion string + RollbackChecksum string +} + +// SummarizeAppConfigGrayStatus builds a safe gray rollout summary from known config versions. +func SummarizeAppConfigGrayStatus(versions []AppConfigVersion) (ConfigGrayStatusSummary, error) { + if len(versions) == 0 { + return ConfigGrayStatusSummary{}, fmt.Errorf("config versions are required") + } + + var owner AppConfigVersion + var ownerSet bool + var active AppConfigVersion + var hasActive bool + var candidate AppConfigVersion + var hasCandidate bool + var rollback AppConfigVersion + var hasRollback bool + + for _, version := range versions { + if err := version.Validate(); err != nil { + return ConfigGrayStatusSummary{}, fmt.Errorf("config version %q: %w", version.Version, err) + } + if !ownerSet { + owner = version + ownerSet = true + } else if err := requireSameConfigOwner(owner, version); err != nil { + return ConfigGrayStatusSummary{}, err + } + + switch version.Status { + case AppConfigVersionStatusActive: + if hasActive { + return ConfigGrayStatusSummary{}, fmt.Errorf("multiple active config versions") + } + active = version + hasActive = true + case AppConfigVersionStatusReleased: + if hasCandidate { + return ConfigGrayStatusSummary{}, fmt.Errorf("multiple released config versions") + } + candidate = version + hasCandidate = true + case AppConfigVersionStatusRollback: + if hasRollback { + return ConfigGrayStatusSummary{}, fmt.Errorf("multiple rollback config versions") + } + rollback = version + hasRollback = true + } + } + if !hasActive { + return ConfigGrayStatusSummary{}, fmt.Errorf("active config version is required") + } + + summary := ConfigGrayStatusSummary{ + TenantID: active.TenantID, + AppID: active.AppID, + ActiveVersion: active.Version, + ActiveChecksum: active.Checksum, + ActiveTrafficPercent: 100, + } + if hasCandidate { + summary.HasCandidate = true + summary.CandidateVersion = candidate.Version + summary.CandidateChecksum = candidate.Checksum + summary.CandidateGrayPercent = candidate.GrayPercent + summary.CandidateTrafficPercent = candidate.GrayPercent + summary.ActiveTrafficPercent = 100 - candidate.GrayPercent + } + if hasRollback { + summary.HasRollback = true + summary.RollbackVersion = rollback.Version + summary.RollbackChecksum = rollback.Checksum + } + return summary, nil +} diff --git a/platform/gray_status_test.go b/platform/gray_status_test.go new file mode 100644 index 0000000000..473ddf46ee --- /dev/null +++ b/platform/gray_status_test.go @@ -0,0 +1,141 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "strings" + "testing" +) + +func TestSummarizeAppConfigGrayStatusReportsCandidateAndRollback(t *testing.T) { + active := validGrayActiveConfigVersion() + active.Checksum = "sha256:active" + candidate := validGrayCandidateConfigVersion() + candidate.GrayPercent = 25 + candidate.Checksum = "sha256:candidate" + rollback := validGrayRollbackConfigVersion() + rollback.Checksum = "sha256:rollback" + draft := validAppConfigVersion() + draft.Version = "draft" + draft.Status = AppConfigVersionStatusDraft + + summary, err := SummarizeAppConfigGrayStatus([]AppConfigVersion{draft, rollback, candidate, active}) + if err != nil { + t.Fatalf("summarize gray status: %v", err) + } + if summary.TenantID != active.TenantID || summary.AppID != active.AppID { + t.Fatalf("unexpected owner: %+v", summary) + } + if summary.ActiveVersion != "v1" || summary.ActiveChecksum != "sha256:active" { + t.Fatalf("unexpected active version summary: %+v", summary) + } + if summary.ActiveTrafficPercent != 75 { + t.Fatalf("expected active traffic 75, got %d", summary.ActiveTrafficPercent) + } + if !summary.HasCandidate || summary.CandidateVersion != "v2" || summary.CandidateChecksum != "sha256:candidate" { + t.Fatalf("unexpected candidate summary: %+v", summary) + } + if summary.CandidateGrayPercent != 25 || summary.CandidateTrafficPercent != 25 { + t.Fatalf("unexpected candidate traffic summary: %+v", summary) + } + if !summary.HasRollback || summary.RollbackVersion != "v0" || summary.RollbackChecksum != "sha256:rollback" { + t.Fatalf("unexpected rollback summary: %+v", summary) + } +} + +func TestSummarizeAppConfigGrayStatusHandlesActiveOnly(t *testing.T) { + active := validGrayActiveConfigVersion() + + summary, err := SummarizeAppConfigGrayStatus([]AppConfigVersion{active}) + if err != nil { + t.Fatalf("summarize active-only gray status: %v", err) + } + if summary.ActiveTrafficPercent != 100 { + t.Fatalf("expected all traffic on active, got %+v", summary) + } + if summary.HasCandidate || summary.CandidateTrafficPercent != 0 || summary.CandidateVersion != "" { + t.Fatalf("expected no candidate fields, got %+v", summary) + } + if summary.HasRollback || summary.RollbackVersion != "" { + t.Fatalf("expected no rollback fields, got %+v", summary) + } +} + +func TestSummarizeAppConfigGrayStatusRejectsInvalidInputs(t *testing.T) { + if _, err := SummarizeAppConfigGrayStatus(nil); err == nil || + !strings.Contains(err.Error(), "config versions are required") { + t.Fatalf("expected empty input error, got %v", err) + } + + invalid := validGrayActiveConfigVersion() + invalid.ConfigBundleJSON = "{" + if _, err := SummarizeAppConfigGrayStatus([]AppConfigVersion{invalid}); err == nil || + !strings.Contains(err.Error(), "config version") { + t.Fatalf("expected invalid version error, got %v", err) + } + + candidate := validGrayCandidateConfigVersion() + if _, err := SummarizeAppConfigGrayStatus([]AppConfigVersion{candidate}); err == nil || + !strings.Contains(err.Error(), "active config version") { + t.Fatalf("expected missing active error, got %v", err) + } + + active := validGrayActiveConfigVersion() + otherActive := validGrayActiveConfigVersion() + otherActive.Version = "v1b" + otherActive.Checksum = "sha256:active-b" + if _, err := SummarizeAppConfigGrayStatus([]AppConfigVersion{active, otherActive}); err == nil || + !strings.Contains(err.Error(), "multiple active") { + t.Fatalf("expected duplicate active error, got %v", err) + } + + otherCandidate := validGrayCandidateConfigVersion() + otherCandidate.Version = "v3" + otherCandidate.Checksum = "sha256:candidate-3" + if _, err := SummarizeAppConfigGrayStatus([]AppConfigVersion{ + active, + validGrayCandidateConfigVersion(), + otherCandidate, + }); err == nil || !strings.Contains(err.Error(), "multiple released") { + t.Fatalf("expected duplicate candidate error, got %v", err) + } + + mismatched := validGrayCandidateConfigVersion() + mismatched.TenantID = "other-tenant" + if _, err := SummarizeAppConfigGrayStatus([]AppConfigVersion{active, mismatched}); err == nil || + !strings.Contains(err.Error(), "tenant_id") { + t.Fatalf("expected tenant mismatch error, got %v", err) + } + + mismatched = validGrayCandidateConfigVersion() + mismatched.AppID = "other-app" + if _, err := SummarizeAppConfigGrayStatus([]AppConfigVersion{active, mismatched}); err == nil || + !strings.Contains(err.Error(), "app_id") { + t.Fatalf("expected app mismatch error, got %v", err) + } + + otherRollback := validGrayRollbackConfigVersion() + otherRollback.Version = "v-1" + otherRollback.Checksum = "sha256:rollback-2" + if _, err := SummarizeAppConfigGrayStatus([]AppConfigVersion{ + active, + validGrayRollbackConfigVersion(), + otherRollback, + }); err == nil || !strings.Contains(err.Error(), "multiple rollback") { + t.Fatalf("expected duplicate rollback error, got %v", err) + } +} + +func validGrayRollbackConfigVersion() AppConfigVersion { + version := validAppConfigVersion() + version.Version = "v0" + version.Status = AppConfigVersionStatusRollback + version.GrayPercent = 0 + return version +} From c59acc43c22bca5d809d4e5c34e1d2adb373d989 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 15:49:48 +0800 Subject: [PATCH 26/95] platform: add operational action audit contracts --- platform/operational_action_audit.go | 216 ++++++++++++++++++++++ platform/operational_action_audit_test.go | 192 +++++++++++++++++++ 2 files changed, 408 insertions(+) create mode 100644 platform/operational_action_audit.go create mode 100644 platform/operational_action_audit_test.go diff --git a/platform/operational_action_audit.go b/platform/operational_action_audit.go new file mode 100644 index 0000000000..5221504ec1 --- /dev/null +++ b/platform/operational_action_audit.go @@ -0,0 +1,216 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "bytes" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "strings" + "time" +) + +// OperationalAction names a high-risk operations or admin action. +type OperationalAction string + +const ( + // OperationalActionDeleteTenant deletes or disables an entire tenant boundary. + OperationalActionDeleteTenant OperationalAction = "delete_tenant" + // OperationalActionSwitchStorageProfile changes the tenant/app storage route. + OperationalActionSwitchStorageProfile OperationalAction = "switch_storage_profile" + // OperationalActionDisableAudit disables or weakens audit capture. + OperationalActionDisableAudit OperationalAction = "disable_audit" + // OperationalActionExpandToolPermission expands tool access for an app. + OperationalActionExpandToolPermission OperationalAction = "expand_tool_permission" + // OperationalActionExecuteDataMigration runs an operational data migration. + OperationalActionExecuteDataMigration OperationalAction = "execute_data_migration" +) + +// OperationalActionDecision is the recorded outcome of an operations action boundary. +type OperationalActionDecision string + +const ( + // OperationalActionDecisionApprovalRequired records a pending confirmation boundary. + OperationalActionDecisionApprovalRequired OperationalActionDecision = "approval_required" + // OperationalActionDecisionApproved records an approved operation. + OperationalActionDecisionApproved OperationalActionDecision = "approved" + // OperationalActionDecisionRejected records a rejected operation. + OperationalActionDecisionRejected OperationalActionDecision = "rejected" + // OperationalActionDecisionExecuted records a completed operation. + OperationalActionDecisionExecuted OperationalActionDecision = "executed" + // OperationalActionDecisionFailed records a failed operation attempt. + OperationalActionDecisionFailed OperationalActionDecision = "failed" +) + +// OperationalActionAuditInput contains safe dimensions for a high-risk operations audit record. +type OperationalActionAuditInput struct { + TenantID string + AppID string + Action OperationalAction + OperationID string + ResourceType string + ResourceID string + ActorUserID string + ActorInternalUserID string + ApproverUserID string + Decision OperationalActionDecision + DecisionReason string + RequestID string + TraceID string + DetailJSON []byte + CreatedAt time.Time +} + +// NewOperationalActionAuditRecord maps an operations action into a safe audit record. +func NewOperationalActionAuditRecord(input OperationalActionAuditInput) (AuditRecord, error) { + normalized, err := input.normalize() + if err != nil { + return AuditRecord{}, err + } + record := AuditRecord{ + TenantID: normalized.TenantID, + AppID: normalized.AppID, + AuditID: normalized.auditID(), + RequestID: normalized.RequestID, + TraceID: normalized.TraceID, + InternalUserID: normalized.ActorInternalUserID, + UserIDHash: normalized.actorHash(), + ToolName: "ops:" + string(normalized.Action), + Decision: string(normalized.Decision), + DecisionReason: normalized.DecisionReason, + RedactedDetailRef: normalized.detailRef(), + RedactionVersion: "platform-operational-action-v1", + CreatedAt: normalized.CreatedAt, + } + if err := record.Validate(); err != nil { + return AuditRecord{}, err + } + return record, nil +} + +func (i OperationalActionAuditInput) normalize() (OperationalActionAuditInput, error) { + i.TenantID = strings.TrimSpace(i.TenantID) + if i.TenantID == "" { + return OperationalActionAuditInput{}, ErrTenantIDRequired + } + i.AppID = strings.TrimSpace(i.AppID) + i.Action = OperationalAction(strings.TrimSpace(string(i.Action))) + if !i.Action.valid() { + return OperationalActionAuditInput{}, fmt.Errorf("invalid operational action %q", i.Action) + } + i.OperationID = strings.TrimSpace(i.OperationID) + if i.OperationID == "" { + return OperationalActionAuditInput{}, fmt.Errorf("operation_id is required") + } + i.ResourceType = strings.TrimSpace(i.ResourceType) + if i.ResourceType == "" { + return OperationalActionAuditInput{}, fmt.Errorf("resource_type is required") + } + i.ResourceID = strings.TrimSpace(i.ResourceID) + if i.ResourceID == "" { + return OperationalActionAuditInput{}, fmt.Errorf("resource_id is required") + } + i.ActorUserID = strings.TrimSpace(i.ActorUserID) + i.ActorInternalUserID = strings.TrimSpace(i.ActorInternalUserID) + if i.ActorUserID == "" && i.ActorInternalUserID == "" { + return OperationalActionAuditInput{}, fmt.Errorf("actor identity is required") + } + i.ApproverUserID = strings.TrimSpace(i.ApproverUserID) + i.Decision = OperationalActionDecision(strings.TrimSpace(string(i.Decision))) + if !i.Decision.valid() { + return OperationalActionAuditInput{}, fmt.Errorf("invalid operational action decision %q", i.Decision) + } + i.DecisionReason = strings.TrimSpace(i.DecisionReason) + i.RequestID = strings.TrimSpace(i.RequestID) + i.TraceID = strings.TrimSpace(i.TraceID) + i.DetailJSON = bytes.TrimSpace(i.DetailJSON) + if len(i.DetailJSON) > 0 && !json.Valid(i.DetailJSON) { + return OperationalActionAuditInput{}, fmt.Errorf("detail_json must be valid json") + } + for field, value := range map[string]string{ + "app_id": i.AppID, + "action": string(i.Action), + "operation_id": i.OperationID, + "resource_type": i.ResourceType, + "actor_internal_user_id": i.ActorInternalUserID, + "decision": string(i.Decision), + "decision_reason": i.DecisionReason, + "request_id": i.RequestID, + "trace_id": i.TraceID, + } { + if err := validateAuditRedactedText(field, value); err != nil { + return OperationalActionAuditInput{}, err + } + } + return i, nil +} + +func (i OperationalActionAuditInput) auditID() string { + return AuditID( + i.TenantID, + i.AppID, + string(i.Action), + i.OperationID, + i.ResourceType, + i.ResourceID, + string(i.Decision), + ) +} + +func (i OperationalActionAuditInput) actorHash() string { + actor := i.ActorUserID + if actor == "" { + actor = i.ActorInternalUserID + } + return UserIDHash(i.TenantID, "ops", actor) +} + +func (i OperationalActionAuditInput) detailRef() string { + parts := []string{ + "resource_type:" + i.ResourceType, + "resource_hash:" + shortHash(i.TenantID, i.ResourceType, i.ResourceID), + } + if i.ApproverUserID != "" { + parts = append(parts, "approver_hash:"+UserIDHash(i.TenantID, "ops", i.ApproverUserID)) + } + if len(i.DetailJSON) > 0 { + sum := sha256.Sum256(i.DetailJSON) + parts = append(parts, fmt.Sprintf("detail_sha256:%s", hex.EncodeToString(sum[:]))) + parts = append(parts, fmt.Sprintf("detail_bytes:%d", len(i.DetailJSON))) + } + return strings.Join(parts, " ") +} + +func (a OperationalAction) valid() bool { + switch a { + case OperationalActionDeleteTenant, + OperationalActionSwitchStorageProfile, + OperationalActionDisableAudit, + OperationalActionExpandToolPermission, + OperationalActionExecuteDataMigration: + return true + default: + return false + } +} + +func (d OperationalActionDecision) valid() bool { + switch d { + case OperationalActionDecisionApprovalRequired, + OperationalActionDecisionApproved, + OperationalActionDecisionRejected, + OperationalActionDecisionExecuted, + OperationalActionDecisionFailed: + return true + default: + return false + } +} diff --git a/platform/operational_action_audit_test.go b/platform/operational_action_audit_test.go new file mode 100644 index 0000000000..1d2c9d4564 --- /dev/null +++ b/platform/operational_action_audit_test.go @@ -0,0 +1,192 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "errors" + "strings" + "testing" + "time" +) + +func TestNewOperationalActionAuditRecordBuildsSafeRecord(t *testing.T) { + createdAt := time.Date(2026, 7, 8, 10, 0, 0, 0, time.UTC) + input := OperationalActionAuditInput{ + TenantID: "tenant", + AppID: "app", + Action: OperationalActionDeleteTenant, + OperationID: "operation-1", + ResourceType: "tenant", + ResourceID: "tenant", + ActorUserID: "admin@example.com", + ActorInternalUserID: "usr_admin", + ApproverUserID: "security@example.com", + Decision: OperationalActionDecisionApproved, + DecisionReason: "secondary confirmation accepted", + RequestID: "request-1", + TraceID: "trace-1", + DetailJSON: []byte(`{"password":"plain","target":"tenant"}`), + CreatedAt: createdAt, + } + + record, err := NewOperationalActionAuditRecord(input) + if err != nil { + t.Fatalf("new operational action audit: %v", err) + } + if record.TenantID != "tenant" || record.AppID != "app" { + t.Fatalf("unexpected owner: %+v", record) + } + if record.ToolName != "ops:delete_tenant" || record.Decision != "approved" { + t.Fatalf("unexpected action decision: %+v", record) + } + if record.UserID != "" || record.UserIDHash == "" || !strings.HasPrefix(record.UserIDHash, "user_hash_") { + t.Fatalf("expected hashed actor without raw user id, got %+v", record) + } + if record.InternalUserID != "usr_admin" { + t.Fatalf("expected internal actor id, got %+v", record) + } + if !record.CreatedAt.Equal(createdAt) { + t.Fatalf("expected created_at to be retained, got %v", record.CreatedAt) + } + if record.AuditID == "" || record.RedactionVersion != "platform-operational-action-v1" { + t.Fatalf("expected audit id and redaction version, got %+v", record) + } + if !strings.Contains(record.RedactedDetailRef, "resource_type:tenant") || + !strings.Contains(record.RedactedDetailRef, "resource_hash:") || + !strings.Contains(record.RedactedDetailRef, "approver_hash:user_hash_") || + !strings.Contains(record.RedactedDetailRef, "detail_sha256:") || + !strings.Contains(record.RedactedDetailRef, "detail_bytes:38") { + t.Fatalf("unexpected redacted detail ref: %q", record.RedactedDetailRef) + } + if strings.Contains(record.RedactedDetailRef, "plain") || + strings.Contains(record.RedactedDetailRef, "password") || + strings.Contains(record.RedactedDetailRef, "tenant\"") || + strings.Contains(record.RedactedDetailRef, "admin@example.com") || + strings.Contains(record.RedactedDetailRef, "security@example.com") { + t.Fatalf("audit detail leaked raw operation context: %q", record.RedactedDetailRef) + } + + again, err := NewOperationalActionAuditRecord(input) + if err != nil { + t.Fatalf("new duplicate operational action audit: %v", err) + } + if record.AuditID != again.AuditID { + t.Fatalf("expected stable audit id, got %q and %q", record.AuditID, again.AuditID) + } + + nextOperation := input + nextOperation.OperationID = "operation-2" + nextRecord, err := NewOperationalActionAuditRecord(nextOperation) + if err != nil { + t.Fatalf("new next operational action audit: %v", err) + } + if record.AuditID == nextRecord.AuditID { + t.Fatalf("expected operation id to scope audit boundary, got %q", record.AuditID) + } +} + +func TestNewOperationalActionAuditRecordRejectsInvalidInputs(t *testing.T) { + base := validOperationalActionAuditInput() + + missingTenant := base + missingTenant.TenantID = " " + if _, err := NewOperationalActionAuditRecord(missingTenant); !errors.Is(err, ErrTenantIDRequired) { + t.Fatalf("expected tenant requirement, got %v", err) + } + + missingAction := base + missingAction.Action = " " + if _, err := NewOperationalActionAuditRecord(missingAction); err == nil || + !strings.Contains(err.Error(), "invalid operational action") { + t.Fatalf("expected action requirement, got %v", err) + } + + unknownAction := base + unknownAction.Action = "drop_prod_database" + if _, err := NewOperationalActionAuditRecord(unknownAction); err == nil || + !strings.Contains(err.Error(), "invalid operational action") { + t.Fatalf("expected unknown action rejection, got %v", err) + } + + missingOperation := base + missingOperation.OperationID = " " + if _, err := NewOperationalActionAuditRecord(missingOperation); err == nil || + !strings.Contains(err.Error(), "operation_id") { + t.Fatalf("expected operation id requirement, got %v", err) + } + + missingResource := base + missingResource.ResourceID = " " + if _, err := NewOperationalActionAuditRecord(missingResource); err == nil || + !strings.Contains(err.Error(), "resource_id") { + t.Fatalf("expected resource id requirement, got %v", err) + } + + missingActor := base + missingActor.ActorUserID = " " + missingActor.ActorInternalUserID = " " + if _, err := NewOperationalActionAuditRecord(missingActor); err == nil || + !strings.Contains(err.Error(), "actor identity") { + t.Fatalf("expected actor identity requirement, got %v", err) + } + + missingDecision := base + missingDecision.Decision = " " + if _, err := NewOperationalActionAuditRecord(missingDecision); err == nil || + !strings.Contains(err.Error(), "invalid operational action decision") { + t.Fatalf("expected decision requirement, got %v", err) + } + + unknownDecision := base + unknownDecision.Decision = "bypassed" + if _, err := NewOperationalActionAuditRecord(unknownDecision); err == nil || + !strings.Contains(err.Error(), "invalid operational action decision") { + t.Fatalf("expected unknown decision rejection, got %v", err) + } + + invalidDetail := base + invalidDetail.DetailJSON = []byte(`{"broken":`) + if _, err := NewOperationalActionAuditRecord(invalidDetail); err == nil || + !strings.Contains(err.Error(), "detail_json") { + t.Fatalf("expected detail json validation, got %v", err) + } +} + +func TestNewOperationalActionAuditRecordRejectsSensitivePublicFields(t *testing.T) { + input := validOperationalActionAuditInput() + input.DecisionReason = "Authorization: Bearer raw-token" + if _, err := NewOperationalActionAuditRecord(input); err == nil || + !strings.Contains(err.Error(), "decision_reason") { + t.Fatalf("expected sensitive decision reason rejection, got %v", err) + } + + input = validOperationalActionAuditInput() + input.Action = "sk-1234567890abcdef" + if _, err := NewOperationalActionAuditRecord(input); err == nil || + !strings.Contains(err.Error(), "invalid operational action") { + t.Fatalf("expected sensitive action rejection, got %v", err) + } +} + +func validOperationalActionAuditInput() OperationalActionAuditInput { + return OperationalActionAuditInput{ + TenantID: "tenant", + AppID: "app", + Action: OperationalActionSwitchStorageProfile, + OperationID: "operation", + ResourceType: "storage_profile", + ResourceID: "profile-a", + ActorUserID: "admin", + Decision: OperationalActionDecisionApprovalRequired, + RequestID: "request", + TraceID: "trace", + DetailJSON: []byte(`{"profile_id":"profile-a"}`), + CreatedAt: time.Now(), + } +} From 644a03c8cc56c1b4214dc504377c44dd0276f5fb Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 15:55:27 +0800 Subject: [PATCH 27/95] platform: add config cache invalidation contracts --- platform/config_cache_invalidation.go | 194 +++++++++++++++++ platform/config_cache_invalidation_test.go | 233 +++++++++++++++++++++ 2 files changed, 427 insertions(+) create mode 100644 platform/config_cache_invalidation.go create mode 100644 platform/config_cache_invalidation_test.go diff --git a/platform/config_cache_invalidation.go b/platform/config_cache_invalidation.go new file mode 100644 index 0000000000..2128f7fe95 --- /dev/null +++ b/platform/config_cache_invalidation.go @@ -0,0 +1,194 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "fmt" + "strings" + "time" +) + +// AppConfigCacheInvalidationReason explains why active config cache must be invalidated. +type AppConfigCacheInvalidationReason string + +const ( + // AppConfigCacheInvalidationReasonActivate follows a normal release activation. + AppConfigCacheInvalidationReasonActivate AppConfigCacheInvalidationReason = "activate" + // AppConfigCacheInvalidationReasonRollback follows an operational rollback. + AppConfigCacheInvalidationReasonRollback AppConfigCacheInvalidationReason = "rollback" +) + +// AppConfigCacheInvalidationInput describes one active config version switch. +type AppConfigCacheInvalidationInput struct { + PreviousVersion AppConfigVersion + NextVersion AppConfigVersion + Reason AppConfigCacheInvalidationReason + OperationID string + TraceID string + CreatedAt time.Time +} + +// AppConfigCacheInvalidation is a safe marker for invalidating active config caches. +type AppConfigCacheInvalidation struct { + TenantID string + AppID string + InvalidationID string + CacheKey string + PreviousVersion string + PreviousChecksum string + NextVersion string + NextChecksum string + Reason AppConfigCacheInvalidationReason + OperationID string + TraceID string + CreatedAt time.Time +} + +// NewAppConfigCacheInvalidation builds a cache invalidation marker for an active config switch. +func NewAppConfigCacheInvalidation(input AppConfigCacheInvalidationInput) (AppConfigCacheInvalidation, error) { + normalized, err := input.normalize() + if err != nil { + return AppConfigCacheInvalidation{}, err + } + marker := AppConfigCacheInvalidation{ + TenantID: strings.TrimSpace(normalized.NextVersion.TenantID), + AppID: strings.TrimSpace(normalized.NextVersion.AppID), + InvalidationID: normalized.invalidationID(), + CacheKey: activeConfigCacheKey(normalized.NextVersion.TenantID, normalized.NextVersion.AppID), + PreviousVersion: strings.TrimSpace(normalized.PreviousVersion.Version), + PreviousChecksum: strings.TrimSpace(normalized.PreviousVersion.Checksum), + NextVersion: strings.TrimSpace(normalized.NextVersion.Version), + NextChecksum: strings.TrimSpace(normalized.NextVersion.Checksum), + Reason: normalized.Reason, + OperationID: normalized.OperationID, + TraceID: normalized.TraceID, + CreatedAt: normalized.CreatedAt, + } + if err := marker.Validate(); err != nil { + return AppConfigCacheInvalidation{}, err + } + return marker, nil +} + +// Validate checks that a cache invalidation marker is safe to emit or store. +func (m AppConfigCacheInvalidation) Validate() error { + if strings.TrimSpace(m.TenantID) == "" { + return ErrTenantIDRequired + } + if strings.TrimSpace(m.AppID) == "" { + return ErrAppIDRequired + } + if strings.TrimSpace(m.InvalidationID) == "" { + return fmt.Errorf("invalidation_id is required") + } + if strings.TrimSpace(m.CacheKey) == "" { + return fmt.Errorf("cache_key is required") + } + if strings.TrimSpace(m.PreviousVersion) == "" { + return fmt.Errorf("previous_version is required") + } + if strings.TrimSpace(m.NextVersion) == "" { + return fmt.Errorf("next_version is required") + } + if strings.TrimSpace(m.PreviousChecksum) == "" || strings.TrimSpace(m.NextChecksum) == "" { + return fmt.Errorf("config checksums are required") + } + if !m.Reason.valid() { + return fmt.Errorf("invalid config cache invalidation reason %q", m.Reason) + } + if strings.TrimSpace(m.OperationID) == "" { + return fmt.Errorf("operation_id is required") + } + if m.CreatedAt.IsZero() { + return fmt.Errorf("created_at is required") + } + for field, value := range map[string]string{ + "invalidation_id": m.InvalidationID, + "cache_key": m.CacheKey, + "previous_version": m.PreviousVersion, + "previous_checksum": m.PreviousChecksum, + "next_version": m.NextVersion, + "next_checksum": m.NextChecksum, + "operation_id": m.OperationID, + "trace_id": m.TraceID, + } { + if err := validateAuditRedactedText(field, value); err != nil { + return err + } + } + return nil +} + +func (i AppConfigCacheInvalidationInput) normalize() (AppConfigCacheInvalidationInput, error) { + if err := i.PreviousVersion.Validate(); err != nil { + return AppConfigCacheInvalidationInput{}, fmt.Errorf("previous config version: %w", err) + } + if err := i.NextVersion.Validate(); err != nil { + return AppConfigCacheInvalidationInput{}, fmt.Errorf("next config version: %w", err) + } + if err := requireSameConfigOwner(i.PreviousVersion, i.NextVersion); err != nil { + return AppConfigCacheInvalidationInput{}, err + } + if i.NextVersion.Status != AppConfigVersionStatusActive { + return AppConfigCacheInvalidationInput{}, fmt.Errorf("next config version status must be active") + } + if strings.TrimSpace(i.PreviousVersion.Version) == strings.TrimSpace(i.NextVersion.Version) { + return AppConfigCacheInvalidationInput{}, fmt.Errorf("config version switch must change version") + } + i.Reason = AppConfigCacheInvalidationReason(strings.TrimSpace(string(i.Reason))) + if !i.Reason.valid() { + return AppConfigCacheInvalidationInput{}, fmt.Errorf("invalid config cache invalidation reason %q", i.Reason) + } + i.OperationID = strings.TrimSpace(i.OperationID) + if i.OperationID == "" { + return AppConfigCacheInvalidationInput{}, fmt.Errorf("operation_id is required") + } + i.TraceID = strings.TrimSpace(i.TraceID) + if i.CreatedAt.IsZero() { + return AppConfigCacheInvalidationInput{}, fmt.Errorf("created_at is required") + } + for field, value := range map[string]string{ + "operation_id": i.OperationID, + "trace_id": i.TraceID, + } { + if err := validateAuditRedactedText(field, value); err != nil { + return AppConfigCacheInvalidationInput{}, err + } + } + return i, nil +} + +func (i AppConfigCacheInvalidationInput) invalidationID() string { + return "config_invalidation_" + shortHash( + strings.TrimSpace(i.PreviousVersion.TenantID), + strings.TrimSpace(i.PreviousVersion.AppID), + strings.TrimSpace(i.PreviousVersion.Version), + strings.TrimSpace(i.NextVersion.Version), + string(i.Reason), + i.OperationID, + ) +} + +func activeConfigCacheKey(tenantID, appID string) string { + return strings.Join([]string{ + "tenant", escapeKeyPart(tenantID), + "app", escapeKeyPart(appID), + "config", "active", + }, ":") +} + +func (r AppConfigCacheInvalidationReason) valid() bool { + switch r { + case AppConfigCacheInvalidationReasonActivate, + AppConfigCacheInvalidationReasonRollback: + return true + default: + return false + } +} diff --git a/platform/config_cache_invalidation_test.go b/platform/config_cache_invalidation_test.go new file mode 100644 index 0000000000..a85319bd1e --- /dev/null +++ b/platform/config_cache_invalidation_test.go @@ -0,0 +1,233 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "errors" + "fmt" + "strings" + "testing" + "time" +) + +func TestNewAppConfigCacheInvalidationBuildsRollbackMarker(t *testing.T) { + createdAt := time.Date(2026, 7, 8, 12, 0, 0, 0, time.UTC) + previous := validLifecycleConfigVersion("v2", AppConfigVersionStatusRollback) + previous.Checksum = "sha256:previous" + next := validLifecycleConfigVersion("v1", AppConfigVersionStatusActive) + next.Checksum = "sha256:next" + + marker, err := NewAppConfigCacheInvalidation(AppConfigCacheInvalidationInput{ + PreviousVersion: previous, + NextVersion: next, + Reason: AppConfigCacheInvalidationReasonRollback, + OperationID: "rollback-1", + TraceID: "trace-1", + CreatedAt: createdAt, + }) + if err != nil { + t.Fatalf("new config cache invalidation: %v", err) + } + if marker.TenantID != "tenant" || marker.AppID != "app" { + t.Fatalf("unexpected owner: %+v", marker) + } + if marker.CacheKey != "tenant:tenant:app:app:config:active" { + t.Fatalf("unexpected cache key: %q", marker.CacheKey) + } + if marker.PreviousVersion != "v2" || marker.PreviousChecksum != "sha256:previous" || + marker.NextVersion != "v1" || marker.NextChecksum != "sha256:next" { + t.Fatalf("unexpected version switch summary: %+v", marker) + } + if marker.Reason != AppConfigCacheInvalidationReasonRollback || + marker.OperationID != "rollback-1" || marker.TraceID != "trace-1" || + !marker.CreatedAt.Equal(createdAt) { + t.Fatalf("unexpected marker metadata: %+v", marker) + } + if !strings.HasPrefix(marker.InvalidationID, "config_invalidation_") { + t.Fatalf("unexpected invalidation id: %q", marker.InvalidationID) + } + serialized := fmt.Sprintf("%+v", marker) + if strings.Contains(serialized, "model_profile_id") || + strings.Contains(serialized, "tool_policy_id") || + strings.Contains(serialized, "api_key_ref") { + t.Fatalf("marker leaked config bundle content: %s", serialized) + } + + again, err := NewAppConfigCacheInvalidation(AppConfigCacheInvalidationInput{ + PreviousVersion: previous, + NextVersion: next, + Reason: AppConfigCacheInvalidationReasonRollback, + OperationID: "rollback-1", + TraceID: "trace-2", + CreatedAt: createdAt.Add(time.Minute), + }) + if err != nil { + t.Fatalf("new duplicate config cache invalidation: %v", err) + } + if marker.InvalidationID != again.InvalidationID { + t.Fatalf("expected stable invalidation id, got %q and %q", marker.InvalidationID, again.InvalidationID) + } + + nextOperation, err := NewAppConfigCacheInvalidation(AppConfigCacheInvalidationInput{ + PreviousVersion: previous, + NextVersion: next, + Reason: AppConfigCacheInvalidationReasonRollback, + OperationID: "rollback-2", + CreatedAt: createdAt, + }) + if err != nil { + t.Fatalf("new next config cache invalidation: %v", err) + } + if marker.InvalidationID == nextOperation.InvalidationID { + t.Fatalf("expected operation id to scope invalidation id, got %q", marker.InvalidationID) + } + + whitespace := previous + whitespace.TenantID = " tenant " + whitespace.AppID = " app " + whitespace.Version = " v2 " + whitespace.Checksum = " sha256:previous " + whitespaceNext := next + whitespaceNext.TenantID = " tenant " + whitespaceNext.AppID = " app " + whitespaceNext.Version = " v1 " + whitespaceNext.Checksum = " sha256:next " + trimmed, err := NewAppConfigCacheInvalidation(AppConfigCacheInvalidationInput{ + PreviousVersion: whitespace, + NextVersion: whitespaceNext, + Reason: AppConfigCacheInvalidationReasonRollback, + OperationID: "rollback-1", + CreatedAt: createdAt, + }) + if err != nil { + t.Fatalf("new trimmed config cache invalidation: %v", err) + } + if marker.InvalidationID != trimmed.InvalidationID { + t.Fatalf("expected trimmed identity to keep stable invalidation id, got %q and %q", marker.InvalidationID, trimmed.InvalidationID) + } + if trimmed.PreviousVersion != "v2" || trimmed.NextVersion != "v1" || + trimmed.PreviousChecksum != "sha256:previous" || trimmed.NextChecksum != "sha256:next" { + t.Fatalf("expected marker fields to be trimmed, got %+v", trimmed) + } +} + +func TestNewAppConfigCacheInvalidationBuildsActivationMarker(t *testing.T) { + previous := validLifecycleConfigVersion("v1", AppConfigVersionStatusRollback) + next := validLifecycleConfigVersion("v2", AppConfigVersionStatusActive) + + marker, err := NewAppConfigCacheInvalidation(AppConfigCacheInvalidationInput{ + PreviousVersion: previous, + NextVersion: next, + Reason: AppConfigCacheInvalidationReasonActivate, + OperationID: "activate-1", + CreatedAt: time.Now(), + }) + if err != nil { + t.Fatalf("new activation invalidation: %v", err) + } + if marker.Reason != AppConfigCacheInvalidationReasonActivate || + marker.PreviousVersion != "v1" || marker.NextVersion != "v2" { + t.Fatalf("unexpected activation marker: %+v", marker) + } +} + +func TestNewAppConfigCacheInvalidationRejectsInvalidInputs(t *testing.T) { + base := validAppConfigCacheInvalidationInput() + + missingTenant := base + missingTenant.NextVersion.TenantID = " " + if _, err := NewAppConfigCacheInvalidation(missingTenant); !errors.Is(err, ErrTenantIDRequired) { + t.Fatalf("expected tenant requirement, got %v", err) + } + + mismatch := base + mismatch.NextVersion.AppID = "other-app" + if _, err := NewAppConfigCacheInvalidation(mismatch); err == nil || + !strings.Contains(err.Error(), "app_id") { + t.Fatalf("expected app mismatch error, got %v", err) + } + + nextNotActive := base + nextNotActive.NextVersion.Status = AppConfigVersionStatusReleased + if _, err := NewAppConfigCacheInvalidation(nextNotActive); err == nil || + !strings.Contains(err.Error(), "active") { + t.Fatalf("expected next active status error, got %v", err) + } + + sameVersion := base + sameVersion.NextVersion.Version = sameVersion.PreviousVersion.Version + if _, err := NewAppConfigCacheInvalidation(sameVersion); err == nil || + !strings.Contains(err.Error(), "change version") { + t.Fatalf("expected version switch validation, got %v", err) + } + + missingReason := base + missingReason.Reason = " " + if _, err := NewAppConfigCacheInvalidation(missingReason); err == nil || + !strings.Contains(err.Error(), "invalid config cache invalidation reason") { + t.Fatalf("expected reason validation, got %v", err) + } + + missingOperation := base + missingOperation.OperationID = " " + if _, err := NewAppConfigCacheInvalidation(missingOperation); err == nil || + !strings.Contains(err.Error(), "operation_id") { + t.Fatalf("expected operation id requirement, got %v", err) + } + + zeroCreatedAt := base + zeroCreatedAt.CreatedAt = time.Time{} + if _, err := NewAppConfigCacheInvalidation(zeroCreatedAt); err == nil || + !strings.Contains(err.Error(), "created_at") { + t.Fatalf("expected created_at requirement, got %v", err) + } + + sensitiveTrace := base + sensitiveTrace.TraceID = "Authorization: Bearer raw-token" + if _, err := NewAppConfigCacheInvalidation(sensitiveTrace); err == nil || + !strings.Contains(err.Error(), "trace_id") { + t.Fatalf("expected sensitive trace id rejection, got %v", err) + } +} + +func TestAppConfigCacheInvalidationValidateRejectsUnsafeMarker(t *testing.T) { + marker := AppConfigCacheInvalidation{ + TenantID: "tenant", + AppID: "app", + InvalidationID: "config_invalidation_id", + CacheKey: "tenant:tenant:app:app:config:active", + PreviousVersion: "v1", + PreviousChecksum: "sha256:previous", + NextVersion: "v2", + NextChecksum: "sha256:next", + Reason: AppConfigCacheInvalidationReasonActivate, + OperationID: "operation", + TraceID: "trace", + CreatedAt: time.Now(), + } + if err := marker.Validate(); err != nil { + t.Fatalf("expected marker to validate: %v", err) + } + + marker.OperationID = "sk-1234567890abcdef" + if err := marker.Validate(); err == nil || !strings.Contains(err.Error(), "operation_id") { + t.Fatalf("expected operation id sensitive content rejection, got %v", err) + } +} + +func validAppConfigCacheInvalidationInput() AppConfigCacheInvalidationInput { + return AppConfigCacheInvalidationInput{ + PreviousVersion: validLifecycleConfigVersion("v1", AppConfigVersionStatusRollback), + NextVersion: validLifecycleConfigVersion("v2", AppConfigVersionStatusActive), + Reason: AppConfigCacheInvalidationReasonActivate, + OperationID: "operation", + TraceID: "trace", + CreatedAt: time.Now(), + } +} From 8219105a004a1cf2d21e67183970a5453d7dfec3 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 16:06:53 +0800 Subject: [PATCH 28/95] platform: add secret rotation status contracts --- platform/secret_rotation_status.go | 281 ++++++++++++++++++++++++ platform/secret_rotation_status_test.go | 246 +++++++++++++++++++++ 2 files changed, 527 insertions(+) create mode 100644 platform/secret_rotation_status.go create mode 100644 platform/secret_rotation_status_test.go diff --git a/platform/secret_rotation_status.go b/platform/secret_rotation_status.go new file mode 100644 index 0000000000..6c193bda91 --- /dev/null +++ b/platform/secret_rotation_status.go @@ -0,0 +1,281 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "fmt" + "strings" + "time" +) + +const secretRotationIDPrefix = "secret_rotation_" + +// SecretRotationStatus describes the lifecycle state of one secret rotation. +type SecretRotationStatus string + +const ( + // SecretRotationStatusPending means the new secret reference has been registered but not verified. + SecretRotationStatusPending SecretRotationStatus = "pending" + // SecretRotationStatusVerifying means dependent systems are validating the new secret reference. + SecretRotationStatusVerifying SecretRotationStatus = "verifying" + // SecretRotationStatusReady means the new secret reference is ready for cutover. + SecretRotationStatusReady SecretRotationStatus = "ready" + // SecretRotationStatusActive means traffic has moved to the new secret reference. + SecretRotationStatusActive SecretRotationStatus = "active" + // SecretRotationStatusRolledBack means traffic has returned to the previous secret reference. + SecretRotationStatusRolledBack SecretRotationStatus = "rolled_back" + // SecretRotationStatusFailed means the rotation failed before completion. + SecretRotationStatusFailed SecretRotationStatus = "failed" +) + +// SecretRotationStatusInput contains safe metadata for one secret rotation status update. +type SecretRotationStatusInput struct { + TenantID string + AppID string + ResourceType string + ResourceID string + SecretField string + PreviousRef string + NextRef string + Status SecretRotationStatus + OperationID string + FailureReason string + TraceID string + UpdatedAt time.Time +} + +// SecretRotationStatusReport is a safe, operations-facing secret rotation status. +type SecretRotationStatusReport struct { + TenantID string + AppID string + RotationID string + ResourceType string + ResourceHash string + SecretField string + PreviousRef string + NextRef string + Status SecretRotationStatus + OperationID string + FailureReason string + TraceID string + UpdatedAt time.Time +} + +// NewSecretRotationStatusReport builds a safe status report for secret rotation observability. +func NewSecretRotationStatusReport(input SecretRotationStatusInput) (SecretRotationStatusReport, error) { + normalized, err := input.normalize() + if err != nil { + return SecretRotationStatusReport{}, err + } + report := SecretRotationStatusReport{ + TenantID: normalized.TenantID, + AppID: normalized.AppID, + RotationID: normalized.rotationID(), + ResourceType: normalized.ResourceType, + ResourceHash: shortHash(normalized.TenantID, normalized.ResourceType, normalized.ResourceID), + SecretField: normalized.SecretField, + PreviousRef: normalized.PreviousRef, + NextRef: normalized.NextRef, + Status: normalized.Status, + OperationID: normalized.OperationID, + FailureReason: normalized.FailureReason, + TraceID: normalized.TraceID, + UpdatedAt: normalized.UpdatedAt, + } + if err := report.Validate(); err != nil { + return SecretRotationStatusReport{}, err + } + return report, nil +} + +// Validate checks that a secret rotation status report is safe to expose or store. +func (r SecretRotationStatusReport) Validate() error { + if strings.TrimSpace(r.TenantID) == "" { + return ErrTenantIDRequired + } + if strings.TrimSpace(r.RotationID) == "" { + return fmt.Errorf("rotation_id is required") + } + if !isSecretRotationID(r.RotationID) { + return fmt.Errorf("rotation_id must be %s followed by a 24 character hex hash", secretRotationIDPrefix) + } + if strings.TrimSpace(r.ResourceType) == "" { + return fmt.Errorf("resource_type is required") + } + if strings.TrimSpace(r.ResourceHash) == "" { + return fmt.Errorf("resource_hash is required") + } + if !isShortHash(r.ResourceHash) { + return fmt.Errorf("resource_hash must be a 24 character hex hash") + } + if strings.TrimSpace(r.SecretField) == "" { + return fmt.Errorf("secret_field is required") + } + if err := validateRotationSecretReference("previous_ref", r.PreviousRef); err != nil { + return err + } + if err := validateRotationSecretReference("next_ref", r.NextRef); err != nil { + return err + } + if strings.TrimSpace(r.NextRef) == "" { + return fmt.Errorf("next_ref is required") + } + if !r.Status.valid() { + return fmt.Errorf("invalid secret rotation status %q", r.Status) + } + if strings.TrimSpace(r.OperationID) == "" { + return fmt.Errorf("operation_id is required") + } + if r.UpdatedAt.IsZero() { + return fmt.Errorf("updated_at is required") + } + for field, value := range map[string]string{ + "app_id": r.AppID, + "rotation_id": r.RotationID, + "resource_type": r.ResourceType, + "resource_hash": r.ResourceHash, + "secret_field": r.SecretField, + "operation_id": r.OperationID, + "failure_reason": r.FailureReason, + "trace_id": r.TraceID, + } { + if err := validateAuditRedactedText(field, value); err != nil { + return err + } + } + return nil +} + +func (i SecretRotationStatusInput) normalize() (SecretRotationStatusInput, error) { + i.TenantID = strings.TrimSpace(i.TenantID) + if i.TenantID == "" { + return SecretRotationStatusInput{}, ErrTenantIDRequired + } + i.AppID = strings.TrimSpace(i.AppID) + i.ResourceType = strings.TrimSpace(i.ResourceType) + if i.ResourceType == "" { + return SecretRotationStatusInput{}, fmt.Errorf("resource_type is required") + } + i.ResourceID = strings.TrimSpace(i.ResourceID) + if i.ResourceID == "" { + return SecretRotationStatusInput{}, fmt.Errorf("resource_id is required") + } + i.SecretField = strings.TrimSpace(i.SecretField) + if i.SecretField == "" { + return SecretRotationStatusInput{}, fmt.Errorf("secret_field is required") + } + i.PreviousRef = strings.TrimSpace(i.PreviousRef) + i.NextRef = strings.TrimSpace(i.NextRef) + if err := validateRotationSecretReference("previous_ref", i.PreviousRef); err != nil { + return SecretRotationStatusInput{}, err + } + if err := validateRotationSecretReference("next_ref", i.NextRef); err != nil { + return SecretRotationStatusInput{}, err + } + if i.NextRef == "" { + return SecretRotationStatusInput{}, fmt.Errorf("next_ref is required") + } + i.Status = SecretRotationStatus(strings.TrimSpace(string(i.Status))) + if !i.Status.valid() { + return SecretRotationStatusInput{}, fmt.Errorf("invalid secret rotation status %q", i.Status) + } + i.OperationID = strings.TrimSpace(i.OperationID) + if i.OperationID == "" { + return SecretRotationStatusInput{}, fmt.Errorf("operation_id is required") + } + i.FailureReason = strings.TrimSpace(i.FailureReason) + i.TraceID = strings.TrimSpace(i.TraceID) + if i.UpdatedAt.IsZero() { + return SecretRotationStatusInput{}, fmt.Errorf("updated_at is required") + } + for field, value := range map[string]string{ + "app_id": i.AppID, + "resource_type": i.ResourceType, + "secret_field": i.SecretField, + "operation_id": i.OperationID, + "failure_reason": i.FailureReason, + "trace_id": i.TraceID, + } { + if err := validateAuditRedactedText(field, value); err != nil { + return SecretRotationStatusInput{}, err + } + } + return i, nil +} + +func (i SecretRotationStatusInput) rotationID() string { + return secretRotationIDPrefix + shortHash( + i.TenantID, + i.AppID, + i.ResourceType, + i.ResourceID, + i.SecretField, + i.OperationID, + ) +} + +func (s SecretRotationStatus) valid() bool { + switch s { + case SecretRotationStatusPending, + SecretRotationStatusVerifying, + SecretRotationStatusReady, + SecretRotationStatusActive, + SecretRotationStatusRolledBack, + SecretRotationStatusFailed: + return true + default: + return false + } +} + +func validateRotationSecretReference(field, value string) error { + value = strings.TrimSpace(value) + if value == "" { + return nil + } + if err := validateSecretReference(field, value); err != nil { + return err + } + if !isAllowedSecretReference(value) { + return fmt.Errorf("%s must use secret://, kms://, or vault:// reference format", field) + } + return nil +} + +func isAllowedSecretReference(value string) bool { + switch { + case strings.HasPrefix(value, "secret://"): + return len(strings.TrimPrefix(value, "secret://")) > 0 + case strings.HasPrefix(value, "kms://"): + return len(strings.TrimPrefix(value, "kms://")) > 0 + case strings.HasPrefix(value, "vault://"): + return len(strings.TrimPrefix(value, "vault://")) > 0 + default: + return false + } +} + +func isSecretRotationID(value string) bool { + if !strings.HasPrefix(value, secretRotationIDPrefix) { + return false + } + return isShortHash(strings.TrimPrefix(value, secretRotationIDPrefix)) +} + +func isShortHash(value string) bool { + if len(value) != 24 { + return false + } + for _, r := range value { + if (r < '0' || r > '9') && (r < 'a' || r > 'f') { + return false + } + } + return true +} diff --git a/platform/secret_rotation_status_test.go b/platform/secret_rotation_status_test.go new file mode 100644 index 0000000000..cebb0e0ed8 --- /dev/null +++ b/platform/secret_rotation_status_test.go @@ -0,0 +1,246 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "errors" + "fmt" + "strings" + "testing" + "time" +) + +func TestNewSecretRotationStatusReportBuildsSafeReport(t *testing.T) { + updatedAt := time.Date(2026, 7, 8, 13, 0, 0, 0, time.UTC) + input := SecretRotationStatusInput{ + TenantID: "tenant", + AppID: "app", + ResourceType: "model_profile", + ResourceID: "profile-a", + SecretField: "api_key_ref", + PreviousRef: "secret://model-key-v1", + NextRef: "kms://tenant/model-key-v2", + Status: SecretRotationStatusReady, + OperationID: "rotation-1", + FailureReason: "verification passed", + TraceID: "trace-1", + UpdatedAt: updatedAt, + } + + report, err := NewSecretRotationStatusReport(input) + if err != nil { + t.Fatalf("new secret rotation status report: %v", err) + } + if report.TenantID != "tenant" || report.AppID != "app" { + t.Fatalf("unexpected owner: %+v", report) + } + if report.ResourceType != "model_profile" || report.ResourceHash == "" || + report.ResourceHash == "profile-a" { + t.Fatalf("expected resource hash without raw resource id, got %+v", report) + } + if report.SecretField != "api_key_ref" || + report.PreviousRef != "secret://model-key-v1" || + report.NextRef != "kms://tenant/model-key-v2" || + report.Status != SecretRotationStatusReady { + t.Fatalf("unexpected report fields: %+v", report) + } + if report.OperationID != "rotation-1" || report.TraceID != "trace-1" || + !report.UpdatedAt.Equal(updatedAt) { + t.Fatalf("unexpected operation metadata: %+v", report) + } + if !strings.HasPrefix(report.RotationID, "secret_rotation_") { + t.Fatalf("unexpected rotation id: %q", report.RotationID) + } + serialized := fmt.Sprintf("%+v", report) + if strings.Contains(serialized, "profile-a") || + strings.Contains(serialized, "plain-secret") { + t.Fatalf("report leaked raw resource or secret content: %s", serialized) + } + + again, err := NewSecretRotationStatusReport(input) + if err != nil { + t.Fatalf("new duplicate secret rotation status report: %v", err) + } + if report.RotationID != again.RotationID { + t.Fatalf("expected stable rotation id, got %q and %q", report.RotationID, again.RotationID) + } + + nextOperation := input + nextOperation.OperationID = "rotation-2" + nextReport, err := NewSecretRotationStatusReport(nextOperation) + if err != nil { + t.Fatalf("new next secret rotation status report: %v", err) + } + if report.RotationID == nextReport.RotationID { + t.Fatalf("expected operation id to scope rotation id, got %q", report.RotationID) + } +} + +func TestNewSecretRotationStatusReportRejectsInvalidInputs(t *testing.T) { + base := validSecretRotationStatusInput() + + missingTenant := base + missingTenant.TenantID = " " + if _, err := NewSecretRotationStatusReport(missingTenant); !errors.Is(err, ErrTenantIDRequired) { + t.Fatalf("expected tenant requirement, got %v", err) + } + + missingResource := base + missingResource.ResourceID = " " + if _, err := NewSecretRotationStatusReport(missingResource); err == nil || + !strings.Contains(err.Error(), "resource_id") { + t.Fatalf("expected resource id requirement, got %v", err) + } + + missingField := base + missingField.SecretField = " " + if _, err := NewSecretRotationStatusReport(missingField); err == nil || + !strings.Contains(err.Error(), "secret_field") { + t.Fatalf("expected secret field requirement, got %v", err) + } + + missingNextRef := base + missingNextRef.NextRef = " " + if _, err := NewSecretRotationStatusReport(missingNextRef); err == nil || + !strings.Contains(err.Error(), "next_ref") { + t.Fatalf("expected next ref requirement, got %v", err) + } + + inlineSecret := base + inlineSecret.NextRef = "sk-1234567890abcdef" + if _, err := NewSecretRotationStatusReport(inlineSecret); !errors.Is(err, ErrInlineSecretRejected) { + t.Fatalf("expected inline secret rejection, got %v", err) + } + + unknownStatus := base + unknownStatus.Status = "skipped" + if _, err := NewSecretRotationStatusReport(unknownStatus); err == nil || + !strings.Contains(err.Error(), "invalid secret rotation status") { + t.Fatalf("expected status validation, got %v", err) + } + + missingOperation := base + missingOperation.OperationID = " " + if _, err := NewSecretRotationStatusReport(missingOperation); err == nil || + !strings.Contains(err.Error(), "operation_id") { + t.Fatalf("expected operation id requirement, got %v", err) + } + + zeroUpdatedAt := base + zeroUpdatedAt.UpdatedAt = time.Time{} + if _, err := NewSecretRotationStatusReport(zeroUpdatedAt); err == nil || + !strings.Contains(err.Error(), "updated_at") { + t.Fatalf("expected updated at requirement, got %v", err) + } + + sensitiveFailure := base + sensitiveFailure.FailureReason = "password=plain" + if _, err := NewSecretRotationStatusReport(sensitiveFailure); err == nil || + !strings.Contains(err.Error(), "failure_reason") { + t.Fatalf("expected sensitive failure reason rejection, got %v", err) + } +} + +func TestSecretRotationStatusReportValidateRejectsUnsafeReport(t *testing.T) { + generated, err := NewSecretRotationStatusReport(validSecretRotationStatusInput()) + if err != nil { + t.Fatalf("new generated report: %v", err) + } + report := SecretRotationStatusReport{ + TenantID: "tenant", + AppID: "app", + RotationID: generated.RotationID, + ResourceType: "channel_binding", + ResourceHash: generated.ResourceHash, + SecretField: "token_ref", + PreviousRef: "secret://token-v1", + NextRef: "secret://token-v2", + Status: SecretRotationStatusActive, + OperationID: "rotation", + FailureReason: "cutover completed", + TraceID: "trace", + UpdatedAt: time.Now(), + } + if err := report.Validate(); err != nil { + t.Fatalf("expected report to validate: %v", err) + } + + report.NextRef = "postgres://user:password@example.com/db" + if err := report.Validate(); !errors.Is(err, ErrInlineSecretRejected) { + t.Fatalf("expected unsafe next ref rejection, got %v", err) + } + + report.NextRef = "secret://token-v2" + report.ResourceHash = "binding-a" + if err := report.Validate(); err == nil || + !strings.Contains(err.Error(), "resource_hash") { + t.Fatalf("expected raw resource hash rejection, got %v", err) + } + + report.ResourceHash = generated.ResourceHash + report.RotationID = "secret_rotation_binding-a" + if err := report.Validate(); err == nil || + !strings.Contains(err.Error(), "rotation_id") { + t.Fatalf("expected unsafe rotation id rejection, got %v", err) + } +} + +func TestNewSecretRotationStatusReportRequiresSafeReferenceFormat(t *testing.T) { + base := validSecretRotationStatusInput() + + plaintext := base + plaintext.NextRef = "ordinary-token-value" + if _, err := NewSecretRotationStatusReport(plaintext); err == nil || + !strings.Contains(err.Error(), "next_ref") { + t.Fatalf("expected plaintext next ref rejection, got %v", err) + } + + unknownScheme := base + unknownScheme.NextRef = "file://tenant/token" + if _, err := NewSecretRotationStatusReport(unknownScheme); err == nil || + !strings.Contains(err.Error(), "next_ref") { + t.Fatalf("expected unknown scheme rejection, got %v", err) + } + + unsafePrevious := base + unsafePrevious.PreviousRef = "plain-previous-token" + if _, err := NewSecretRotationStatusReport(unsafePrevious); err == nil || + !strings.Contains(err.Error(), "previous_ref") { + t.Fatalf("expected plaintext previous ref rejection, got %v", err) + } + + for _, nextRef := range []string{ + "secret://token-v2", + "kms://tenant/token-v2", + "vault://secret/data/token-v2", + } { + input := base + input.PreviousRef = "" + input.NextRef = nextRef + if _, err := NewSecretRotationStatusReport(input); err != nil { + t.Fatalf("expected %q to validate: %v", nextRef, err) + } + } +} + +func validSecretRotationStatusInput() SecretRotationStatusInput { + return SecretRotationStatusInput{ + TenantID: "tenant", + AppID: "app", + ResourceType: "channel_binding", + ResourceID: "binding-a", + SecretField: "token_ref", + PreviousRef: "secret://token-v1", + NextRef: "secret://token-v2", + Status: SecretRotationStatusPending, + OperationID: "rotation", + TraceID: "trace", + UpdatedAt: time.Now(), + } +} From e66bf9caa54e9c4bf02f8a36ec7000fca9e246f0 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 16:19:39 +0800 Subject: [PATCH 29/95] platform: add backend migration status contracts --- platform/backend_migration_status.go | 442 ++++++++++++++++++++++ platform/backend_migration_status_test.go | 382 +++++++++++++++++++ 2 files changed, 824 insertions(+) create mode 100644 platform/backend_migration_status.go create mode 100644 platform/backend_migration_status_test.go diff --git a/platform/backend_migration_status.go b/platform/backend_migration_status.go new file mode 100644 index 0000000000..67ab59b735 --- /dev/null +++ b/platform/backend_migration_status.go @@ -0,0 +1,442 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "fmt" + "strings" + "time" +) + +const backendMigrationIDPrefix = "backend_migration_" + +// BackendMigrationResource names the storage resource being migrated. +type BackendMigrationResource string + +const ( + // BackendMigrationResourceSession covers session event storage migrations. + BackendMigrationResourceSession BackendMigrationResource = "session" + // BackendMigrationResourceMemory covers memory store migrations. + BackendMigrationResourceMemory BackendMigrationResource = "memory" + // BackendMigrationResourceArtifact covers artifact object storage migrations. + BackendMigrationResourceArtifact BackendMigrationResource = "artifact" + // BackendMigrationResourceKnowledge covers knowledge/vector store migrations. + BackendMigrationResourceKnowledge BackendMigrationResource = "knowledge" + // BackendMigrationResourceAudit covers audit sink migrations. + BackendMigrationResourceAudit BackendMigrationResource = "audit" +) + +// BackendMigrationStatus describes the lifecycle state of one backend migration task. +type BackendMigrationStatus string + +const ( + // BackendMigrationStatusPending means the task is registered but has not started. + BackendMigrationStatusPending BackendMigrationStatus = "pending" + // BackendMigrationStatusRunning means records are being copied or dual-written. + BackendMigrationStatusRunning BackendMigrationStatus = "running" + // BackendMigrationStatusVerifying means source and target data are being compared. + BackendMigrationStatusVerifying BackendMigrationStatus = "verifying" + // BackendMigrationStatusReady means verification is clean enough for cutover. + BackendMigrationStatusReady BackendMigrationStatus = "ready" + // BackendMigrationStatusCompleted means cutover completed successfully. + BackendMigrationStatusCompleted BackendMigrationStatus = "completed" + // BackendMigrationStatusRolledBack means traffic returned to the source backend. + BackendMigrationStatusRolledBack BackendMigrationStatus = "rolled_back" + // BackendMigrationStatusFailed means the task failed before a safe completion. + BackendMigrationStatusFailed BackendMigrationStatus = "failed" +) + +// BackendMigrationStatusInput contains safe metadata for one backend migration update. +type BackendMigrationStatusInput struct { + TenantID string + AppID string + ProfileID string + Resource BackendMigrationResource + SourceBackendID string + TargetBackendID string + MigrationMode StorageMigrationMode + Status BackendMigrationStatus + OperationID string + SourceRecordCount int64 + TargetRecordCount int64 + VerifiedRecordCount int64 + MismatchCount int64 + LastRecordID string + SampleSetRef string + SampledTopKQueries int64 + MatchedTopKQueries int64 + FailureReason string + TraceID string + UpdatedAt time.Time +} + +// BackendMigrationStatusReport is a safe operations-facing backend migration task status. +type BackendMigrationStatusReport struct { + TenantID string + AppID string + ProfileID string + MigrationID string + Resource BackendMigrationResource + SourceBackendID string + TargetBackendID string + MigrationMode StorageMigrationMode + Status BackendMigrationStatus + OperationID string + SourceRecordCount int64 + TargetRecordCount int64 + LagRecordCount int64 + VerifiedRecordCount int64 + MismatchCount int64 + LastRecordID string + SampleSetRef string + SampledTopKQueries int64 + MatchedTopKQueries int64 + FailureReason string + TraceID string + UpdatedAt time.Time +} + +// NewBackendMigrationStatusReport builds a safe status report for backend migration observability. +func NewBackendMigrationStatusReport(input BackendMigrationStatusInput) (BackendMigrationStatusReport, error) { + normalized, err := input.normalize() + if err != nil { + return BackendMigrationStatusReport{}, err + } + report := BackendMigrationStatusReport{ + TenantID: normalized.TenantID, + AppID: normalized.AppID, + ProfileID: normalized.ProfileID, + MigrationID: normalized.migrationID(), + Resource: normalized.Resource, + SourceBackendID: normalized.SourceBackendID, + TargetBackendID: normalized.TargetBackendID, + MigrationMode: normalized.MigrationMode, + Status: normalized.Status, + OperationID: normalized.OperationID, + SourceRecordCount: normalized.SourceRecordCount, + TargetRecordCount: normalized.TargetRecordCount, + LagRecordCount: backendMigrationLag(normalized.SourceRecordCount, normalized.TargetRecordCount), + VerifiedRecordCount: normalized.VerifiedRecordCount, + MismatchCount: normalized.MismatchCount, + LastRecordID: normalized.LastRecordID, + SampleSetRef: normalized.SampleSetRef, + SampledTopKQueries: normalized.SampledTopKQueries, + MatchedTopKQueries: normalized.MatchedTopKQueries, + FailureReason: normalized.FailureReason, + TraceID: normalized.TraceID, + UpdatedAt: normalized.UpdatedAt, + } + if err := report.Validate(); err != nil { + return BackendMigrationStatusReport{}, err + } + return report, nil +} + +// Validate checks that a backend migration status report is safe to expose or store. +func (r BackendMigrationStatusReport) Validate() error { + if strings.TrimSpace(r.TenantID) == "" { + return ErrTenantIDRequired + } + if strings.TrimSpace(r.ProfileID) == "" { + return fmt.Errorf("profile_id is required") + } + if strings.TrimSpace(r.MigrationID) == "" { + return fmt.Errorf("migration_id is required") + } + if !isBackendMigrationID(r.MigrationID) { + return fmt.Errorf("migration_id must be %s followed by a 24 character hex hash", backendMigrationIDPrefix) + } + if r.MigrationID != backendMigrationID( + r.TenantID, + r.AppID, + r.ProfileID, + r.Resource, + r.SourceBackendID, + r.TargetBackendID, + r.OperationID, + ) { + return fmt.Errorf("migration_id does not match backend migration identity") + } + if !r.Resource.valid() { + return fmt.Errorf("invalid backend migration resource %q", r.Resource) + } + if strings.TrimSpace(r.SourceBackendID) == "" { + return fmt.Errorf("source_backend_id is required") + } + if strings.TrimSpace(r.TargetBackendID) == "" { + return fmt.Errorf("target_backend_id is required") + } + if strings.TrimSpace(r.SourceBackendID) == strings.TrimSpace(r.TargetBackendID) { + return fmt.Errorf("source_backend_id and target_backend_id must differ") + } + mode, err := NormalizeStorageMigrationMode(string(r.MigrationMode)) + if err != nil { + return err + } + if !IsActiveStorageMigrationMode(mode) { + return fmt.Errorf("migration_mode must be an active migration mode") + } + if !r.Status.valid() { + return fmt.Errorf("invalid backend migration status %q", r.Status) + } + if strings.TrimSpace(r.OperationID) == "" { + return fmt.Errorf("operation_id is required") + } + if err := validateBackendMigrationCounts(r); err != nil { + return err + } + if err := validateBackendMigrationStatusGate(r); err != nil { + return err + } + if r.Status == BackendMigrationStatusFailed && strings.TrimSpace(r.FailureReason) == "" { + return fmt.Errorf("failure_reason is required for failed backend migration status") + } + if r.UpdatedAt.IsZero() { + return fmt.Errorf("updated_at is required") + } + for field, value := range map[string]string{ + "app_id": r.AppID, + "profile_id": r.ProfileID, + "migration_id": r.MigrationID, + "source_backend_id": r.SourceBackendID, + "target_backend_id": r.TargetBackendID, + "operation_id": r.OperationID, + "last_record_id": r.LastRecordID, + "sample_set_ref": r.SampleSetRef, + "failure_reason": r.FailureReason, + "trace_id": r.TraceID, + } { + if err := validateAuditRedactedText(field, value); err != nil { + return err + } + } + return nil +} + +func (i BackendMigrationStatusInput) normalize() (BackendMigrationStatusInput, error) { + i.TenantID = strings.TrimSpace(i.TenantID) + if i.TenantID == "" { + return BackendMigrationStatusInput{}, ErrTenantIDRequired + } + i.AppID = strings.TrimSpace(i.AppID) + i.ProfileID = strings.TrimSpace(i.ProfileID) + if i.ProfileID == "" { + return BackendMigrationStatusInput{}, fmt.Errorf("profile_id is required") + } + i.Resource = BackendMigrationResource(strings.TrimSpace(string(i.Resource))) + if !i.Resource.valid() { + return BackendMigrationStatusInput{}, fmt.Errorf("invalid backend migration resource %q", i.Resource) + } + i.SourceBackendID = strings.TrimSpace(i.SourceBackendID) + if i.SourceBackendID == "" { + return BackendMigrationStatusInput{}, fmt.Errorf("source_backend_id is required") + } + i.TargetBackendID = strings.TrimSpace(i.TargetBackendID) + if i.TargetBackendID == "" { + return BackendMigrationStatusInput{}, fmt.Errorf("target_backend_id is required") + } + if i.SourceBackendID == i.TargetBackendID { + return BackendMigrationStatusInput{}, fmt.Errorf("source_backend_id and target_backend_id must differ") + } + mode, err := NormalizeStorageMigrationMode(string(i.MigrationMode)) + if err != nil { + return BackendMigrationStatusInput{}, err + } + if !IsActiveStorageMigrationMode(mode) { + return BackendMigrationStatusInput{}, fmt.Errorf("migration_mode must be an active migration mode") + } + i.MigrationMode = mode + i.Status = BackendMigrationStatus(strings.TrimSpace(string(i.Status))) + if !i.Status.valid() { + return BackendMigrationStatusInput{}, fmt.Errorf("invalid backend migration status %q", i.Status) + } + i.OperationID = strings.TrimSpace(i.OperationID) + if i.OperationID == "" { + return BackendMigrationStatusInput{}, fmt.Errorf("operation_id is required") + } + i.LastRecordID = strings.TrimSpace(i.LastRecordID) + i.SampleSetRef = strings.TrimSpace(i.SampleSetRef) + i.FailureReason = strings.TrimSpace(i.FailureReason) + i.TraceID = strings.TrimSpace(i.TraceID) + report := BackendMigrationStatusReport{ + SourceRecordCount: i.SourceRecordCount, + TargetRecordCount: i.TargetRecordCount, + LagRecordCount: backendMigrationLag(i.SourceRecordCount, i.TargetRecordCount), + VerifiedRecordCount: i.VerifiedRecordCount, + MismatchCount: i.MismatchCount, + SampledTopKQueries: i.SampledTopKQueries, + MatchedTopKQueries: i.MatchedTopKQueries, + } + if err := validateBackendMigrationCounts(report); err != nil { + return BackendMigrationStatusInput{}, err + } + report.MigrationMode = i.MigrationMode + report.Status = i.Status + if err := validateBackendMigrationStatusGate(report); err != nil { + return BackendMigrationStatusInput{}, err + } + if i.Status == BackendMigrationStatusFailed && i.FailureReason == "" { + return BackendMigrationStatusInput{}, fmt.Errorf("failure_reason is required for failed backend migration status") + } + if i.UpdatedAt.IsZero() { + return BackendMigrationStatusInput{}, fmt.Errorf("updated_at is required") + } + for field, value := range map[string]string{ + "app_id": i.AppID, + "profile_id": i.ProfileID, + "source_backend_id": i.SourceBackendID, + "target_backend_id": i.TargetBackendID, + "operation_id": i.OperationID, + "last_record_id": i.LastRecordID, + "sample_set_ref": i.SampleSetRef, + "failure_reason": i.FailureReason, + "trace_id": i.TraceID, + } { + if err := validateAuditRedactedText(field, value); err != nil { + return BackendMigrationStatusInput{}, err + } + } + return i, nil +} + +func (i BackendMigrationStatusInput) migrationID() string { + return backendMigrationID( + i.TenantID, + i.AppID, + i.ProfileID, + i.Resource, + i.SourceBackendID, + i.TargetBackendID, + i.OperationID, + ) +} + +func backendMigrationID( + tenantID string, + appID string, + profileID string, + resource BackendMigrationResource, + sourceBackendID string, + targetBackendID string, + operationID string, +) string { + return backendMigrationIDPrefix + shortHash( + strings.TrimSpace(tenantID), + strings.TrimSpace(appID), + strings.TrimSpace(profileID), + string(resource), + strings.TrimSpace(sourceBackendID), + strings.TrimSpace(targetBackendID), + strings.TrimSpace(operationID), + ) +} + +func validateBackendMigrationCounts(r BackendMigrationStatusReport) error { + for field, value := range map[string]int64{ + "source_record_count": r.SourceRecordCount, + "target_record_count": r.TargetRecordCount, + "lag_record_count": r.LagRecordCount, + "verified_record_count": r.VerifiedRecordCount, + "mismatch_count": r.MismatchCount, + "sampled_topk_queries": r.SampledTopKQueries, + "matched_topk_queries": r.MatchedTopKQueries, + } { + if value < 0 { + return fmt.Errorf("%s must be non-negative", field) + } + } + if r.MismatchCount > r.VerifiedRecordCount { + return fmt.Errorf("mismatch_count must be less than or equal to verified_record_count") + } + if r.MatchedTopKQueries > r.SampledTopKQueries { + return fmt.Errorf("matched_topk_queries must be less than or equal to sampled_topk_queries") + } + if expected := backendMigrationLag(r.SourceRecordCount, r.TargetRecordCount); r.LagRecordCount != expected { + return fmt.Errorf("lag_record_count must equal source_record_count minus target_record_count when positive") + } + return nil +} + +func validateBackendMigrationStatusGate(r BackendMigrationStatusReport) error { + switch r.Status { + case BackendMigrationStatusReady, BackendMigrationStatusCompleted: + if r.SourceRecordCount != r.TargetRecordCount { + return fmt.Errorf("%s backend migration status requires source_record_count to equal target_record_count", r.Status) + } + if r.VerifiedRecordCount != r.SourceRecordCount { + return fmt.Errorf("%s backend migration status requires verified_record_count to equal source_record_count", r.Status) + } + if r.LagRecordCount != 0 { + return fmt.Errorf("%s backend migration status requires zero lag_record_count", r.Status) + } + if r.MismatchCount != 0 { + return fmt.Errorf("%s backend migration status requires zero mismatch_count", r.Status) + } + if r.Resource == BackendMigrationResourceKnowledge && r.SampledTopKQueries == 0 { + return fmt.Errorf("%s knowledge backend migration status requires sampled_topk_queries", r.Status) + } + if r.Resource == BackendMigrationResourceKnowledge && strings.TrimSpace(r.SampleSetRef) == "" { + return fmt.Errorf("%s knowledge backend migration status requires sample_set_ref", r.Status) + } + if r.MatchedTopKQueries != r.SampledTopKQueries { + return fmt.Errorf("%s backend migration status requires all sampled topK queries to match", r.Status) + } + case BackendMigrationStatusRolledBack: + if r.MigrationMode != StorageMigrationModeRollback { + return fmt.Errorf("rolled_back backend migration status requires rollback migration_mode") + } + } + if r.Status == BackendMigrationStatusCompleted && + r.MigrationMode != StorageMigrationModeCutover { + return fmt.Errorf("completed backend migration status requires cutover migration_mode") + } + return nil +} + +func backendMigrationLag(sourceCount, targetCount int64) int64 { + if sourceCount <= targetCount { + return 0 + } + return sourceCount - targetCount +} + +func (r BackendMigrationResource) valid() bool { + switch r { + case BackendMigrationResourceSession, + BackendMigrationResourceMemory, + BackendMigrationResourceArtifact, + BackendMigrationResourceKnowledge, + BackendMigrationResourceAudit: + return true + default: + return false + } +} + +func (s BackendMigrationStatus) valid() bool { + switch s { + case BackendMigrationStatusPending, + BackendMigrationStatusRunning, + BackendMigrationStatusVerifying, + BackendMigrationStatusReady, + BackendMigrationStatusCompleted, + BackendMigrationStatusRolledBack, + BackendMigrationStatusFailed: + return true + default: + return false + } +} + +func isBackendMigrationID(value string) bool { + if !strings.HasPrefix(value, backendMigrationIDPrefix) { + return false + } + return isShortHash(strings.TrimPrefix(value, backendMigrationIDPrefix)) +} diff --git a/platform/backend_migration_status_test.go b/platform/backend_migration_status_test.go new file mode 100644 index 0000000000..bbfc5f15db --- /dev/null +++ b/platform/backend_migration_status_test.go @@ -0,0 +1,382 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "errors" + "fmt" + "strings" + "testing" + "time" +) + +func TestNewBackendMigrationStatusReportBuildsSafeReport(t *testing.T) { + updatedAt := time.Date(2026, 7, 8, 14, 0, 0, 0, time.UTC) + input := BackendMigrationStatusInput{ + TenantID: "tenant", + AppID: "app", + ProfileID: "profile-a", + Resource: BackendMigrationResourceSession, + SourceBackendID: "redis-primary", + TargetBackendID: "sql-primary", + MigrationMode: StorageMigrationModeDualWrite, + Status: BackendMigrationStatusVerifying, + OperationID: "migration-1", + SourceRecordCount: 100, + TargetRecordCount: 97, + VerifiedRecordCount: 97, + MismatchCount: 1, + LastRecordID: "event-100", + SampleSetRef: "sample://tenant/app/session-migration-1", + SampledTopKQueries: 10, + MatchedTopKQueries: 9, + TraceID: "trace-1", + UpdatedAt: updatedAt, + } + + report, err := NewBackendMigrationStatusReport(input) + if err != nil { + t.Fatalf("new backend migration status report: %v", err) + } + if report.TenantID != "tenant" || report.AppID != "app" || report.ProfileID != "profile-a" { + t.Fatalf("unexpected owner/profile: %+v", report) + } + if report.Resource != BackendMigrationResourceSession || + report.SourceBackendID != "redis-primary" || + report.TargetBackendID != "sql-primary" || + report.MigrationMode != StorageMigrationModeDualWrite || + report.Status != BackendMigrationStatusVerifying { + t.Fatalf("unexpected routing fields: %+v", report) + } + if report.SourceRecordCount != 100 || + report.TargetRecordCount != 97 || + report.LagRecordCount != 3 || + report.VerifiedRecordCount != 97 || + report.MismatchCount != 1 { + t.Fatalf("unexpected count summary: %+v", report) + } + if report.SampledTopKQueries != 10 || report.MatchedTopKQueries != 9 { + t.Fatalf("unexpected topK summary: %+v", report) + } + if report.OperationID != "migration-1" || report.TraceID != "trace-1" || + !report.UpdatedAt.Equal(updatedAt) { + t.Fatalf("unexpected operation metadata: %+v", report) + } + if !strings.HasPrefix(report.MigrationID, backendMigrationIDPrefix) { + t.Fatalf("unexpected migration id: %q", report.MigrationID) + } + + again, err := NewBackendMigrationStatusReport(input) + if err != nil { + t.Fatalf("new duplicate backend migration status report: %v", err) + } + if report.MigrationID != again.MigrationID { + t.Fatalf("expected stable migration id, got %q and %q", report.MigrationID, again.MigrationID) + } + + nextOperation := input + nextOperation.OperationID = "migration-2" + nextReport, err := NewBackendMigrationStatusReport(nextOperation) + if err != nil { + t.Fatalf("new next backend migration status report: %v", err) + } + if report.MigrationID == nextReport.MigrationID { + t.Fatalf("expected operation id to scope migration id, got %q", report.MigrationID) + } + + serialized := fmt.Sprintf("%+v", report) + if strings.Contains(serialized, "password=plain") { + t.Fatalf("report leaked sensitive content: %s", serialized) + } +} + +func TestNewBackendMigrationStatusReportRejectsInvalidInputs(t *testing.T) { + base := validBackendMigrationStatusInput() + + missingTenant := base + missingTenant.TenantID = " " + if _, err := NewBackendMigrationStatusReport(missingTenant); !errors.Is(err, ErrTenantIDRequired) { + t.Fatalf("expected tenant requirement, got %v", err) + } + + missingProfile := base + missingProfile.ProfileID = " " + if _, err := NewBackendMigrationStatusReport(missingProfile); err == nil || + !strings.Contains(err.Error(), "profile_id") { + t.Fatalf("expected profile id requirement, got %v", err) + } + + unknownResource := base + unknownResource.Resource = "cache" + if _, err := NewBackendMigrationStatusReport(unknownResource); err == nil || + !strings.Contains(err.Error(), "invalid backend migration resource") { + t.Fatalf("expected resource validation, got %v", err) + } + + sameBackend := base + sameBackend.TargetBackendID = sameBackend.SourceBackendID + if _, err := NewBackendMigrationStatusReport(sameBackend); err == nil || + !strings.Contains(err.Error(), "must differ") { + t.Fatalf("expected source/target mismatch requirement, got %v", err) + } + + normalMode := base + normalMode.MigrationMode = StorageMigrationModeNormal + if _, err := NewBackendMigrationStatusReport(normalMode); err == nil || + !strings.Contains(err.Error(), "active migration mode") { + t.Fatalf("expected active migration mode requirement, got %v", err) + } + + unknownStatus := base + unknownStatus.Status = "paused" + if _, err := NewBackendMigrationStatusReport(unknownStatus); err == nil || + !strings.Contains(err.Error(), "invalid backend migration status") { + t.Fatalf("expected status validation, got %v", err) + } + + missingOperation := base + missingOperation.OperationID = " " + if _, err := NewBackendMigrationStatusReport(missingOperation); err == nil || + !strings.Contains(err.Error(), "operation_id") { + t.Fatalf("expected operation id requirement, got %v", err) + } + + negativeCount := base + negativeCount.TargetRecordCount = -1 + if _, err := NewBackendMigrationStatusReport(negativeCount); err == nil || + !strings.Contains(err.Error(), "target_record_count") { + t.Fatalf("expected non-negative count validation, got %v", err) + } + + tooManyMismatches := base + tooManyMismatches.MismatchCount = tooManyMismatches.VerifiedRecordCount + 1 + if _, err := NewBackendMigrationStatusReport(tooManyMismatches); err == nil || + !strings.Contains(err.Error(), "mismatch_count") { + t.Fatalf("expected mismatch bound validation, got %v", err) + } + + tooManyTopKMatches := base + tooManyTopKMatches.MatchedTopKQueries = tooManyTopKMatches.SampledTopKQueries + 1 + if _, err := NewBackendMigrationStatusReport(tooManyTopKMatches); err == nil || + !strings.Contains(err.Error(), "matched_topk_queries") { + t.Fatalf("expected topK bound validation, got %v", err) + } + + failedWithoutReason := base + failedWithoutReason.Status = BackendMigrationStatusFailed + failedWithoutReason.FailureReason = " " + if _, err := NewBackendMigrationStatusReport(failedWithoutReason); err == nil || + !strings.Contains(err.Error(), "failure_reason") { + t.Fatalf("expected failed status reason requirement, got %v", err) + } + + sensitiveReason := base + sensitiveReason.Status = BackendMigrationStatusFailed + sensitiveReason.FailureReason = "password=plain" + if _, err := NewBackendMigrationStatusReport(sensitiveReason); err == nil || + !strings.Contains(err.Error(), "failure_reason") { + t.Fatalf("expected sensitive failure reason rejection, got %v", err) + } + + zeroUpdatedAt := base + zeroUpdatedAt.UpdatedAt = time.Time{} + if _, err := NewBackendMigrationStatusReport(zeroUpdatedAt); err == nil || + !strings.Contains(err.Error(), "updated_at") { + t.Fatalf("expected updated at requirement, got %v", err) + } +} + +func TestBackendMigrationStatusReportValidateRejectsUnsafeReport(t *testing.T) { + generated, err := NewBackendMigrationStatusReport(validBackendMigrationStatusInput()) + if err != nil { + t.Fatalf("new generated report: %v", err) + } + report := generated + if err := report.Validate(); err != nil { + t.Fatalf("expected report to validate: %v", err) + } + + report.MigrationID = "backend_migration_profile-a" + if err := report.Validate(); err == nil || + !strings.Contains(err.Error(), "migration_id") { + t.Fatalf("expected unsafe migration id rejection, got %v", err) + } + + report = generated + report.OperationID = "other-migration" + if err := report.Validate(); err == nil || + !strings.Contains(err.Error(), "migration_id") { + t.Fatalf("expected stale migration id rejection, got %v", err) + } + + report = generated + report.LagRecordCount = -1 + if err := report.Validate(); err == nil || + !strings.Contains(err.Error(), "lag_record_count") { + t.Fatalf("expected unsafe lag count rejection, got %v", err) + } + + report = generated + report.LagRecordCount = 0 + if err := report.Validate(); err == nil || + !strings.Contains(err.Error(), "lag_record_count") { + t.Fatalf("expected inconsistent lag count rejection, got %v", err) + } + + report = generated + report.FailureReason = "token: plain" + if err := report.Validate(); err == nil || + !strings.Contains(err.Error(), "failure_reason") { + t.Fatalf("expected sensitive failure reason rejection, got %v", err) + } +} + +func TestNewBackendMigrationStatusReportEnforcesStatusGates(t *testing.T) { + readyWithLag := validBackendMigrationStatusInput() + readyWithLag.Status = BackendMigrationStatusReady + if _, err := NewBackendMigrationStatusReport(readyWithLag); err == nil || + !strings.Contains(err.Error(), "source_record_count") { + t.Fatalf("expected ready status count equality rejection, got %v", err) + } + + readyWithMismatch := validBackendMigrationStatusInput() + readyWithMismatch.Status = BackendMigrationStatusReady + readyWithMismatch.TargetRecordCount = readyWithMismatch.SourceRecordCount + readyWithMismatch.VerifiedRecordCount = readyWithMismatch.SourceRecordCount + readyWithMismatch.MismatchCount = 1 + if _, err := NewBackendMigrationStatusReport(readyWithMismatch); err == nil || + !strings.Contains(err.Error(), "mismatch_count") { + t.Fatalf("expected ready status mismatch rejection, got %v", err) + } + + readyWithMissingVerification := validBackendMigrationStatusInput() + readyWithMissingVerification.Status = BackendMigrationStatusReady + readyWithMissingVerification.TargetRecordCount = readyWithMissingVerification.SourceRecordCount + readyWithMissingVerification.VerifiedRecordCount = 0 + if _, err := NewBackendMigrationStatusReport(readyWithMissingVerification); err == nil || + !strings.Contains(err.Error(), "verified_record_count") { + t.Fatalf("expected ready status verification count rejection, got %v", err) + } + + readyWithTopKGap := validBackendMigrationStatusInput() + readyWithTopKGap.Status = BackendMigrationStatusReady + readyWithTopKGap.TargetRecordCount = readyWithTopKGap.SourceRecordCount + readyWithTopKGap.VerifiedRecordCount = readyWithTopKGap.SourceRecordCount + readyWithTopKGap.MatchedTopKQueries = readyWithTopKGap.SampledTopKQueries - 1 + if _, err := NewBackendMigrationStatusReport(readyWithTopKGap); err == nil || + !strings.Contains(err.Error(), "topK") { + t.Fatalf("expected ready status topK rejection, got %v", err) + } + + completedWrongMode := validBackendMigrationStatusInput() + completedWrongMode.Status = BackendMigrationStatusCompleted + completedWrongMode.TargetRecordCount = completedWrongMode.SourceRecordCount + completedWrongMode.VerifiedRecordCount = completedWrongMode.SourceRecordCount + if _, err := NewBackendMigrationStatusReport(completedWrongMode); err == nil || + !strings.Contains(err.Error(), "cutover") { + t.Fatalf("expected completed status cutover mode requirement, got %v", err) + } + + completed := validBackendMigrationStatusInput() + completed.Status = BackendMigrationStatusCompleted + completed.MigrationMode = StorageMigrationModeCutover + completed.TargetRecordCount = completed.SourceRecordCount + completed.VerifiedRecordCount = completed.SourceRecordCount + if _, err := NewBackendMigrationStatusReport(completed); err != nil { + t.Fatalf("expected completed cutover status to validate: %v", err) + } + + knowledgeWithoutTopKSamples := validBackendMigrationStatusInput() + knowledgeWithoutTopKSamples.Resource = BackendMigrationResourceKnowledge + knowledgeWithoutTopKSamples.Status = BackendMigrationStatusReady + knowledgeWithoutTopKSamples.TargetRecordCount = knowledgeWithoutTopKSamples.SourceRecordCount + knowledgeWithoutTopKSamples.VerifiedRecordCount = knowledgeWithoutTopKSamples.SourceRecordCount + knowledgeWithoutTopKSamples.SampledTopKQueries = 0 + knowledgeWithoutTopKSamples.MatchedTopKQueries = 0 + if _, err := NewBackendMigrationStatusReport(knowledgeWithoutTopKSamples); err == nil || + !strings.Contains(err.Error(), "sampled_topk_queries") { + t.Fatalf("expected knowledge topK sample requirement, got %v", err) + } + + knowledgeWithoutSampleRef := validBackendMigrationStatusInput() + knowledgeWithoutSampleRef.Resource = BackendMigrationResourceKnowledge + knowledgeWithoutSampleRef.Status = BackendMigrationStatusReady + knowledgeWithoutSampleRef.TargetRecordCount = knowledgeWithoutSampleRef.SourceRecordCount + knowledgeWithoutSampleRef.VerifiedRecordCount = knowledgeWithoutSampleRef.SourceRecordCount + knowledgeWithoutSampleRef.SampleSetRef = " " + if _, err := NewBackendMigrationStatusReport(knowledgeWithoutSampleRef); err == nil || + !strings.Contains(err.Error(), "sample_set_ref") { + t.Fatalf("expected knowledge sample ref requirement, got %v", err) + } + + knowledgeReady := validBackendMigrationStatusInput() + knowledgeReady.Resource = BackendMigrationResourceKnowledge + knowledgeReady.Status = BackendMigrationStatusReady + knowledgeReady.TargetRecordCount = knowledgeReady.SourceRecordCount + knowledgeReady.VerifiedRecordCount = knowledgeReady.SourceRecordCount + if _, err := NewBackendMigrationStatusReport(knowledgeReady); err != nil { + t.Fatalf("expected knowledge ready status to validate: %v", err) + } + + rolledBackWrongMode := validBackendMigrationStatusInput() + rolledBackWrongMode.Status = BackendMigrationStatusRolledBack + if _, err := NewBackendMigrationStatusReport(rolledBackWrongMode); err == nil || + !strings.Contains(err.Error(), "rollback") { + t.Fatalf("expected rolled_back status rollback mode requirement, got %v", err) + } + + rolledBack := validBackendMigrationStatusInput() + rolledBack.Status = BackendMigrationStatusRolledBack + rolledBack.MigrationMode = StorageMigrationModeRollback + if _, err := NewBackendMigrationStatusReport(rolledBack); err != nil { + t.Fatalf("expected rolled_back rollback status to validate: %v", err) + } +} + +func TestNewBackendMigrationStatusReportSupportsAcceptanceResources(t *testing.T) { + for _, resource := range []BackendMigrationResource{ + BackendMigrationResourceSession, + BackendMigrationResourceMemory, + BackendMigrationResourceArtifact, + BackendMigrationResourceKnowledge, + BackendMigrationResourceAudit, + } { + t.Run(string(resource), func(t *testing.T) { + input := validBackendMigrationStatusInput() + input.Resource = resource + if _, err := NewBackendMigrationStatusReport(input); err != nil { + t.Fatalf("expected resource %q to validate: %v", resource, err) + } + }) + } +} + +func validBackendMigrationStatusInput() BackendMigrationStatusInput { + return BackendMigrationStatusInput{ + TenantID: "tenant", + AppID: "app", + ProfileID: "profile", + Resource: BackendMigrationResourceSession, + SourceBackendID: "redis", + TargetBackendID: "sql", + MigrationMode: StorageMigrationModeShadowRead, + Status: BackendMigrationStatusRunning, + OperationID: "migration", + SourceRecordCount: 20, + TargetRecordCount: 18, + VerifiedRecordCount: 18, + MismatchCount: 0, + LastRecordID: "event-20", + SampleSetRef: "sample://tenant/app/migration", + SampledTopKQueries: 5, + MatchedTopKQueries: 5, + TraceID: "trace", + UpdatedAt: time.Now(), + } +} From 274747cec0039db6fa30726a9ed3e3603c08f2d3 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 16:29:51 +0800 Subject: [PATCH 30/95] platform: add storage router status summary --- platform/storagerouter/status.go | 181 ++++++++++++++++++++++++++ platform/storagerouter/status_test.go | 163 +++++++++++++++++++++++ 2 files changed, 344 insertions(+) create mode 100644 platform/storagerouter/status.go create mode 100644 platform/storagerouter/status_test.go diff --git a/platform/storagerouter/status.go b/platform/storagerouter/status.go new file mode 100644 index 0000000000..0964e1ee2d --- /dev/null +++ b/platform/storagerouter/status.go @@ -0,0 +1,181 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package storagerouter + +import ( + "context" + "fmt" + "strings" + "unicode" + + "trpc.group/trpc-go/trpc-agent-go/platform" +) + +// ResourceStatus describes the routing readiness of one storage resource. +type ResourceStatus string + +const ( + // ResourceStatusReady means the selected backend is registered and has the resource service. + ResourceStatusReady ResourceStatus = "ready" + // ResourceStatusBackendMissing means the profile points at an unregistered backend. + ResourceStatusBackendMissing ResourceStatus = "backend_missing" + // ResourceStatusServiceMissing means the backend is registered without the requested service. + ResourceStatusServiceMissing ResourceStatus = "service_missing" + // ResourceStatusBackendTenantMismatch means the selected backend belongs to another tenant. + ResourceStatusBackendTenantMismatch ResourceStatus = "backend_tenant_mismatch" +) + +// ResourceStatusEntry is a safe operations-facing status for one routed resource. +type ResourceStatusEntry struct { + Resource platform.BackendMigrationResource + BackendID string + Status ResourceStatus + Reason string +} + +// StatusSummary is a safe operations-facing view of one storage profile route. +type StatusSummary struct { + TenantID string + ProfileID string + MigrationMode platform.StorageMigrationMode + IsMigrating bool + Resources []ResourceStatusEntry + ReadyCount int + MissingCount int +} + +// Status summarizes the backend readiness for one registered storage profile. +func (r *InMemoryRouter) Status( + ctx context.Context, + tenantID string, + profileID string, +) (StatusSummary, error) { + profile, err := r.Profile(ctx, tenantID, profileID) + if err != nil { + return StatusSummary{}, err + } + mode, err := platform.NormalizeStorageMigrationMode(profile.MigrationMode) + if err != nil { + return StatusSummary{}, err + } + summary := StatusSummary{ + TenantID: profile.TenantID, + ProfileID: profile.ProfileID, + MigrationMode: mode, + IsMigrating: platform.IsActiveStorageMigrationMode(mode), + } + + for _, resource := range []struct { + kind resourceKind + resource platform.BackendMigrationResource + }{ + {kind: resourceSession, resource: platform.BackendMigrationResourceSession}, + {kind: resourceMemory, resource: platform.BackendMigrationResourceMemory}, + {kind: resourceArtifact, resource: platform.BackendMigrationResourceArtifact}, + {kind: resourceKnowledge, resource: platform.BackendMigrationResourceKnowledge}, + {kind: resourceAudit, resource: platform.BackendMigrationResourceAudit}, + } { + if err := ctx.Err(); err != nil { + return StatusSummary{}, err + } + entry := r.resourceStatus(ctx, profile, resource.kind, resource.resource) + summary.Resources = append(summary.Resources, entry) + if entry.Status == ResourceStatusReady { + summary.ReadyCount++ + } else { + summary.MissingCount++ + } + } + return summary, nil +} + +func (r *InMemoryRouter) resourceStatus( + ctx context.Context, + profile platform.StorageProfile, + kind resourceKind, + resource platform.BackendMigrationResource, +) ResourceStatusEntry { + entry := ResourceStatusEntry{ + Resource: resource, + } + backendID := backendIDFor(profile, kind) + entry.BackendID = safeBackendIDForStatus(backendID) + if strings.TrimSpace(backendID) == "" { + entry.Status = ResourceStatusBackendMissing + entry.Reason = fmt.Sprintf("%s backend is not configured", resource) + return entry + } + if entry.BackendID == "" { + entry.Status = ResourceStatusBackendMissing + entry.Reason = fmt.Sprintf("%s backend id is unsafe to expose", resource) + return entry + } + + r.mu.RLock() + backend, ok := r.backends[backendKey{tenantID: profile.TenantID, backendID: backendID}] + r.mu.RUnlock() + if !ok { + entry.Status = ResourceStatusBackendMissing + entry.Reason = fmt.Sprintf("%s backend is not registered", resource) + return entry + } + if backend.TenantID != profile.TenantID { + entry.Status = ResourceStatusBackendTenantMismatch + entry.Reason = fmt.Sprintf("%s backend belongs to another tenant", resource) + return entry + } + if !backendHasResource(backend, kind) { + entry.Status = ResourceStatusServiceMissing + entry.Reason = fmt.Sprintf("%s service is not registered on selected backend", resource) + return entry + } + entry.Status = ResourceStatusReady + return entry +} + +func backendHasResource(backend BackendSet, kind resourceKind) bool { + switch kind { + case resourceSession: + return backend.Session != nil + case resourceMemory: + return backend.Memory != nil + case resourceArtifact: + return backend.Artifact != nil + case resourceKnowledge: + return backend.Knowledge != nil + case resourceAudit: + return backend.Audit != nil + default: + return false + } +} + +func safeBackendIDForStatus(backendID string) string { + backendID = strings.TrimSpace(backendID) + if backendID == "" || strings.Contains(backendID, "://") || + strings.ContainsAny(backendID, "=@/\\") { + return "" + } + redactor, err := platform.NewRedactor() + if err != nil || redactor.Redact(backendID) != backendID { + return "" + } + for _, r := range backendID { + if unicode.IsLetter(r) || unicode.IsDigit(r) { + continue + } + switch r { + case '-', '_', '.': + continue + default: + return "" + } + } + return backendID +} diff --git a/platform/storagerouter/status_test.go b/platform/storagerouter/status_test.go new file mode 100644 index 0000000000..0c706e6673 --- /dev/null +++ b/platform/storagerouter/status_test.go @@ -0,0 +1,163 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package storagerouter + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + artifactmemory "trpc.group/trpc-go/trpc-agent-go/artifact/inmemory" + memoryinmemory "trpc.group/trpc-go/trpc-agent-go/memory/inmemory" + "trpc.group/trpc-go/trpc-agent-go/platform" + sessioninmemory "trpc.group/trpc-go/trpc-agent-go/session/inmemory" +) + +func TestRouterStatusReportsAllResourcesReady(t *testing.T) { + ctx := context.Background() + router := NewInMemoryRouter() + p := profile("tenant-a", "profile-a", "hot") + p.MigrationMode = string(platform.StorageMigrationModeDualWrite) + require.NoError(t, router.RegisterProfile(p)) + require.NoError(t, router.RegisterBackend(BackendSet{ + TenantID: "tenant-a", + BackendID: "hot", + Session: sessioninmemory.NewSessionService(), + Memory: memoryinmemory.NewMemoryService(), + Artifact: artifactmemory.NewService(), + Knowledge: &stubKnowledge{}, + Audit: platform.NewInMemoryAuditSink(), + })) + + summary, err := router.Status(ctx, "tenant-a", "profile-a") + require.NoError(t, err) + + assert.Equal(t, "tenant-a", summary.TenantID) + assert.Equal(t, "profile-a", summary.ProfileID) + assert.Equal(t, platform.StorageMigrationModeDualWrite, summary.MigrationMode) + assert.True(t, summary.IsMigrating) + assert.Equal(t, 5, summary.ReadyCount) + assert.Equal(t, 0, summary.MissingCount) + require.Len(t, summary.Resources, 5) + for _, resource := range summary.Resources { + assert.Equal(t, "hot", resource.BackendID) + assert.Equal(t, ResourceStatusReady, resource.Status) + assert.Empty(t, resource.Reason) + } +} + +func TestRouterStatusReportsMissingBackendAndService(t *testing.T) { + ctx := context.Background() + router := NewInMemoryRouter() + p := profile("tenant-a", "profile-a", "hot") + p.MemoryBackend = "missing" + p.ArtifactBackend = "" + require.NoError(t, router.RegisterProfile(p)) + require.NoError(t, router.RegisterBackend(BackendSet{ + TenantID: "tenant-a", + BackendID: "hot", + Session: sessioninmemory.NewSessionService(), + Audit: platform.NewInMemoryAuditSink(), + })) + + summary, err := router.Status(ctx, "tenant-a", "profile-a") + require.NoError(t, err) + + assert.False(t, summary.IsMigrating) + assert.Equal(t, 2, summary.ReadyCount) + assert.Equal(t, 3, summary.MissingCount) + assertResourceStatus(t, summary, platform.BackendMigrationResourceSession, "hot", ResourceStatusReady) + assertResourceStatus(t, summary, platform.BackendMigrationResourceMemory, "missing", ResourceStatusBackendMissing) + assertResourceStatus(t, summary, platform.BackendMigrationResourceArtifact, "", ResourceStatusBackendMissing) + assertResourceStatus(t, summary, platform.BackendMigrationResourceKnowledge, "hot", ResourceStatusServiceMissing) + assertResourceStatus(t, summary, platform.BackendMigrationResourceAudit, "hot", ResourceStatusReady) +} + +func TestRouterStatusRedactsUnsafeBackendIDs(t *testing.T) { + ctx := context.Background() + router := NewInMemoryRouter() + p := profile("tenant-a", "profile-a", "hot") + p.SessionBackend = "postgres://user:password@localhost/db" + p.MemoryBackend = "sk-1234567890abcdef" + p.ArtifactBackend = "safe.backend-1" + require.NoError(t, router.RegisterProfile(p)) + require.NoError(t, router.RegisterBackend(BackendSet{ + TenantID: "tenant-a", + BackendID: "safe.backend-1", + Artifact: artifactmemory.NewService(), + })) + + summary, err := router.Status(ctx, "tenant-a", "profile-a") + require.NoError(t, err) + + assertResourceStatus(t, summary, platform.BackendMigrationResourceSession, "", ResourceStatusBackendMissing) + assertResourceStatus(t, summary, platform.BackendMigrationResourceMemory, "", ResourceStatusBackendMissing) + assertResourceStatus(t, summary, platform.BackendMigrationResourceArtifact, "safe.backend-1", ResourceStatusReady) + for _, entry := range summary.Resources { + assert.NotContains(t, entry.BackendID, "password") + assert.NotContains(t, entry.Reason, "password") + assert.NotContains(t, entry.BackendID, "sk-") + assert.NotContains(t, entry.Reason, "sk-") + } +} + +func TestRouterStatusHonorsContextCancellation(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + router := NewInMemoryRouter() + + _, err := router.Status(ctx, "tenant-a", "profile-a") + + require.True(t, errors.Is(err, context.Canceled)) +} + +func TestRouterStatusReturnsContextCancellationAfterProfileLookup(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + router := NewInMemoryRouter() + require.NoError(t, router.RegisterProfile(profile("tenant-a", "profile-a", "hot"))) + cancel() + + _, err := router.Status(ctx, "tenant-a", "profile-a") + + require.True(t, errors.Is(err, context.Canceled)) +} + +func TestRouterStatusRejectsUnknownProfile(t *testing.T) { + router := NewInMemoryRouter() + + _, err := router.Status(context.Background(), "tenant-a", "profile-a") + + require.ErrorIs(t, err, ErrProfileNotFound) +} + +func assertResourceStatus( + t *testing.T, + summary StatusSummary, + resource platform.BackendMigrationResource, + backendID string, + status ResourceStatus, +) { + t.Helper() + for _, entry := range summary.Resources { + if entry.Resource == resource { + assert.Equal(t, backendID, entry.BackendID) + assert.Equal(t, status, entry.Status) + if status == ResourceStatusReady { + assert.Empty(t, entry.Reason) + } else { + assert.NotEmpty(t, entry.Reason) + } + return + } + } + t.Fatalf("missing resource status for %q", resource) +} From 826a71849b10c7f8f9da31b7b12832e7062da4cb Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 16:45:49 +0800 Subject: [PATCH 31/95] platform: add config operation summary contracts --- platform/config_operation_summary.go | 378 ++++++++++++++++++++++ platform/config_operation_summary_test.go | 354 ++++++++++++++++++++ 2 files changed, 732 insertions(+) create mode 100644 platform/config_operation_summary.go create mode 100644 platform/config_operation_summary_test.go diff --git a/platform/config_operation_summary.go b/platform/config_operation_summary.go new file mode 100644 index 0000000000..489b99f82f --- /dev/null +++ b/platform/config_operation_summary.go @@ -0,0 +1,378 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "fmt" + "strings" + "time" +) + +const configOperationSummaryIDPrefix = "config_operation_" + +// AppConfigOperation names an operations-facing config switch. +type AppConfigOperation string + +const ( + // AppConfigOperationActivate promotes a released config version to active. + AppConfigOperationActivate AppConfigOperation = "activate" + // AppConfigOperationRollback promotes a rollback config version to active. + AppConfigOperationRollback AppConfigOperation = "rollback" +) + +// AppConfigOperationSummaryInput describes one planned or completed config operation. +type AppConfigOperationSummaryInput struct { + Operation AppConfigOperation + PreviousActive AppConfigVersion + NextActive AppConfigVersion + ResultVersions []AppConfigVersion + OperationID string + TraceID string + CreatedAt time.Time +} + +// AppConfigOperationSummary is a safe operations-facing summary of one config switch. +type AppConfigOperationSummary struct { + TenantID string + AppID string + SummaryID string + Operation AppConfigOperation + OperationID string + PreviousVersion string + PreviousChecksum string + NextVersion string + NextChecksum string + DiffChangeCount int + CacheInvalidation AppConfigCacheInvalidation + GrayStatus ConfigGrayStatusSummary + RequiresCacheFlush bool + TraceID string + CreatedAt time.Time +} + +// NewAppConfigOperationSummary builds a safe config operation summary from existing contracts. +func NewAppConfigOperationSummary(input AppConfigOperationSummaryInput) (AppConfigOperationSummary, error) { + normalized, err := input.normalize() + if err != nil { + return AppConfigOperationSummary{}, err + } + diff, err := DiffAppConfigVersions(normalized.PreviousActive, normalized.NextActive) + if err != nil { + return AppConfigOperationSummary{}, err + } + invalidation, err := NewAppConfigCacheInvalidation(AppConfigCacheInvalidationInput{ + PreviousVersion: normalized.PreviousActive, + NextVersion: normalized.NextActive, + Reason: normalized.invalidationReason(), + OperationID: normalized.OperationID, + TraceID: normalized.TraceID, + CreatedAt: normalized.CreatedAt, + }) + if err != nil { + return AppConfigOperationSummary{}, err + } + grayStatus, err := SummarizeAppConfigGrayStatus(normalized.ResultVersions) + if err != nil { + return AppConfigOperationSummary{}, err + } + summary := AppConfigOperationSummary{ + TenantID: normalized.NextActive.TenantID, + AppID: normalized.NextActive.AppID, + SummaryID: normalized.summaryID(), + Operation: normalized.Operation, + OperationID: normalized.OperationID, + PreviousVersion: normalized.PreviousActive.Version, + PreviousChecksum: normalized.PreviousActive.Checksum, + NextVersion: normalized.NextActive.Version, + NextChecksum: normalized.NextActive.Checksum, + DiffChangeCount: len(diff.Changes), + CacheInvalidation: invalidation, + GrayStatus: grayStatus, + RequiresCacheFlush: true, + TraceID: normalized.TraceID, + CreatedAt: normalized.CreatedAt, + } + if err := summary.Validate(); err != nil { + return AppConfigOperationSummary{}, err + } + return summary, nil +} + +// Validate checks that a config operation summary is safe to expose or store. +func (s AppConfigOperationSummary) Validate() error { + if strings.TrimSpace(s.TenantID) == "" { + return ErrTenantIDRequired + } + if strings.TrimSpace(s.AppID) == "" { + return ErrAppIDRequired + } + if strings.TrimSpace(s.SummaryID) == "" { + return fmt.Errorf("summary_id is required") + } + if !isConfigOperationSummaryID(s.SummaryID) { + return fmt.Errorf("summary_id must be %s followed by a 24 character hex hash", configOperationSummaryIDPrefix) + } + if s.SummaryID != configOperationSummaryID( + s.TenantID, + s.AppID, + s.Operation, + s.PreviousVersion, + s.NextVersion, + s.OperationID, + ) { + return fmt.Errorf("summary_id does not match config operation identity") + } + if !s.Operation.valid() { + return fmt.Errorf("invalid config operation %q", s.Operation) + } + if strings.TrimSpace(s.OperationID) == "" { + return fmt.Errorf("operation_id is required") + } + if strings.TrimSpace(s.PreviousVersion) == "" { + return fmt.Errorf("previous_version is required") + } + if strings.TrimSpace(s.NextVersion) == "" { + return fmt.Errorf("next_version is required") + } + if strings.TrimSpace(s.PreviousVersion) == strings.TrimSpace(s.NextVersion) { + return fmt.Errorf("config operation must change active version") + } + if strings.TrimSpace(s.PreviousChecksum) == "" || + strings.TrimSpace(s.NextChecksum) == "" { + return fmt.Errorf("config checksums are required") + } + if s.DiffChangeCount <= 0 { + return fmt.Errorf("diff_change_count must be positive") + } + if !s.RequiresCacheFlush { + return fmt.Errorf("requires_cache_flush must be true") + } + if s.CreatedAt.IsZero() { + return fmt.Errorf("created_at is required") + } + if err := s.CacheInvalidation.Validate(); err != nil { + return fmt.Errorf("cache_invalidation: %w", err) + } + if err := validateConfigOperationInvalidation(s); err != nil { + return err + } + if err := validateConfigOperationOwner(s, s.CacheInvalidation.TenantID, s.CacheInvalidation.AppID); err != nil { + return err + } + if err := validateConfigOperationOwner(s, s.GrayStatus.TenantID, s.GrayStatus.AppID); err != nil { + return err + } + if err := validateConfigOperationGrayStatus(s); err != nil { + return err + } + if s.GrayStatus.ActiveVersion != s.NextVersion || + s.GrayStatus.ActiveChecksum != s.NextChecksum { + return fmt.Errorf("gray_status active version must match next active version") + } + for field, value := range map[string]string{ + "summary_id": s.SummaryID, + "operation_id": s.OperationID, + "previous_version": s.PreviousVersion, + "previous_checksum": s.PreviousChecksum, + "next_version": s.NextVersion, + "next_checksum": s.NextChecksum, + "trace_id": s.TraceID, + } { + if err := validateAuditRedactedText(field, value); err != nil { + return err + } + } + return nil +} + +func validateConfigOperationInvalidation(s AppConfigOperationSummary) error { + expectedReason := AppConfigCacheInvalidationReasonActivate + if s.Operation == AppConfigOperationRollback { + expectedReason = AppConfigCacheInvalidationReasonRollback + } + if s.CacheInvalidation.Reason != expectedReason { + return fmt.Errorf("cache_invalidation reason must match config operation") + } + if s.CacheInvalidation.OperationID != s.OperationID { + return fmt.Errorf("cache_invalidation operation_id must match config operation") + } + if s.CacheInvalidation.PreviousVersion != s.PreviousVersion || + s.CacheInvalidation.PreviousChecksum != s.PreviousChecksum || + s.CacheInvalidation.NextVersion != s.NextVersion || + s.CacheInvalidation.NextChecksum != s.NextChecksum { + return fmt.Errorf("cache_invalidation version summary must match config operation") + } + if s.CacheInvalidation.TraceID != s.TraceID { + return fmt.Errorf("cache_invalidation trace_id must match config operation") + } + if !s.CacheInvalidation.CreatedAt.Equal(s.CreatedAt) { + return fmt.Errorf("cache_invalidation created_at must match config operation") + } + return nil +} + +func validateConfigOperationGrayStatus(s AppConfigOperationSummary) error { + for field, value := range map[string]string{ + "gray_active_version": s.GrayStatus.ActiveVersion, + "gray_active_checksum": s.GrayStatus.ActiveChecksum, + "gray_candidate_version": s.GrayStatus.CandidateVersion, + "gray_candidate_checksum": s.GrayStatus.CandidateChecksum, + "gray_rollback_version": s.GrayStatus.RollbackVersion, + "gray_rollback_checksum": s.GrayStatus.RollbackChecksum, + } { + if err := validateAuditRedactedText(field, value); err != nil { + return err + } + } + if s.GrayStatus.ActiveTrafficPercent < 0 || s.GrayStatus.ActiveTrafficPercent > 100 || + s.GrayStatus.CandidateGrayPercent < 0 || s.GrayStatus.CandidateGrayPercent > 100 || + s.GrayStatus.CandidateTrafficPercent < 0 || s.GrayStatus.CandidateTrafficPercent > 100 { + return fmt.Errorf("gray_status traffic percentages must be between 0 and 100") + } + if s.GrayStatus.HasCandidate { + if strings.TrimSpace(s.GrayStatus.CandidateVersion) == "" || + strings.TrimSpace(s.GrayStatus.CandidateChecksum) == "" { + return fmt.Errorf("gray_status candidate version and checksum are required") + } + if s.GrayStatus.CandidateGrayPercent != s.GrayStatus.CandidateTrafficPercent { + return fmt.Errorf("gray_status candidate traffic must match candidate gray percent") + } + if s.GrayStatus.ActiveTrafficPercent != 100-s.GrayStatus.CandidateTrafficPercent { + return fmt.Errorf("gray_status active traffic must complement candidate traffic") + } + } else if s.GrayStatus.CandidateVersion != "" || + s.GrayStatus.CandidateChecksum != "" || + s.GrayStatus.CandidateGrayPercent != 0 || + s.GrayStatus.CandidateTrafficPercent != 0 { + return fmt.Errorf("gray_status candidate fields require has_candidate") + } else if s.GrayStatus.ActiveTrafficPercent != 100 { + return fmt.Errorf("gray_status active traffic must be 100 when there is no candidate") + } + if !s.GrayStatus.HasRollback { + return fmt.Errorf("gray_status rollback version is required") + } + if s.GrayStatus.RollbackVersion != s.PreviousVersion || + s.GrayStatus.RollbackChecksum != s.PreviousChecksum { + return fmt.Errorf("gray_status rollback version must match previous active version") + } + return nil +} + +func (i AppConfigOperationSummaryInput) normalize() (AppConfigOperationSummaryInput, error) { + i.Operation = AppConfigOperation(strings.TrimSpace(string(i.Operation))) + if !i.Operation.valid() { + return AppConfigOperationSummaryInput{}, fmt.Errorf("invalid config operation %q", i.Operation) + } + if err := i.PreviousActive.Validate(); err != nil { + return AppConfigOperationSummaryInput{}, fmt.Errorf("previous active config version: %w", err) + } + if i.PreviousActive.Status != AppConfigVersionStatusRollback { + return AppConfigOperationSummaryInput{}, fmt.Errorf("previous active config version status must be rollback") + } + if err := i.NextActive.Validate(); err != nil { + return AppConfigOperationSummaryInput{}, fmt.Errorf("next active config version: %w", err) + } + if i.NextActive.Status != AppConfigVersionStatusActive { + return AppConfigOperationSummaryInput{}, fmt.Errorf("next active config version status must be active") + } + if err := requireSameConfigOwner(i.PreviousActive, i.NextActive); err != nil { + return AppConfigOperationSummaryInput{}, err + } + if strings.TrimSpace(i.PreviousActive.Version) == strings.TrimSpace(i.NextActive.Version) { + return AppConfigOperationSummaryInput{}, fmt.Errorf("config operation must change active version") + } + i.OperationID = strings.TrimSpace(i.OperationID) + if i.OperationID == "" { + return AppConfigOperationSummaryInput{}, fmt.Errorf("operation_id is required") + } + i.TraceID = strings.TrimSpace(i.TraceID) + if i.CreatedAt.IsZero() { + return AppConfigOperationSummaryInput{}, fmt.Errorf("created_at is required") + } + if len(i.ResultVersions) == 0 { + return AppConfigOperationSummaryInput{}, fmt.Errorf("result_versions are required") + } + for _, version := range i.ResultVersions { + if err := requireSameConfigOwner(i.NextActive, version); err != nil { + return AppConfigOperationSummaryInput{}, err + } + } + for field, value := range map[string]string{ + "operation_id": i.OperationID, + "trace_id": i.TraceID, + } { + if err := validateAuditRedactedText(field, value); err != nil { + return AppConfigOperationSummaryInput{}, err + } + } + return i, nil +} + +func (i AppConfigOperationSummaryInput) invalidationReason() AppConfigCacheInvalidationReason { + if i.Operation == AppConfigOperationRollback { + return AppConfigCacheInvalidationReasonRollback + } + return AppConfigCacheInvalidationReasonActivate +} + +func (i AppConfigOperationSummaryInput) summaryID() string { + return configOperationSummaryID( + i.NextActive.TenantID, + i.NextActive.AppID, + i.Operation, + i.PreviousActive.Version, + i.NextActive.Version, + i.OperationID, + ) +} + +func configOperationSummaryID( + tenantID string, + appID string, + operation AppConfigOperation, + previousVersion string, + nextVersion string, + operationID string, +) string { + return configOperationSummaryIDPrefix + shortHash( + strings.TrimSpace(tenantID), + strings.TrimSpace(appID), + string(operation), + strings.TrimSpace(previousVersion), + strings.TrimSpace(nextVersion), + strings.TrimSpace(operationID), + ) +} + +func isConfigOperationSummaryID(value string) bool { + if !strings.HasPrefix(value, configOperationSummaryIDPrefix) { + return false + } + return isShortHash(strings.TrimPrefix(value, configOperationSummaryIDPrefix)) +} + +func validateConfigOperationOwner(s AppConfigOperationSummary, tenantID, appID string) error { + if strings.TrimSpace(tenantID) != strings.TrimSpace(s.TenantID) { + return fmt.Errorf("config operation tenant_id must match") + } + if strings.TrimSpace(appID) != strings.TrimSpace(s.AppID) { + return fmt.Errorf("config operation app_id must match") + } + return nil +} + +func (o AppConfigOperation) valid() bool { + switch o { + case AppConfigOperationActivate, + AppConfigOperationRollback: + return true + default: + return false + } +} diff --git a/platform/config_operation_summary_test.go b/platform/config_operation_summary_test.go new file mode 100644 index 0000000000..a5a962f6d2 --- /dev/null +++ b/platform/config_operation_summary_test.go @@ -0,0 +1,354 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "errors" + "fmt" + "strings" + "testing" + "time" +) + +func TestNewAppConfigOperationSummaryBuildsActivationSummary(t *testing.T) { + createdAt := time.Date(2026, 7, 8, 15, 0, 0, 0, time.UTC) + previous := validLifecycleConfigVersion("v1", AppConfigVersionStatusRollback) + previous.Checksum = "sha256:previous" + next := validLifecycleConfigVersion("v2", AppConfigVersionStatusActive) + next.Checksum = "sha256:next" + candidate := validLifecycleConfigVersion("v3", AppConfigVersionStatusReleased) + candidate.Checksum = "sha256:candidate" + candidate.GrayPercent = 20 + + summary, err := NewAppConfigOperationSummary(AppConfigOperationSummaryInput{ + Operation: AppConfigOperationActivate, + PreviousActive: previous, + NextActive: next, + ResultVersions: []AppConfigVersion{previous, next, candidate}, + OperationID: "activate-1", + TraceID: "trace-1", + CreatedAt: createdAt, + }) + if err != nil { + t.Fatalf("new activation operation summary: %v", err) + } + if summary.TenantID != "tenant" || summary.AppID != "app" { + t.Fatalf("unexpected owner: %+v", summary) + } + if summary.Operation != AppConfigOperationActivate || + summary.OperationID != "activate-1" || + !summary.CreatedAt.Equal(createdAt) { + t.Fatalf("unexpected operation metadata: %+v", summary) + } + if summary.PreviousVersion != "v1" || summary.PreviousChecksum != "sha256:previous" || + summary.NextVersion != "v2" || summary.NextChecksum != "sha256:next" { + t.Fatalf("unexpected version summary: %+v", summary) + } + if !strings.HasPrefix(summary.SummaryID, configOperationSummaryIDPrefix) { + t.Fatalf("unexpected summary id: %q", summary.SummaryID) + } + if summary.DiffChangeCount <= 0 || !summary.RequiresCacheFlush { + t.Fatalf("expected positive diff and cache flush: %+v", summary) + } + if summary.CacheInvalidation.Reason != AppConfigCacheInvalidationReasonActivate || + summary.CacheInvalidation.OperationID != "activate-1" || + summary.CacheInvalidation.NextVersion != "v2" { + t.Fatalf("unexpected cache invalidation marker: %+v", summary.CacheInvalidation) + } + if summary.GrayStatus.ActiveVersion != "v2" || + !summary.GrayStatus.HasCandidate || + summary.GrayStatus.CandidateVersion != "v3" || + summary.GrayStatus.CandidateTrafficPercent != 20 { + t.Fatalf("unexpected gray status: %+v", summary.GrayStatus) + } + serialized := fmt.Sprintf("%+v", summary) + if strings.Contains(serialized, "model_profile_id") || + strings.Contains(serialized, "tool_policy_id") || + strings.Contains(serialized, "api_key_ref") { + t.Fatalf("summary leaked config bundle content: %s", serialized) + } + + again, err := NewAppConfigOperationSummary(AppConfigOperationSummaryInput{ + Operation: AppConfigOperationActivate, + PreviousActive: previous, + NextActive: next, + ResultVersions: []AppConfigVersion{previous, next, candidate}, + OperationID: "activate-1", + TraceID: "trace-2", + CreatedAt: createdAt.Add(time.Minute), + }) + if err != nil { + t.Fatalf("new duplicate operation summary: %v", err) + } + if summary.SummaryID != again.SummaryID { + t.Fatalf("expected stable summary id, got %q and %q", summary.SummaryID, again.SummaryID) + } + + nextOperation, err := NewAppConfigOperationSummary(AppConfigOperationSummaryInput{ + Operation: AppConfigOperationActivate, + PreviousActive: previous, + NextActive: next, + ResultVersions: []AppConfigVersion{previous, next, candidate}, + OperationID: "activate-2", + CreatedAt: createdAt, + }) + if err != nil { + t.Fatalf("new next operation summary: %v", err) + } + if summary.SummaryID == nextOperation.SummaryID { + t.Fatalf("expected operation id to scope summary id, got %q", summary.SummaryID) + } +} + +func TestNewAppConfigOperationSummaryBuildsRollbackSummary(t *testing.T) { + previous := validLifecycleConfigVersion("v2", AppConfigVersionStatusRollback) + previous.Checksum = "sha256:previous-active" + next := validLifecycleConfigVersion("v1", AppConfigVersionStatusActive) + next.Checksum = "sha256:rollback-active" + + summary, err := NewAppConfigOperationSummary(AppConfigOperationSummaryInput{ + Operation: AppConfigOperationRollback, + PreviousActive: previous, + NextActive: next, + ResultVersions: []AppConfigVersion{previous, next}, + OperationID: "rollback-1", + CreatedAt: time.Now(), + }) + if err != nil { + t.Fatalf("new rollback operation summary: %v", err) + } + if summary.Operation != AppConfigOperationRollback || + summary.CacheInvalidation.Reason != AppConfigCacheInvalidationReasonRollback || + summary.GrayStatus.ActiveVersion != "v1" { + t.Fatalf("unexpected rollback summary: %+v", summary) + } +} + +func TestNewAppConfigOperationSummaryRejectsInvalidInputs(t *testing.T) { + base := validAppConfigOperationSummaryInput() + + unknownOperation := base + unknownOperation.Operation = "promote" + if _, err := NewAppConfigOperationSummary(unknownOperation); err == nil || + !strings.Contains(err.Error(), "invalid config operation") { + t.Fatalf("expected operation validation, got %v", err) + } + + missingTenant := base + missingTenant.NextActive.TenantID = " " + if _, err := NewAppConfigOperationSummary(missingTenant); !errors.Is(err, ErrTenantIDRequired) { + t.Fatalf("expected tenant requirement, got %v", err) + } + + previousNotRollback := base + previousNotRollback.PreviousActive.Status = AppConfigVersionStatusActive + if _, err := NewAppConfigOperationSummary(previousNotRollback); err == nil || + !strings.Contains(err.Error(), "rollback") { + t.Fatalf("expected previous rollback status requirement, got %v", err) + } + + nextNotActive := base + nextNotActive.NextActive.Status = AppConfigVersionStatusReleased + if _, err := NewAppConfigOperationSummary(nextNotActive); err == nil || + !strings.Contains(err.Error(), "active") { + t.Fatalf("expected next active status requirement, got %v", err) + } + + sameVersion := base + sameVersion.NextActive.Version = sameVersion.PreviousActive.Version + if _, err := NewAppConfigOperationSummary(sameVersion); err == nil || + !strings.Contains(err.Error(), "change active version") { + t.Fatalf("expected version switch validation, got %v", err) + } + + missingOperation := base + missingOperation.OperationID = " " + if _, err := NewAppConfigOperationSummary(missingOperation); err == nil || + !strings.Contains(err.Error(), "operation_id") { + t.Fatalf("expected operation id requirement, got %v", err) + } + + missingResults := base + missingResults.ResultVersions = nil + if _, err := NewAppConfigOperationSummary(missingResults); err == nil || + !strings.Contains(err.Error(), "result_versions") { + t.Fatalf("expected result versions requirement, got %v", err) + } + + mismatchedResult := base + mismatchedResult.ResultVersions = append([]AppConfigVersion(nil), base.ResultVersions...) + mismatchedResult.ResultVersions[0].AppID = "other-app" + if _, err := NewAppConfigOperationSummary(mismatchedResult); err == nil || + !strings.Contains(err.Error(), "app_id") { + t.Fatalf("expected result version app mismatch, got %v", err) + } + + zeroCreatedAt := base + zeroCreatedAt.CreatedAt = time.Time{} + if _, err := NewAppConfigOperationSummary(zeroCreatedAt); err == nil || + !strings.Contains(err.Error(), "created_at") { + t.Fatalf("expected created at requirement, got %v", err) + } + + sensitiveTrace := validAppConfigOperationSummaryInput() + sensitiveTrace.TraceID = "Authorization: Bearer raw-token" + if _, err := NewAppConfigOperationSummary(sensitiveTrace); err == nil || + !strings.Contains(err.Error(), "trace_id") { + t.Fatalf("expected sensitive trace rejection, got %v", err) + } +} + +func TestAppConfigOperationSummaryValidateRejectsUnsafeSummary(t *testing.T) { + generated, err := NewAppConfigOperationSummary(validAppConfigOperationSummaryInput()) + if err != nil { + t.Fatalf("new generated config operation summary: %v", err) + } + summary := generated + if err := summary.Validate(); err != nil { + t.Fatalf("expected summary to validate: %v", err) + } + + summary.SummaryID = "config_operation_v1" + if err := summary.Validate(); err == nil || + !strings.Contains(err.Error(), "summary_id") { + t.Fatalf("expected unsafe summary id rejection, got %v", err) + } + + summary = generated + summary.OperationID = "activate-2" + if err := summary.Validate(); err == nil || + !strings.Contains(err.Error(), "summary_id") { + t.Fatalf("expected stale summary id rejection, got %v", err) + } + + summary = generated + summary.RequiresCacheFlush = false + if err := summary.Validate(); err == nil || + !strings.Contains(err.Error(), "requires_cache_flush") { + t.Fatalf("expected cache flush requirement, got %v", err) + } + + summary = generated + summary.DiffChangeCount = 0 + if err := summary.Validate(); err == nil || + !strings.Contains(err.Error(), "diff_change_count") { + t.Fatalf("expected positive diff count requirement, got %v", err) + } + + summary = generated + summary.GrayStatus.ActiveVersion = "other-version" + if err := summary.Validate(); err == nil || + !strings.Contains(err.Error(), "gray_status") { + t.Fatalf("expected gray status consistency rejection, got %v", err) + } + + summary = generated + summary.CacheInvalidation.Reason = AppConfigCacheInvalidationReasonRollback + if err := summary.Validate(); err == nil || + !strings.Contains(err.Error(), "cache_invalidation") { + t.Fatalf("expected cache invalidation reason mismatch rejection, got %v", err) + } + + summary = generated + summary.CacheInvalidation.OperationID = "activate-2" + if err := summary.Validate(); err == nil || + !strings.Contains(err.Error(), "cache_invalidation") { + t.Fatalf("expected cache invalidation operation mismatch rejection, got %v", err) + } + + summary = generated + summary.CacheInvalidation.NextVersion = "v3" + if err := summary.Validate(); err == nil || + !strings.Contains(err.Error(), "cache_invalidation") { + t.Fatalf("expected cache invalidation version mismatch rejection, got %v", err) + } + + summary = generated + summary.CacheInvalidation.CreatedAt = summary.CreatedAt.Add(time.Minute) + if err := summary.Validate(); err == nil || + !strings.Contains(err.Error(), "cache_invalidation") { + t.Fatalf("expected cache invalidation time mismatch rejection, got %v", err) + } + + summary = generated + summary.GrayStatus.HasCandidate = true + summary.GrayStatus.CandidateVersion = "candidate" + summary.GrayStatus.CandidateChecksum = "Authorization: Bearer raw-token" + if err := summary.Validate(); err == nil || + !strings.Contains(err.Error(), "gray_candidate_checksum") { + t.Fatalf("expected sensitive gray candidate rejection, got %v", err) + } + + summary = generated + summary.GrayStatus.HasRollback = false + summary.GrayStatus.RollbackVersion = "" + summary.GrayStatus.RollbackChecksum = "" + if err := summary.Validate(); err == nil || + !strings.Contains(err.Error(), "rollback") { + t.Fatalf("expected missing rollback status rejection, got %v", err) + } + + summary = generated + summary.GrayStatus.RollbackVersion = "other-version" + if err := summary.Validate(); err == nil || + !strings.Contains(err.Error(), "rollback") { + t.Fatalf("expected rollback status mismatch rejection, got %v", err) + } + + summary = generated + summary.GrayStatus.HasCandidate = true + summary.GrayStatus.CandidateVersion = "v3" + summary.GrayStatus.CandidateChecksum = "sha256:candidate" + summary.GrayStatus.CandidateGrayPercent = 20 + summary.GrayStatus.CandidateTrafficPercent = 10 + if err := summary.Validate(); err == nil || + !strings.Contains(err.Error(), "candidate traffic") { + t.Fatalf("expected candidate traffic mismatch rejection, got %v", err) + } + + summary = generated + summary.GrayStatus.ActiveTrafficPercent = 50 + if err := summary.Validate(); err == nil || + !strings.Contains(err.Error(), "active traffic") { + t.Fatalf("expected no-candidate active traffic rejection, got %v", err) + } + + summary = generated + summary.CacheInvalidation.OperationID = "sk-1234567890abcdef" + if err := summary.Validate(); err == nil || + !strings.Contains(err.Error(), "cache_invalidation") { + t.Fatalf("expected unsafe cache invalidation rejection, got %v", err) + } +} + +func TestNewAppConfigOperationSummaryRequiresRollbackInResultVersions(t *testing.T) { + input := validAppConfigOperationSummaryInput() + input.ResultVersions = []AppConfigVersion{input.NextActive} + + if _, err := NewAppConfigOperationSummary(input); err == nil || + !strings.Contains(err.Error(), "rollback") { + t.Fatalf("expected missing rollback result rejection, got %v", err) + } +} + +func validAppConfigOperationSummaryInput() AppConfigOperationSummaryInput { + previous := validLifecycleConfigVersion("v1", AppConfigVersionStatusRollback) + previous.Checksum = "sha256:previous" + next := validLifecycleConfigVersion("v2", AppConfigVersionStatusActive) + next.Checksum = "sha256:next" + return AppConfigOperationSummaryInput{ + Operation: AppConfigOperationActivate, + PreviousActive: previous, + NextActive: next, + ResultVersions: []AppConfigVersion{previous, next}, + OperationID: "activate-1", + TraceID: "trace", + CreatedAt: time.Now(), + } +} From ce636fec140e587f9bcc788fc35a05816456f338 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 16:58:21 +0800 Subject: [PATCH 32/95] platform/toolpolicy: add approval summary contract --- platform/toolpolicy/policy.go | 250 ++++++++++++++++++++++++++--- platform/toolpolicy/policy_test.go | 134 +++++++++++++++- 2 files changed, 361 insertions(+), 23 deletions(-) diff --git a/platform/toolpolicy/policy.go b/platform/toolpolicy/policy.go index 2a9804ae8c..74451f5530 100644 --- a/platform/toolpolicy/policy.go +++ b/platform/toolpolicy/policy.go @@ -13,6 +13,8 @@ import ( "crypto/sha256" "encoding/hex" "fmt" + "regexp" + "strconv" "strings" "time" @@ -32,6 +34,28 @@ type Policy struct { now func() time.Time } +// ApprovalSummary is the safe approval-facing summary of one tool call. +type ApprovalSummary struct { + TenantID string + AppID string + PolicyID string + ToolName string + ToolCallID string + Decision tool.PermissionAction + Reason string + ArgumentsDigest string + ArgumentsBytes int + RequiresApproval bool + ReadOnly bool + Destructive bool + OpenWorld bool + ConcurrencySafe bool + SearchOrRead bool + MaxResultSize int + RedactionVersion string + CreatedAt time.Time +} + // Option configures Policy. type Option func(*Policy) @@ -125,7 +149,11 @@ func (p *Policy) CheckToolPermission( } decision, reason, audit := p.decide(req, name) if audit { - p.writeAudit(ctx, req, name, string(decision.Action), reason) + summary, err := p.ApprovalSummary(req, decision, reason) + if err != nil { + return tool.PermissionDecision{}, err + } + p.writeAudit(ctx, summary) } return decision, nil } @@ -143,7 +171,11 @@ func (p *Policy) beforeTool() tool.BeforeToolCallbackStructured { } decision, reason, audit := p.decideNameOnly(req, req.ToolName) if audit { - p.writeAudit(ctx, req, req.ToolName, string(decision.Action), reason) + summary, err := p.ApprovalSummary(req, decision, reason) + if err != nil { + return nil, err + } + p.writeAudit(ctx, summary) } var err error if err != nil { @@ -228,7 +260,11 @@ func (r *Reviewer) Review(ctx context.Context, req *review.Request) (*review.Dec } decision, reason, audit := r.policy.decideReviewer(permissionReq, permissionReq.ToolName) if audit { - r.policy.writeAudit(ctx, permissionReq, permissionReq.ToolName, string(decision.Action), reason) + summary, err := r.policy.ApprovalSummary(permissionReq, decision, reason) + if err != nil { + return nil, err + } + r.policy.writeAudit(ctx, summary) } var err error decision, err = tool.NormalizePermissionDecision(decision) @@ -379,34 +415,206 @@ func contains(items []string, target string) bool { return false } -func (p *Policy) writeAudit( - ctx context.Context, +// ApprovalSummary builds a redacted approval summary suitable for audit, +// approval messages, and logs. Raw tool arguments are never included. +func (p *Policy) ApprovalSummary( req *tool.PermissionRequest, - toolName string, - decision string, + decision tool.PermissionDecision, reason string, -) { +) (ApprovalSummary, error) { + if p == nil { + return ApprovalSummary{}, fmt.Errorf("policy is nil") + } + if req == nil { + return ApprovalSummary{}, fmt.Errorf("permission request is nil") + } + name := strings.TrimSpace(req.ToolName) + if name == "" && req.Declaration != nil { + name = strings.TrimSpace(req.Declaration.Name) + } + if name == "" { + return ApprovalSummary{}, fmt.Errorf("tool_name is required") + } + decision, err := tool.NormalizePermissionDecision(decision) + if err != nil { + return ApprovalSummary{}, err + } + reason = strings.TrimSpace(reason) + if reason == "" { + reason = strings.TrimSpace(decision.Reason) + } + if err := platformSafeText("tool_name", name); err != nil { + return ApprovalSummary{}, err + } + if err := platformSafeText("tool_call_id", req.ToolCallID); err != nil { + return ApprovalSummary{}, err + } + if err := platformSafeText("reason", reason); err != nil { + return ApprovalSummary{}, err + } + argumentsDigest, argumentsBytes := argumentDigest(req.Arguments) + summary := ApprovalSummary{ + TenantID: strings.TrimSpace(p.policy.TenantID), + AppID: strings.TrimSpace(p.policy.AppID), + PolicyID: strings.TrimSpace(p.policy.PolicyID), + ToolName: name, + ToolCallID: strings.TrimSpace(req.ToolCallID), + Decision: decision.Action, + Reason: reason, + ArgumentsDigest: argumentsDigest, + ArgumentsBytes: argumentsBytes, + RequiresApproval: decision.Action == tool.PermissionActionAsk, + ReadOnly: req.Metadata.ReadOnly, + Destructive: req.Metadata.Destructive, + OpenWorld: req.Metadata.OpenWorld, + ConcurrencySafe: req.Metadata.ConcurrencySafe, + SearchOrRead: req.Metadata.SearchOrRead, + MaxResultSize: req.Metadata.MaxResultSize, + RedactionVersion: "platform-toolpolicy-v1", + CreatedAt: p.now(), + } + if err := summary.Validate(); err != nil { + return ApprovalSummary{}, err + } + return summary, nil +} + +// Validate checks that the summary is safe to expose outside the tool runtime. +func (s ApprovalSummary) Validate() error { + if strings.TrimSpace(s.ToolName) == "" { + return fmt.Errorf("tool_name is required") + } + if err := platformSafeText("tool_name", s.ToolName); err != nil { + return err + } + if err := platformSafeText("tool_call_id", s.ToolCallID); err != nil { + return err + } + if err := platformSafeText("reason", s.Reason); err != nil { + return err + } + if s.ArgumentsBytes < 0 { + return fmt.Errorf("arguments_bytes must be greater than or equal to 0") + } + if s.ArgumentsBytes == 0 { + if s.ArgumentsDigest != "" { + return fmt.Errorf("arguments_digest must be empty when arguments_bytes is 0") + } + } else if !validSHA256Digest(s.ArgumentsDigest) { + return fmt.Errorf("arguments_digest must be sha256 followed by a 64 character hex digest") + } + if s.MaxResultSize < 0 { + return fmt.Errorf("max_result_size must be greater than or equal to 0") + } + switch s.Decision { + case tool.PermissionActionAllow: + if s.RequiresApproval { + return fmt.Errorf("requires_approval must be false for allow decisions") + } + case tool.PermissionActionDeny: + if s.RequiresApproval { + return fmt.Errorf("requires_approval must be false for deny decisions") + } + case tool.PermissionActionAsk: + if !s.RequiresApproval { + return fmt.Errorf("requires_approval must be true for ask decisions") + } + case "": + return fmt.Errorf("decision is required") + default: + return fmt.Errorf("invalid decision %q", s.Decision) + } + if strings.TrimSpace(s.RedactionVersion) == "" { + return fmt.Errorf("redaction_version is required") + } + if s.CreatedAt.IsZero() { + return fmt.Errorf("created_at is required") + } + if err := platformSafeText("detail_ref", s.DetailRef()); err != nil { + return err + } + return nil +} + +func (p *Policy) writeAudit(ctx context.Context, summary ApprovalSummary) { if p.audit == nil { return } - argsSummary := argumentSummary(req.Arguments) + detailRef := summary.DetailRef() _ = p.audit.WriteAudit(ctx, platform.AuditRecord{ - AuditID: platform.AuditID(p.policy.TenantID, p.policy.AppID, toolName, req.ToolCallID, decision, argsSummary), - TenantID: p.policy.TenantID, - AppID: p.policy.AppID, - ToolName: toolName, - Decision: decision, - DecisionReason: reason, - RedactedDetailRef: argsSummary, - RedactionVersion: "platform-toolpolicy-v1", - CreatedAt: p.now(), + AuditID: platform.AuditID(summary.TenantID, summary.AppID, summary.ToolName, summary.ToolCallID, string(summary.Decision), detailRef), + TenantID: summary.TenantID, + AppID: summary.AppID, + ToolName: summary.ToolName, + Decision: string(summary.Decision), + DecisionReason: summary.Reason, + RedactedDetailRef: detailRef, + RedactionVersion: summary.RedactionVersion, + CreatedAt: summary.CreatedAt, }) } -func argumentSummary(args []byte) string { +// DetailRef returns compact non-secret detail that can be stored in audit logs. +func (s ApprovalSummary) DetailRef() string { + parts := []string{ + "tool:" + s.ToolName, + "decision:" + string(s.Decision), + } + if s.ToolCallID != "" { + parts = append(parts, "tool_call_id:"+s.ToolCallID) + } + if s.ArgumentsDigest != "" { + parts = append(parts, "args:"+s.ArgumentsDigest) + parts = append(parts, "args_bytes:"+strconv.Itoa(s.ArgumentsBytes)) + } + if s.RequiresApproval { + parts = append(parts, "requires_approval:true") + } + if s.ReadOnly { + parts = append(parts, "read_only:true") + } + if s.Destructive { + parts = append(parts, "destructive:true") + } + if s.OpenWorld { + parts = append(parts, "open_world:true") + } + return strings.Join(parts, " ") +} + +func argumentDigest(args []byte) (string, int) { if len(args) == 0 { - return "" + return "", 0 } sum := sha256.Sum256(args) - return fmt.Sprintf("sha256:%s bytes:%d", hex.EncodeToString(sum[:]), len(args)) + return "sha256:" + hex.EncodeToString(sum[:]), len(args) +} + +func argumentSummary(args []byte) string { + digest, size := argumentDigest(args) + if digest == "" { + return "" + } + return fmt.Sprintf("%s bytes:%d", digest, size) +} + +var sha256DigestPattern = regexp.MustCompile(`^sha256:[a-f0-9]{64}$`) + +func validSHA256Digest(value string) bool { + return sha256DigestPattern.MatchString(value) +} + +func platformSafeText(field, value string) error { + value = strings.TrimSpace(value) + if value == "" { + return nil + } + redactor, err := platform.NewRedactor() + if err != nil { + return fmt.Errorf("%s: redactor unavailable: %w", field, err) + } + if redactor.Redact(value) != value { + return fmt.Errorf("%s contains unredacted sensitive content", field) + } + return nil } diff --git a/platform/toolpolicy/policy_test.go b/platform/toolpolicy/policy_test.go index 9baa26b207..1b899e0940 100644 --- a/platform/toolpolicy/policy_test.go +++ b/platform/toolpolicy/policy_test.go @@ -118,8 +118,138 @@ func TestPolicyAllowsHighRiskWithAuditAndRedactsArguments(t *testing.T) { strings.Contains(record.RedactedDetailRef, "Authorization") { t.Fatalf("audit leaked raw argument content: %q", record.RedactedDetailRef) } - if !strings.HasPrefix(record.RedactedDetailRef, "sha256:") { - t.Fatalf("expected digest summary, got %q", record.RedactedDetailRef) + if !strings.Contains(record.RedactedDetailRef, "args:sha256:") || + !strings.Contains(record.RedactedDetailRef, "args_bytes:") || + !strings.Contains(record.RedactedDetailRef, "decision:allow") { + t.Fatalf("expected safe detail summary, got %q", record.RedactedDetailRef) + } +} + +func TestPolicyBuildsApprovalSummaryWithoutRawArguments(t *testing.T) { + now := time.Unix(200, 0) + p := newPolicy( + t, + platform.ToolPolicy{ + TenantID: "tenant", + AppID: "app", + PolicyID: "policy", + DangerousToolAction: platform.DangerousToolActionAsk, + HighRiskTools: []string{"workspace_write"}, + }, + WithNow(func() time.Time { return now }), + ) + req := request( + "workspace_write", + tool.ToolMetadata{ + Destructive: true, + OpenWorld: true, + ConcurrencySafe: false, + MaxResultSize: 4096, + }, + []byte(`{"path":"/private/file","api_key":"sk-secret"}`), + ) + req.ToolCallID = "call-1" + + decision, err := p.CheckToolPermission(context.Background(), req) + if err != nil { + t.Fatalf("CheckToolPermission: %v", err) + } + summary, err := p.ApprovalSummary(req, decision, decision.Reason) + if err != nil { + t.Fatalf("ApprovalSummary: %v", err) + } + if summary.TenantID != "tenant" || + summary.AppID != "app" || + summary.PolicyID != "policy" || + summary.ToolName != "workspace_write" || + summary.ToolCallID != "call-1" || + summary.Decision != tool.PermissionActionAsk || + !summary.RequiresApproval || + !summary.Destructive || + !summary.OpenWorld || + summary.MaxResultSize != 4096 || + !summary.CreatedAt.Equal(now) { + t.Fatalf("unexpected summary: %+v", summary) + } + if summary.ArgumentsBytes == 0 || !strings.HasPrefix(summary.ArgumentsDigest, "sha256:") { + t.Fatalf("expected argument digest, got %+v", summary) + } + detail := summary.DetailRef() + if strings.Contains(detail, "sk-secret") || + strings.Contains(detail, "/private/file") || + strings.Contains(detail, "api_key") { + t.Fatalf("summary detail leaked raw arguments: %q", detail) + } + if !strings.Contains(detail, "requires_approval:true") || + !strings.Contains(detail, "destructive:true") || + !strings.Contains(detail, "open_world:true") { + t.Fatalf("summary detail missing risk markers: %q", detail) + } +} + +func TestApprovalSummaryValidationRejectsUnsafeOrInconsistentFields(t *testing.T) { + now := time.Unix(300, 0) + valid := ApprovalSummary{ + ToolName: "workspace_write", + ToolCallID: "call-1", + Decision: tool.PermissionActionAsk, + Reason: "high-risk tool requires approval", + ArgumentsDigest: "sha256:9f86d081884c7d659a2feaa0c55ad015a3bf4f1b2b0b822cd15d6c15b0f00a08", + ArgumentsBytes: 16, + RequiresApproval: true, + RedactionVersion: "platform-toolpolicy-v1", + CreatedAt: now, + } + if err := valid.Validate(); err != nil { + t.Fatalf("Validate valid summary: %v", err) + } + + unsafe := valid + unsafe.Reason = "sk-secret-token" + if err := unsafe.Validate(); err == nil || !strings.Contains(err.Error(), "reason") { + t.Fatalf("expected unsafe reason rejection, got %v", err) + } + + unsafeToolName := valid + unsafeToolName.ToolName = "api_key: sk-secret-token" + if err := unsafeToolName.Validate(); err == nil || !strings.Contains(err.Error(), "tool_name") { + t.Fatalf("expected unsafe tool name rejection, got %v", err) + } + + unsafeToolCallID := valid + unsafeToolCallID.ToolCallID = "token=sk-secret-token" + if err := unsafeToolCallID.Validate(); err == nil || !strings.Contains(err.Error(), "tool_call_id") { + t.Fatalf("expected unsafe tool call id rejection, got %v", err) + } + + wrongApproval := valid + wrongApproval.RequiresApproval = false + if err := wrongApproval.Validate(); err == nil || !strings.Contains(err.Error(), "requires_approval") { + t.Fatalf("expected ask approval invariant, got %v", err) + } + + wrongDigest := valid + wrongDigest.ArgumentsDigest = "raw-json" + if err := wrongDigest.Validate(); err == nil || !strings.Contains(err.Error(), "arguments_digest") { + t.Fatalf("expected digest prefix rejection, got %v", err) + } + + unsafeDigest := valid + unsafeDigest.ArgumentsDigest = "sha256:sk-secret-token" + if err := unsafeDigest.Validate(); err == nil || !strings.Contains(err.Error(), "arguments_digest") { + t.Fatalf("expected digest hex rejection, got %v", err) + } + + noArguments := valid + noArguments.ArgumentsBytes = 0 + if err := noArguments.Validate(); err == nil || !strings.Contains(err.Error(), "arguments_digest") { + t.Fatalf("expected empty-arguments digest rejection, got %v", err) + } + + allowNeedsApproval := valid + allowNeedsApproval.Decision = tool.PermissionActionAllow + if err := allowNeedsApproval.Validate(); err == nil || !strings.Contains(err.Error(), "requires_approval") { + t.Fatalf("expected allow approval invariant, got %v", err) } } From bc201e065c0d8d1446d2a27db42a8ca3f7d424ca Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Thu, 9 Jul 2026 12:38:32 +0800 Subject: [PATCH 33/95] fix(platform): enforce binding ACL and redact outbox errors --- platform/channeladapter/outbox.go | 6 +- platform/channeladapter/outbox_test.go | 74 ++++++++++++++++++ platform/gateway/errors.go | 4 + platform/gateway/service.go | 33 ++++++++ platform/gateway/service_test.go | 104 +++++++++++++++++++++++++ 5 files changed, 220 insertions(+), 1 deletion(-) diff --git a/platform/channeladapter/outbox.go b/platform/channeladapter/outbox.go index 0ca18a5e11..1885bb9535 100644 --- a/platform/channeladapter/outbox.go +++ b/platform/channeladapter/outbox.go @@ -428,7 +428,11 @@ func errorString(err error) string { if err == nil { return "" } - return err.Error() + redactor, redactorErr := platform.NewRedactor() + if redactorErr != nil { + return err.Error() + } + return redactor.Redact(err.Error()) } func sameOutboundIdentity(existing platform.OutboundMessage, next platform.OutboundMessage) bool { diff --git a/platform/channeladapter/outbox_test.go b/platform/channeladapter/outbox_test.go index f9ee7c8584..e714a55578 100644 --- a/platform/channeladapter/outbox_test.go +++ b/platform/channeladapter/outbox_test.go @@ -11,6 +11,7 @@ package channeladapter import ( "context" "errors" + "strings" "testing" "time" @@ -158,6 +159,79 @@ func TestOutboxFailureSchedulesRetryThenDeadLetter(t *testing.T) { } } +func TestOutboxRedactsFailureDetails(t *testing.T) { + ctx := context.Background() + store := NewInMemoryOutboxStore() + msg := outbound("reply-1") + _, _, err := store.Enqueue(ctx, msg, RetryPolicy{MaxAttempts: 2}) + if err != nil { + t.Fatalf("enqueue: %v", err) + } + now := time.Now().Add(time.Hour) + claimed, err := store.ClaimDue(ctx, now, 1, time.Minute) + if err != nil { + t.Fatalf("claim due: %v", err) + } + + failed, err := store.MarkFailed( + ctx, + msg.DedupKey, + claimed[0].LeaseToken, + errors.New("provider failed Authorization: Bearer raw-token postgres://user:pass@example/db api_key=sk-1234567890abcdef"), + now, + ) + if err != nil { + t.Fatalf("mark failed: %v", err) + } + + for _, leaked := range []string{"raw-token", ":pass@", "sk-1234567890abcdef"} { + if strings.Contains(failed.LastError, leaked) { + t.Fatalf("failure detail leaked %q: %q", leaked, failed.LastError) + } + } + if !strings.Contains(failed.LastError, "Authorization: ****") || + !strings.Contains(failed.LastError, "user:****@") || + !strings.Contains(failed.LastError, "api_key=****") { + t.Fatalf("expected redacted diagnostic, got %q", failed.LastError) + } +} + +func TestOutboxRedactsDeadLetterDetails(t *testing.T) { + ctx := context.Background() + store := NewInMemoryOutboxStore() + msg := outbound("reply-1") + _, _, err := store.Enqueue(ctx, msg, RetryPolicy{MaxAttempts: 1}) + if err != nil { + t.Fatalf("enqueue: %v", err) + } + now := time.Now().Add(time.Hour) + claimed, err := store.ClaimDue(ctx, now, 1, time.Minute) + if err != nil { + t.Fatalf("claim due: %v", err) + } + + dead, err := store.MarkDeadLetter( + ctx, + msg.DedupKey, + claimed[0].LeaseToken, + errors.New("permanent token=secret-value cookie=session-secret"), + now, + ) + if err != nil { + t.Fatalf("mark dead letter: %v", err) + } + + for _, leaked := range []string{"secret-value", "session-secret"} { + if strings.Contains(dead.LastError, leaked) { + t.Fatalf("dead-letter detail leaked %q: %q", leaked, dead.LastError) + } + } + if !strings.Contains(dead.LastError, "token=****") || + !strings.Contains(dead.LastError, "cookie=****") { + t.Fatalf("expected redacted diagnostic, got %q", dead.LastError) + } +} + func TestOutboxListsAndRequeuesDeadLetter(t *testing.T) { ctx := context.Background() store := NewInMemoryOutboxStore() diff --git a/platform/gateway/errors.go b/platform/gateway/errors.go index 7908bfd139..9db62b54db 100644 --- a/platform/gateway/errors.go +++ b/platform/gateway/errors.go @@ -17,6 +17,10 @@ var ( ErrRuntimeInactive = errors.New("gateway runtime inactive") // ErrRuntimeMismatch indicates that a runtime's tenant, app, binding, or inbound identifiers do not match. ErrRuntimeMismatch = errors.New("gateway runtime identifiers mismatch") + // ErrBindingAccessDenied indicates that a binding policy rejects the inbound sender or conversation. + ErrBindingAccessDenied = errors.New("gateway binding access denied") + // ErrBindingMentionRequired indicates that a group/thread message did not mention the agent. + ErrBindingMentionRequired = errors.New("gateway binding mention required") // ErrUnsupportedMessageType indicates that the gateway batch only supports text input. ErrUnsupportedMessageType = errors.New("gateway only supports text messages") // ErrEmptyText indicates that a text message does not contain usable text. diff --git a/platform/gateway/service.go b/platform/gateway/service.go index 01e09f01b3..f61d730b12 100644 --- a/platform/gateway/service.go +++ b/platform/gateway/service.go @@ -122,6 +122,10 @@ func (s *Service) HandleInbound( s.writeAudit(ctx, auditFromMessage(msg, "", "", "reject", err.Error(), start, err)) return Result{}, err } + if err := authorizeBinding(runtime.Binding, msg); err != nil { + s.writeAudit(ctx, auditFromMessage(msg, "", "", "reject", err.Error(), start, err)) + return Result{}, err + } text, err := inboundText(msg) if err != nil { s.writeAudit(ctx, auditFromMessage(msg, "", "", "reject", err.Error(), start, err)) @@ -285,6 +289,35 @@ func (s *Service) duplicateResult( return result, nil } +func authorizeBinding(binding platform.ChannelBinding, msg platform.InboundMessage) error { + if !containsAllowed(binding.AllowedUsers, msg.ExternalUserID) { + return ErrBindingAccessDenied + } + if msg.ConversationType != platform.ConversationTypeDM && + !containsAllowed(binding.AllowedGroups, msg.ExternalGroupID) { + return ErrBindingAccessDenied + } + if binding.RequiredMention && + msg.ConversationType != platform.ConversationTypeDM && + !msg.RequiredMentionSeen { + return ErrBindingMentionRequired + } + return nil +} + +func containsAllowed(allowed []string, value string) bool { + if len(allowed) == 0 { + return true + } + value = strings.TrimSpace(value) + for _, candidate := range allowed { + if strings.TrimSpace(candidate) == value { + return true + } + } + return false +} + func inboundText(msg platform.InboundMessage) (string, error) { if msg.MessageType != platform.MessageTypeText { return "", ErrUnsupportedMessageType diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go index 72f5afacbc..2fbfabaca2 100644 --- a/platform/gateway/service_test.go +++ b/platform/gateway/service_test.go @@ -384,6 +384,110 @@ func TestServiceHandleInboundRejectsUnsupportedMessage(t *testing.T) { assert.NotEqual(t, "user-1", audit.Records()[0].UserID) } +func TestServiceHandleInboundRejectsDisallowedUser(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "unused"} + runtime := validRuntime("tenant-a", r) + runtime.Binding.AllowedUsers = []string{"allowed-user"} + require.NoError(t, registry.Register(runtime)) + audit := platform.NewInMemoryAuditSink() + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + WithAuditSink(audit), + ) + + _, err := svc.HandleInbound(ctx, inbound("tenant-a", "msg-1", "blocked-user", "hello")) + + require.ErrorIs(t, err, ErrBindingAccessDenied) + assert.Empty(t, r.calls) + require.Len(t, audit.Records(), 1) + assert.Equal(t, "reject", audit.Records()[0].Decision) + assert.Equal(t, ErrBindingAccessDenied.Error(), audit.Records()[0].DecisionReason) +} + +func TestServiceHandleInboundRejectsDisallowedGroup(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "unused"} + runtime := validRuntime("tenant-a", r) + runtime.Binding.AllowedGroups = []string{"allowed-group"} + require.NoError(t, registry.Register(runtime)) + audit := platform.NewInMemoryAuditSink() + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + WithAuditSink(audit), + ) + msg := inbound("tenant-a", "msg-1", "user-1", "hello") + msg.ConversationType = platform.ConversationTypeGroup + msg.ExternalGroupID = "blocked-group" + + _, err := svc.HandleInbound(ctx, msg) + + require.ErrorIs(t, err, ErrBindingAccessDenied) + assert.Empty(t, r.calls) + require.Len(t, audit.Records(), 1) + assert.Equal(t, "reject", audit.Records()[0].Decision) +} + +func TestServiceHandleInboundRejectsMissingRequiredMention(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "unused"} + runtime := validRuntime("tenant-a", r) + runtime.Binding.RequiredMention = true + require.NoError(t, registry.Register(runtime)) + audit := platform.NewInMemoryAuditSink() + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + WithAuditSink(audit), + ) + msg := inbound("tenant-a", "msg-1", "user-1", "hello") + msg.ConversationType = platform.ConversationTypeGroup + msg.ExternalGroupID = "group-1" + msg.RequiredMentionSeen = false + + _, err := svc.HandleInbound(ctx, msg) + + require.ErrorIs(t, err, ErrBindingMentionRequired) + assert.Empty(t, r.calls) + require.Len(t, audit.Records(), 1) + assert.Equal(t, "reject", audit.Records()[0].Decision) + assert.Equal(t, ErrBindingMentionRequired.Error(), audit.Records()[0].DecisionReason) +} + +func TestServiceHandleInboundAllowsAuthorizedGroupMention(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "authorized"} + runtime := validRuntime("tenant-a", r) + runtime.Binding.AllowedUsers = []string{"user-1"} + runtime.Binding.AllowedGroups = []string{"group-1"} + runtime.Binding.RequiredMention = true + require.NoError(t, registry.Register(runtime)) + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + ) + msg := inbound("tenant-a", "msg-1", "user-1", "hello") + msg.ConversationType = platform.ConversationTypeGroup + msg.ExternalGroupID = "group-1" + msg.RequiredMentionSeen = true + + result, err := svc.HandleInbound(ctx, msg) + + require.NoError(t, err) + assert.Equal(t, "authorized", result.Outbound.Content) + require.Len(t, r.calls, 1) +} + func TestServiceHandleInboundRunnerErrorDoesNotComplete(t *testing.T) { ctx := context.Background() registry := NewInMemoryRegistry() From 3d57ce89cea6bae0509af39d9d6b55fb97e6133f Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Thu, 9 Jul 2026 12:41:30 +0800 Subject: [PATCH 34/95] fix(platform): redact gateway audit error reasons --- platform/gateway/service.go | 13 ++++++++++++- platform/gateway/service_test.go | 27 +++++++++++++++++++++++++++ 2 files changed, 39 insertions(+), 1 deletion(-) diff --git a/platform/gateway/service.go b/platform/gateway/service.go index f61d730b12..6dd611ed00 100644 --- a/platform/gateway/service.go +++ b/platform/gateway/service.go @@ -424,7 +424,7 @@ func auditFromMessage( MessageID: msg.PlatformMessageID, RequestID: requestIDFor(msg), Decision: decision, - DecisionReason: reason, + DecisionReason: redactAuditReason(reason), LatencyMS: time.Since(start).Milliseconds(), CreatedAt: time.Now(), } @@ -434,6 +434,17 @@ func auditFromMessage( return record } +func redactAuditReason(reason string) string { + if reason == "" { + return "" + } + redactor, err := platform.NewRedactor() + if err != nil { + return reason + } + return redactor.Redact(reason) +} + func (s *Service) writeAudit(ctx context.Context, record platform.AuditRecord) { if s.auditSink == nil { return diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go index 2fbfabaca2..96daf29088 100644 --- a/platform/gateway/service_test.go +++ b/platform/gateway/service_test.go @@ -508,6 +508,33 @@ func TestServiceHandleInboundRunnerErrorDoesNotComplete(t *testing.T) { assert.Empty(t, record.ResultRef) } +func TestServiceHandleInboundRunnerErrorRedactsAuditReason(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + runnerErr := errors.New("runner failed Authorization: Bearer raw-token api_key=sk-1234567890abcdef") + r := &recordingRunner{runErr: runnerErr} + registerRuntime(t, registry, "tenant-a", r) + audit := platform.NewInMemoryAuditSink() + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + WithAuditSink(audit), + ) + msg := inbound("tenant-a", "msg-1", "user-1", "hello") + + _, err := svc.HandleInbound(ctx, msg) + + require.ErrorIs(t, err, runnerErr) + records := audit.Records() + require.Len(t, records, 1) + assert.Equal(t, "runner_error", records[0].Decision) + assert.NotContains(t, records[0].DecisionReason, "raw-token") + assert.NotContains(t, records[0].DecisionReason, "sk-1234567890abcdef") + assert.Contains(t, records[0].DecisionReason, "Authorization: ****") + assert.Contains(t, records[0].DecisionReason, "api_key=****") +} + func TestServiceHandleInboundUsesRequestIDAndStreamsText(t *testing.T) { ctx := context.Background() registry := NewInMemoryRegistry() From 597b46c5b7bf85ee9714e9ac697bebd23eafe0f6 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Thu, 9 Jul 2026 14:24:42 +0800 Subject: [PATCH 35/95] fix(toolpolicy): enforce auditable policy identity --- platform/toolpolicy/policy.go | 66 +++++++++++++++++----- platform/toolpolicy/policy_test.go | 89 +++++++++++++++++++++++++++--- 2 files changed, 135 insertions(+), 20 deletions(-) diff --git a/platform/toolpolicy/policy.go b/platform/toolpolicy/policy.go index 74451f5530..11d4ab58a6 100644 --- a/platform/toolpolicy/policy.go +++ b/platform/toolpolicy/policy.go @@ -76,7 +76,8 @@ func WithAuditSink(sink platform.AuditSink) Option { } } -// WithRedactor sets the redactor used before writing tool arguments to audit. +// WithRedactor overrides the configured redactor. Approval summaries still +// store only argument digests and never include raw tool arguments. func WithRedactor(redactor *platform.Redactor) Option { return func(p *Policy) { if redactor != nil { @@ -99,6 +100,9 @@ func New(policy platform.ToolPolicy, opts ...Option) (*Policy, error) { if err := validate(policy); err != nil { return nil, err } + if err := validateRuntimeIdentity(policy); err != nil { + return nil, err + } redactor, err := platform.NewRedactor(policy.ArgumentRedactionRules...) if err != nil { return nil, fmt.Errorf("newing platform tool policy: redaction rules: %w", err) @@ -153,7 +157,9 @@ func (p *Policy) CheckToolPermission( if err != nil { return tool.PermissionDecision{}, err } - p.writeAudit(ctx, summary) + if err := p.writeAudit(ctx, summary); err != nil { + return tool.PermissionDecision{}, err + } } return decision, nil } @@ -175,13 +181,11 @@ func (p *Policy) beforeTool() tool.BeforeToolCallbackStructured { if err != nil { return nil, err } - p.writeAudit(ctx, summary) - } - var err error - if err != nil { - return nil, err + if err := p.writeAudit(ctx, summary); err != nil { + return nil, err + } } - decision, err = tool.NormalizePermissionDecision(decision) + decision, err := tool.NormalizePermissionDecision(decision) if err != nil { return nil, err } @@ -264,7 +268,9 @@ func (r *Reviewer) Review(ctx context.Context, req *review.Request) (*review.Dec if err != nil { return nil, err } - r.policy.writeAudit(ctx, summary) + if err := r.policy.writeAudit(ctx, summary); err != nil { + return nil, err + } } var err error decision, err = tool.NormalizePermissionDecision(decision) @@ -373,6 +379,19 @@ func validate(policy platform.ToolPolicy) error { } } +func validateRuntimeIdentity(policy platform.ToolPolicy) error { + if strings.TrimSpace(policy.TenantID) == "" { + return fmt.Errorf("tenant_id is required") + } + if strings.TrimSpace(policy.AppID) == "" { + return fmt.Errorf("app_id is required") + } + if strings.TrimSpace(policy.PolicyID) == "" { + return fmt.Errorf("policy_id is required") + } + return nil +} + func isHighRisk(policy platform.ToolPolicy, req *tool.PermissionRequest, name string) bool { if contains(normalizedList(policy.HighRiskTools), name) { return true @@ -481,6 +500,24 @@ func (p *Policy) ApprovalSummary( // Validate checks that the summary is safe to expose outside the tool runtime. func (s ApprovalSummary) Validate() error { + if strings.TrimSpace(s.TenantID) == "" { + return fmt.Errorf("tenant_id is required") + } + if strings.TrimSpace(s.AppID) == "" { + return fmt.Errorf("app_id is required") + } + if strings.TrimSpace(s.PolicyID) == "" { + return fmt.Errorf("policy_id is required") + } + if err := platformSafeText("tenant_id", s.TenantID); err != nil { + return err + } + if err := platformSafeText("app_id", s.AppID); err != nil { + return err + } + if err := platformSafeText("policy_id", s.PolicyID); err != nil { + return err + } if strings.TrimSpace(s.ToolName) == "" { return fmt.Errorf("tool_name is required") } @@ -536,12 +573,12 @@ func (s ApprovalSummary) Validate() error { return nil } -func (p *Policy) writeAudit(ctx context.Context, summary ApprovalSummary) { +func (p *Policy) writeAudit(ctx context.Context, summary ApprovalSummary) error { if p.audit == nil { - return + return nil } detailRef := summary.DetailRef() - _ = p.audit.WriteAudit(ctx, platform.AuditRecord{ + if err := p.audit.WriteAudit(ctx, platform.AuditRecord{ AuditID: platform.AuditID(summary.TenantID, summary.AppID, summary.ToolName, summary.ToolCallID, string(summary.Decision), detailRef), TenantID: summary.TenantID, AppID: summary.AppID, @@ -551,7 +588,10 @@ func (p *Policy) writeAudit(ctx context.Context, summary ApprovalSummary) { RedactedDetailRef: detailRef, RedactionVersion: summary.RedactionVersion, CreatedAt: summary.CreatedAt, - }) + }); err != nil { + return fmt.Errorf("write tool policy audit: %w", err) + } + return nil } // DetailRef returns compact non-secret detail that can be stored in audit logs. diff --git a/platform/toolpolicy/policy_test.go b/platform/toolpolicy/policy_test.go index 1b899e0940..5e4818f57d 100644 --- a/platform/toolpolicy/policy_test.go +++ b/platform/toolpolicy/policy_test.go @@ -10,6 +10,7 @@ package toolpolicy import ( "context" + "errors" "strings" "testing" "time" @@ -190,6 +191,9 @@ func TestPolicyBuildsApprovalSummaryWithoutRawArguments(t *testing.T) { func TestApprovalSummaryValidationRejectsUnsafeOrInconsistentFields(t *testing.T) { now := time.Unix(300, 0) valid := ApprovalSummary{ + TenantID: "tenant", + AppID: "app", + PolicyID: "policy", ToolName: "workspace_write", ToolCallID: "call-1", Decision: tool.PermissionActionAsk, @@ -362,12 +366,63 @@ func TestPolicyDeniesDestructiveMetadata(t *testing.T) { } func TestPolicyRejectsInvalidDangerousAction(t *testing.T) { - _, err := New(platform.ToolPolicy{DangerousToolAction: platform.DangerousToolAction("bad")}) + _, err := New(defaultPolicy(platform.ToolPolicy{ + DangerousToolAction: platform.DangerousToolAction("bad"), + })) if err == nil { t.Fatalf("expected invalid action to fail") } } +func TestPolicyRejectsMissingRuntimeIdentity(t *testing.T) { + for _, tc := range []struct { + name string + policy platform.ToolPolicy + want string + }{ + { + name: "tenant", + policy: platform.ToolPolicy{AppID: "app", PolicyID: "policy"}, + want: "tenant_id", + }, + { + name: "app", + policy: platform.ToolPolicy{TenantID: "tenant", PolicyID: "policy"}, + want: "app_id", + }, + { + name: "policy", + policy: platform.ToolPolicy{TenantID: "tenant", AppID: "app"}, + want: "policy_id", + }, + } { + t.Run(tc.name, func(t *testing.T) { + if _, err := New(tc.policy); err == nil || !strings.Contains(err.Error(), tc.want) { + t.Fatalf("expected %s validation error, got %v", tc.want, err) + } + }) + } +} + +func TestPolicyReturnsAuditSinkErrors(t *testing.T) { + p := newPolicy( + t, + platform.ToolPolicy{ + DangerousToolAction: platform.DangerousToolActionAllowWithAudit, + HighRiskTools: []string{"http_post"}, + }, + WithAuditSink(failingAuditSink{}), + ) + + _, err := p.CheckToolPermission( + context.Background(), + request("http_post", tool.ToolMetadata{}, []byte(`{"url":"https://example.com"}`)), + ) + if err == nil || !strings.Contains(err.Error(), "write tool policy audit") { + t.Fatalf("expected audit sink error, got %v", err) + } +} + func TestApprovalOptionsMapPolicy(t *testing.T) { opts, err := ApprovalOptions(platform.ToolPolicy{ ToolWhitelist: []string{"search", "shell"}, @@ -470,9 +525,9 @@ func TestPolicyRegisterDoesNotTreatUnknownMetadataAsHighRisk(t *testing.T) { } func TestReviewerMapsPolicyDecisionToApprovalDecision(t *testing.T) { - reviewer, err := NewReviewer(platform.ToolPolicy{ + reviewer, err := NewReviewer(defaultPolicy(platform.ToolPolicy{ ToolWhitelist: []string{"search"}, - }) + })) if err != nil { t.Fatalf("NewReviewer: %v", err) } @@ -489,10 +544,10 @@ func TestReviewerMapsPolicyDecisionToApprovalDecision(t *testing.T) { } func TestReviewerApprovesAskDecisionForApprovalPluginFlow(t *testing.T) { - reviewer, err := NewReviewer(platform.ToolPolicy{ + reviewer, err := NewReviewer(defaultPolicy(platform.ToolPolicy{ DangerousToolAction: platform.DangerousToolActionAsk, HighRiskTools: []string{"workspace_write"}, - }) + })) if err != nil { t.Fatalf("NewReviewer: %v", err) } @@ -509,11 +564,11 @@ func TestReviewerApprovesAskDecisionForApprovalPluginFlow(t *testing.T) { } func TestApprovalOptionsWithReviewerAllowsWhitelistedHighRiskAsk(t *testing.T) { - policy := platform.ToolPolicy{ + policy := defaultPolicy(platform.ToolPolicy{ ToolWhitelist: []string{"workspace_write"}, DangerousToolAction: platform.DangerousToolActionAsk, HighRiskTools: []string{"workspace_write"}, - } + }) reviewer, err := NewReviewer(policy) if err != nil { t.Fatalf("NewReviewer: %v", err) @@ -539,6 +594,7 @@ func TestApprovalOptionsWithReviewerAllowsWhitelistedHighRiskAsk(t *testing.T) { func newPolicy(t *testing.T, policy platform.ToolPolicy, opts ...Option) *Policy { t.Helper() + policy = defaultPolicy(policy) p, err := New(policy, opts...) if err != nil { t.Fatalf("New: %v", err) @@ -546,6 +602,19 @@ func newPolicy(t *testing.T, policy platform.ToolPolicy, opts ...Option) *Policy return p } +func defaultPolicy(policy platform.ToolPolicy) platform.ToolPolicy { + if strings.TrimSpace(policy.TenantID) == "" { + policy.TenantID = "tenant" + } + if strings.TrimSpace(policy.AppID) == "" { + policy.AppID = "app" + } + if strings.TrimSpace(policy.PolicyID) == "" { + policy.PolicyID = "policy" + } + return policy +} + func request(name string, metadata tool.ToolMetadata, args ...[]byte) *tool.PermissionRequest { var payload []byte if len(args) > 0 { @@ -564,3 +633,9 @@ type allowReviewer struct{} func (allowReviewer) Review(context.Context, *review.Request) (*review.Decision, error) { return &review.Decision{Approved: true}, nil } + +type failingAuditSink struct{} + +func (failingAuditSink) WriteAudit(context.Context, platform.AuditRecord) error { + return errors.New("audit unavailable") +} From b8b802e5e7c5b5bf18bd6e9b26bf55012799ed70 Mon Sep 17 00:00:00 2001 From: xnlemon Date: Thu, 9 Jul 2026 11:23:11 +0800 Subject: [PATCH 36/95] platform: harden identity and idempotency contracts --- platform/idempotency.go | 51 +++++-- platform/identity.go | 114 ++++++++++++---- platform/redaction.go | 4 +- platform/types_test.go | 291 +++++++++++++++++++++++++++++++++++++++- platform/validation.go | 19 ++- 5 files changed, 428 insertions(+), 51 deletions(-) diff --git a/platform/idempotency.go b/platform/idempotency.go index 8e9333ae58..2f27f79cc4 100644 --- a/platform/idempotency.go +++ b/platform/idempotency.go @@ -10,6 +10,8 @@ package platform import ( "context" + "fmt" + "strings" "sync" "time" ) @@ -49,16 +51,14 @@ func (s *InMemoryIdempotencyStore) Start( if err := ctx.Err(); err != nil { return IdempotencyRecord{}, false, err } - key := record.IdempotencyKey - if key == "" { - key = IdempotencyKey( - record.TenantID, - record.Channel, - record.AccountID, - record.PlatformMessageID, - ) - record.IdempotencyKey = key + key, err := canonicalIdempotencyKey(record) + if err != nil { + return IdempotencyRecord{}, false, err + } + if record.IdempotencyKey != "" && record.IdempotencyKey != key { + return IdempotencyRecord{}, false, fmt.Errorf("idempotency_key does not match canonical key") } + record.IdempotencyKey = key s.mu.Lock() defer s.mu.Unlock() if existing, ok := s.records[key]; ok { @@ -78,7 +78,7 @@ func (s *InMemoryIdempotencyStore) Complete( key string, resultRef string, ) (IdempotencyRecord, error) { - return s.update(ctx, key, IdempotencyStatusCompleted, resultRef) + return s.update(ctx, key, IdempotencyStatusProcessing, IdempotencyStatusCompleted, resultRef) } // MarkReplyFailed marks a completed record as needing outbound retry. @@ -87,7 +87,7 @@ func (s *InMemoryIdempotencyStore) MarkReplyFailed( key string, resultRef string, ) (IdempotencyRecord, error) { - return s.update(ctx, key, IdempotencyStatusReplyFailed, resultRef) + return s.update(ctx, key, IdempotencyStatusCompleted, IdempotencyStatusReplyFailed, resultRef) } // Get returns the record for key. @@ -107,6 +107,7 @@ func (s *InMemoryIdempotencyStore) Get( func (s *InMemoryIdempotencyStore) update( ctx context.Context, key string, + from IdempotencyStatus, status IdempotencyStatus, resultRef string, ) (IdempotencyRecord, error) { @@ -119,9 +120,37 @@ func (s *InMemoryIdempotencyStore) update( if !ok { return IdempotencyRecord{}, ErrIdempotencyRecordNotFound } + if record.Status != from { + return IdempotencyRecord{}, fmt.Errorf( + "invalid idempotency transition from %q to %q", + record.Status, + status, + ) + } record.Status = status record.ResultRef = resultRef record.UpdatedAt = s.now() s.records[key] = record return record, nil } + +func canonicalIdempotencyKey(record IdempotencyRecord) (string, error) { + if strings.TrimSpace(record.TenantID) == "" { + return "", ErrTenantIDRequired + } + if strings.TrimSpace(record.Channel) == "" { + return "", ErrChannelRequired + } + if strings.TrimSpace(record.AccountID) == "" { + return "", ErrAccountIDRequired + } + if strings.TrimSpace(record.PlatformMessageID) == "" { + return "", ErrPlatformMessageIDRequired + } + return IdempotencyKey( + record.TenantID, + record.Channel, + record.AccountID, + record.PlatformMessageID, + ), nil +} diff --git a/platform/identity.go b/platform/identity.go index 9ae3916a83..6195d99ec2 100644 --- a/platform/identity.go +++ b/platform/identity.go @@ -12,10 +12,14 @@ import ( "crypto/sha256" "encoding/hex" "fmt" - "net/url" "strings" ) +type stableIDPart struct { + name string + value string +} + // InternalUserID returns a stable tenant-scoped user identifier. func InternalUserID(tenantID, channel, externalUserID string) string { return "usr_" + shortHash(tenantID, channel, externalUserID) @@ -33,12 +37,13 @@ func AuditID(parts ...string) string { // IdempotencyKey returns the canonical duplicate-delivery key. func IdempotencyKey(tenantID, channel, accountID, platformMessageID string) string { - return strings.Join([]string{ - "tenant", escapeKeyPart(tenantID), - "channel", escapeKeyPart(channel), - "account", escapeKeyPart(accountID), - "message", escapeKeyPart(platformMessageID), - }, ":") + return stableID( + "idem", + stableIDPart{"tenant", tenantID}, + stableIDPart{"channel", channel}, + stableIDPart{"account", accountID}, + stableIDPart{"message", platformMessageID}, + ) } // SessionIDForInbound returns the stable session id for one inbound message. @@ -46,15 +51,32 @@ func SessionIDForInbound(msg InboundMessage) (string, error) { if err := msg.Validate(); err != nil { return "", err } - return SessionID( - msg.TenantID, - msg.AppID, - msg.Channel, + parts := []stableIDPart{ + {"tenant", msg.TenantID}, + {"app", msg.AppID}, + {"binding", msg.BindingID}, + {"channel", msg.Channel}, + {"account", msg.ChannelAccountID}, + } + if msg.MessageType == MessageTypeEvent { + parts = append(parts, + stableIDPart{"message_type", string(msg.MessageType)}, + stableIDPart{"event_type", msg.RawEventType}, + stableIDPart{"message", msg.PlatformMessageID}, + ) + return stableID("ses", parts...), nil + } + conversationParts, err := sessionConversationParts( msg.ConversationType, msg.ExternalUserID, msg.ExternalGroupID, msg.ThreadID, ) + if err != nil { + return "", err + } + parts = append(parts, conversationParts...) + return stableID("ses", parts...), nil } // SessionID returns the stable tenant/app/channel-scoped session id. @@ -76,36 +98,63 @@ func SessionID( if strings.TrimSpace(channel) == "" { return "", ErrChannelRequired } - prefix := fmt.Sprintf( - "tenant:%s:app:%s:channel:%s", - escapeKeyPart(tenantID), - escapeKeyPart(appID), - escapeKeyPart(channel), + parts := []stableIDPart{ + {"tenant", tenantID}, + {"app", appID}, + {"channel", channel}, + } + conversationParts, err := sessionConversationParts( + conversationType, + externalUserID, + externalGroupID, + threadID, ) + if err != nil { + return "", err + } + parts = append(parts, conversationParts...) + return stableID("ses", parts...), nil +} + +func sessionConversationParts( + conversationType ConversationType, + externalUserID string, + externalGroupID string, + threadID string, +) ([]stableIDPart, error) { switch conversationType { case ConversationTypeDM: if strings.TrimSpace(externalUserID) == "" { - return "", ErrExternalUserIDRequired + return nil, ErrExternalUserIDRequired } - return prefix + ":dm:" + escapeKeyPart(externalUserID), nil + return []stableIDPart{ + {"conversation_type", string(ConversationTypeDM)}, + {"user", externalUserID}, + }, nil case ConversationTypeGroup: if strings.TrimSpace(externalGroupID) == "" { - return "", ErrExternalGroupIDRequired + return nil, ErrExternalGroupIDRequired } - return prefix + ":group:" + escapeKeyPart(externalGroupID), nil + return []stableIDPart{ + {"conversation_type", string(ConversationTypeGroup)}, + {"group", externalGroupID}, + }, nil case ConversationTypeThread: if strings.TrimSpace(externalGroupID) == "" { - return "", ErrExternalGroupIDRequired + return nil, ErrExternalGroupIDRequired } if strings.TrimSpace(threadID) == "" { - return "", fmt.Errorf("thread_id is required") + return nil, fmt.Errorf("thread_id is required") } - return prefix + ":group:" + escapeKeyPart(externalGroupID) + - ":thread:" + escapeKeyPart(threadID), nil + return []stableIDPart{ + {"conversation_type", string(ConversationTypeThread)}, + {"group", externalGroupID}, + {"thread", threadID}, + }, nil case "": - return "", ErrConversationTypeRequired + return nil, ErrConversationTypeRequired default: - return "", ErrInvalidConversationType + return nil, ErrInvalidConversationType } } @@ -114,6 +163,15 @@ func shortHash(parts ...string) string { return hex.EncodeToString(sum[:])[:24] } -func escapeKeyPart(value string) string { - return url.PathEscape(strings.TrimSpace(value)) +func stableID(prefix string, parts ...stableIDPart) string { + hash := sha256.New() + writeStablePart(hash, "prefix", prefix) + for _, part := range parts { + writeStablePart(hash, strings.TrimSpace(part.name), strings.TrimSpace(part.value)) + } + return prefix + "_" + hex.EncodeToString(hash.Sum(nil))[:32] +} + +func writeStablePart(hash interface{ Write([]byte) (int, error) }, name, value string) { + fmt.Fprintf(hash, "%d:%s=%d:%s;", len(name), name, len(value), value) } diff --git a/platform/redaction.go b/platform/redaction.go index ccbc327017..4e87412614 100644 --- a/platform/redaction.go +++ b/platform/redaction.go @@ -16,11 +16,13 @@ import ( var defaultRedactionPatterns = []*regexp.Regexp{ regexp.MustCompile(`(?i)(Authorization:\s*Basic\s+)[A-Za-z0-9._~+/\-]+=*`), regexp.MustCompile(`(?i)(Bearer\s+)[A-Za-z0-9._~+/\-]+=*`), + regexp.MustCompile(`(?im)(authorization\s*:\s*)[^\r\n]+`), + regexp.MustCompile(`(?im)(authorization\s*=\s*)[^\r\n]+`), regexp.MustCompile(`(?i)(api[_-]?key|token|secret|password|passwd|authorization|cookie)=([^&\s]+)`), regexp.MustCompile(`(?i)(api[_-]?key|token|secret|password|passwd|authorization|cookie):\s*([^,\s]+)`), regexp.MustCompile(`(?i)("(?:api[_-]?key|token|secret|password|passwd|authorization|cookie)"\s*:\s*")([^"]+)(")`), regexp.MustCompile(`(?i)(sk-[A-Za-z0-9._~+/\-]{8,})`), - regexp.MustCompile(`(?i)(://[^:\s/]+:)([^@\s]+)(@)`), + regexp.MustCompile(`(?i)(://[^\s/?#@]+:)([^\s/?#]+)(@[^\s/?#@]+)`), regexp.MustCompile(`(?s)-----BEGIN [A-Z ]*PRIVATE KEY-----.*?-----END [A-Z ]*PRIVATE KEY-----`), } diff --git a/platform/types_test.go b/platform/types_test.go index 2aee01424f..557a6e96a3 100644 --- a/platform/types_test.go +++ b/platform/types_test.go @@ -18,6 +18,7 @@ import ( func TestSessionIDForInboundIsTenantScoped(t *testing.T) { base := InboundMessage{ AppID: "support", + BindingID: "telegram-bot-1", Channel: "telegram", ChannelAccountID: "bot-1", PlatformMessageID: "msg-1", @@ -41,8 +42,8 @@ func TestSessionIDForInboundIsTenantScoped(t *testing.T) { if sessionA == sessionB { t.Fatalf("sessions should differ across tenants: %q", sessionA) } - if !strings.Contains(sessionA, "tenant:tenant-a:app:support:channel:telegram:dm:same-user") { - t.Fatalf("unexpected dm session id: %q", sessionA) + if !strings.HasPrefix(sessionA, "ses_") || strings.Contains(sessionA, ":") { + t.Fatalf("session id should be opaque and delimiter-safe, got %q", sessionA) } } @@ -51,16 +52,51 @@ func TestSessionIDForInboundSupportsGroupAndThread(t *testing.T) { if err != nil { t.Fatalf("group session: %v", err) } - if groupID != "tenant:tenant:app:app:channel:wecom:group:room%201" { - t.Fatalf("unexpected group id: %q", groupID) + groupIDAgain, err := SessionID("tenant", "app", "wecom", ConversationTypeGroup, "user", "room 1", "") + if err != nil { + t.Fatalf("group session again: %v", err) + } + if groupID != groupIDAgain { + t.Fatalf("expected stable group id, got %q and %q", groupID, groupIDAgain) + } + if !strings.HasPrefix(groupID, "ses_") || strings.Contains(groupID, ":") { + t.Fatalf("group session id should be opaque and delimiter-safe, got %q", groupID) } threadID, err := SessionID("tenant", "app", "telegram", ConversationTypeThread, "user", "chat", "topic/7") if err != nil { t.Fatalf("thread session: %v", err) } - if threadID != "tenant:tenant:app:app:channel:telegram:group:chat:thread:topic%2F7" { - t.Fatalf("unexpected thread id: %q", threadID) + if !strings.HasPrefix(threadID, "ses_") || strings.Contains(threadID, ":") { + t.Fatalf("thread session id should be opaque and delimiter-safe, got %q", threadID) + } + if groupID == threadID { + t.Fatalf("group and thread sessions should differ: %q", groupID) + } +} + +func TestStableIDsAreDelimiterSafe(t *testing.T) { + sessionA, err := SessionID("tenant", "app", "chat:dm:user", ConversationTypeDM, "leaf", "", "") + if err != nil { + t.Fatalf("session A: %v", err) + } + sessionB, err := SessionID("tenant", "app", "chat", ConversationTypeDM, "user:dm:leaf", "", "") + if err != nil { + t.Fatalf("session B: %v", err) + } + if sessionA == sessionB { + t.Fatalf("delimiter-bearing session parts should not collide: %q", sessionA) + } + + keyA := IdempotencyKey("tenant", "chat:account:bot", "primary", "msg") + keyB := IdempotencyKey("tenant", "chat", "bot:account:primary", "msg") + if keyA == keyB { + t.Fatalf("delimiter-bearing idempotency parts should not collide: %q", keyA) + } + for _, key := range []string{keyA, keyB} { + if !strings.HasPrefix(key, "idem_") || strings.Contains(key, ":") { + t.Fatalf("idempotency key should be opaque and delimiter-safe, got %q", key) + } } } @@ -68,6 +104,7 @@ func TestValidateInboundRequiresGroupForGroupConversation(t *testing.T) { msg := InboundMessage{ TenantID: "tenant", AppID: "app", + BindingID: "binding", Channel: "telegram", ChannelAccountID: "bot", PlatformMessageID: "msg", @@ -79,6 +116,76 @@ func TestValidateInboundRequiresGroupForGroupConversation(t *testing.T) { } } +func TestValidateInboundRequiresBindingID(t *testing.T) { + msg := InboundMessage{ + TenantID: "tenant", + AppID: "app", + Channel: "telegram", + ChannelAccountID: "bot", + PlatformMessageID: "msg", + ExternalUserID: "user", + ConversationType: ConversationTypeDM, + MessageType: MessageTypeText, + } + if err := msg.Validate(); !errors.Is(err, ErrBindingIDRequired) { + t.Fatalf("expected ErrBindingIDRequired, got %v", err) + } +} + +func TestValidateInboundEventDoesNotRequireConversationIdentity(t *testing.T) { + msg := InboundMessage{ + TenantID: "tenant", + AppID: "app", + BindingID: "binding", + Channel: "telegram", + ChannelAccountID: "bot", + PlatformMessageID: "event-1", + MessageType: MessageTypeEvent, + RawEventType: "app_mention", + } + if err := msg.Validate(); err != nil { + t.Fatalf("event should not require user or conversation identity, got %v", err) + } +} + +func TestSessionIDForInboundIncludesBindingAndAccountScope(t *testing.T) { + base := InboundMessage{ + TenantID: "tenant", + AppID: "app", + BindingID: "binding-a", + Channel: "telegram", + ChannelAccountID: "bot-a", + PlatformMessageID: "msg", + ExternalUserID: "same-user", + ConversationType: ConversationTypeDM, + MessageType: MessageTypeText, + } + bindingA, err := SessionIDForInbound(base) + if err != nil { + t.Fatalf("binding A: %v", err) + } + + bindingVariant := base + bindingVariant.BindingID = "binding-b" + bindingB, err := SessionIDForInbound(bindingVariant) + if err != nil { + t.Fatalf("binding B: %v", err) + } + if bindingA == bindingB { + t.Fatalf("sessions should differ across bindings: %q", bindingA) + } + + accountVariant := base + accountVariant.ChannelAccountID = "bot-b" + accountB, err := SessionIDForInbound(accountVariant) + if err != nil { + t.Fatalf("account B: %v", err) + } + if bindingA == accountB { + t.Fatalf("sessions should differ across channel accounts: %q", bindingA) + } +} + func TestInternalUserIDIsStableAndTenantScoped(t *testing.T) { a1 := InternalUserID("tenant-a", "telegram", "42") a2 := InternalUserID("tenant-a", "telegram", "42") @@ -143,6 +250,127 @@ func TestIdempotencyStoreDoesNotRestartCompletedMessage(t *testing.T) { } } +func TestIdempotencyStartRejectsMissingKeyFields(t *testing.T) { + tests := []struct { + name string + record IdempotencyRecord + want error + }{ + { + name: "tenant", + record: IdempotencyRecord{ + Channel: "telegram", + AccountID: "bot", + PlatformMessageID: "msg", + }, + want: ErrTenantIDRequired, + }, + { + name: "channel", + record: IdempotencyRecord{ + TenantID: "tenant", + AccountID: "bot", + PlatformMessageID: "msg", + }, + want: ErrChannelRequired, + }, + { + name: "account", + record: IdempotencyRecord{ + TenantID: "tenant", + Channel: "telegram", + PlatformMessageID: "msg", + }, + want: ErrAccountIDRequired, + }, + { + name: "message", + record: IdempotencyRecord{ + TenantID: "tenant", + Channel: "telegram", + AccountID: "bot", + }, + want: ErrPlatformMessageIDRequired, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + store := NewInMemoryIdempotencyStore() + _, started, err := store.Start(context.Background(), tt.record) + if !errors.Is(err, tt.want) { + t.Fatalf("expected %v, got %v", tt.want, err) + } + if started { + t.Fatalf("invalid record should not start") + } + }) + } +} + +func TestIdempotencyStartRejectsMismatchedCallerKey(t *testing.T) { + store := NewInMemoryIdempotencyStore() + record := IdempotencyRecord{ + TenantID: "tenant", + Channel: "telegram", + AccountID: "bot-a", + PlatformMessageID: "msg", + IdempotencyKey: IdempotencyKey("tenant", "telegram", "bot-b", "msg"), + } + _, started, err := store.Start(context.Background(), record) + if err == nil { + t.Fatalf("expected mismatched caller-supplied idempotency key to fail") + } + if started { + t.Fatalf("mismatched caller-supplied key should not start") + } +} + +func TestIdempotencyStoreEnforcesStateTransitions(t *testing.T) { + ctx := context.Background() + processingStore := NewInMemoryIdempotencyStore() + processing, started, err := processingStore.Start(ctx, IdempotencyRecord{ + TenantID: "tenant", + Channel: "telegram", + AccountID: "bot", + PlatformMessageID: "msg-processing", + }) + if err != nil { + t.Fatalf("start processing: %v", err) + } + if !started { + t.Fatalf("processing record should start") + } + if _, err := processingStore.MarkReplyFailed(ctx, processing.IdempotencyKey, "outbound-1"); err == nil { + t.Fatalf("reply failure should only be allowed after completion") + } + + completedStore := NewInMemoryIdempotencyStore() + completed, started, err := completedStore.Start(ctx, IdempotencyRecord{ + TenantID: "tenant", + Channel: "telegram", + AccountID: "bot", + PlatformMessageID: "msg-completed", + }) + if err != nil { + t.Fatalf("start completed: %v", err) + } + if !started { + t.Fatalf("completed record should start") + } + if _, err := completedStore.Complete(ctx, completed.IdempotencyKey, "outbound-1"); err != nil { + t.Fatalf("complete: %v", err) + } + if _, err := completedStore.Complete(ctx, completed.IdempotencyKey, "outbound-2"); err == nil { + t.Fatalf("completed record should not be completed again") + } + if _, err := completedStore.MarkReplyFailed(ctx, completed.IdempotencyKey, "outbound-1"); err != nil { + t.Fatalf("mark reply failed from completed: %v", err) + } + if _, err := completedStore.Complete(ctx, completed.IdempotencyKey, "outbound-3"); err == nil { + t.Fatalf("reply-failed record should not transition back to completed") + } +} + func TestBindingRejectsInlineSecrets(t *testing.T) { binding := ChannelBinding{ TenantID: "tenant", @@ -188,6 +416,21 @@ func TestBindingRejectsTelegramBotToken(t *testing.T) { } } +func TestBindingRejectsURLUserinfoWithMultipleAtSigns(t *testing.T) { + binding := ChannelBinding{ + TenantID: "tenant", + AppID: "app", + BindingID: "binding", + Channel: "telegram", + AccountID: "bot", + WebhookPath: "/channels/telegram/binding/callback", + SecretRef: "postgres://svc@example.com:password@db/prod", + } + if err := binding.Validate(); !errors.Is(err, ErrInlineSecretRejected) { + t.Fatalf("expected inline URL credential rejection, got %v", err) + } +} + func TestBindingAllowsURISecretReferences(t *testing.T) { binding := ChannelBinding{ TenantID: "tenant", @@ -260,6 +503,42 @@ func TestRedactorMasksSecrets(t *testing.T) { } } +func TestRedactorMasksNonBearerAuthorizationCredentials(t *testing.T) { + redactor, err := NewRedactor() + if err != nil { + t.Fatalf("NewRedactor: %v", err) + } + input := "Authorization: Token top-secret\nAuthorization=Digest username=\"bob\", response=\"abc123\"\n" + got := redactor.Redact(input) + for _, leaked := range []string{"top-secret", "username=\"bob\"", "abc123"} { + if strings.Contains(got, leaked) { + t.Fatalf("redacted output leaked %q: %q", leaked, got) + } + } + if !strings.Contains(got, "Authorization: ****") { + t.Fatalf("expected header authorization mask, got %q", got) + } + if !strings.Contains(got, "Authorization=****") { + t.Fatalf("expected key-value authorization mask, got %q", got) + } +} + +func TestRedactorMasksURLUserinfoWithMultipleAtSigns(t *testing.T) { + redactor, err := NewRedactor() + if err != nil { + t.Fatalf("NewRedactor: %v", err) + } + input := "db=postgres://user:pa@ss@word@example.com/db" + got := redactor.Redact(input) + if strings.Contains(got, "pa@ss@word") || + strings.Contains(got, "ss@word@example.com") { + t.Fatalf("redacted output leaked URL password fragments: %q", got) + } + if !strings.Contains(got, "postgres://user:****@example.com/db") { + t.Fatalf("expected URL userinfo password mask, got %q", got) + } +} + func TestAuditSinkStoresSnapshot(t *testing.T) { sink := NewInMemoryAuditSink() record := AuditRecord{ diff --git a/platform/validation.go b/platform/validation.go index 662ee6c54f..a5452a1c4f 100644 --- a/platform/validation.go +++ b/platform/validation.go @@ -11,6 +11,7 @@ package platform import ( "fmt" "math" + "net/url" "strconv" "strings" ) @@ -119,6 +120,9 @@ func (m InboundMessage) Validate() error { if strings.TrimSpace(m.AppID) == "" { return ErrAppIDRequired } + if strings.TrimSpace(m.BindingID) == "" { + return ErrBindingIDRequired + } if strings.TrimSpace(m.Channel) == "" { return ErrChannelRequired } @@ -128,6 +132,9 @@ func (m InboundMessage) Validate() error { if strings.TrimSpace(m.PlatformMessageID) == "" { return ErrPlatformMessageIDRequired } + if m.MessageType == MessageTypeEvent { + return nil + } if strings.TrimSpace(m.ExternalUserID) == "" { return ErrExternalUserIDRequired } @@ -289,13 +296,15 @@ func validateSecretReference(field, value string) error { } func hasInlineURLCredential(value string) bool { - scheme := strings.Index(value, "://") - at := strings.Index(value, "@") - if scheme < 0 || at < 0 || at < scheme { + parsed, err := url.Parse(value) + if err != nil || parsed.Scheme == "" || parsed.Host == "" || parsed.User == nil { return false } - credential := value[scheme+3 : at] - return strings.Contains(credential, ":") + if parsed.User.Username() != "" { + return true + } + _, hasPassword := parsed.User.Password() + return hasPassword } func looksLikeRawSecret(value string) bool { From db7f629415b01a318d630ada59de67fe99c5385a Mon Sep 17 00:00:00 2001 From: xnlemon Date: Thu, 9 Jul 2026 11:41:52 +0800 Subject: [PATCH 37/95] platform: tighten routing identity contracts --- platform/idempotency.go | 17 +-- platform/identity.go | 62 +++++--- platform/redaction.go | 28 +++- platform/types_test.go | 327 +++++++++++++++++++++++++++++++++++++++- platform/validation.go | 121 +++++++++------ 5 files changed, 465 insertions(+), 90 deletions(-) diff --git a/platform/idempotency.go b/platform/idempotency.go index 2f27f79cc4..005b2aae95 100644 --- a/platform/idempotency.go +++ b/platform/idempotency.go @@ -11,7 +11,6 @@ package platform import ( "context" "fmt" - "strings" "sync" "time" ) @@ -135,17 +134,17 @@ func (s *InMemoryIdempotencyStore) update( } func canonicalIdempotencyKey(record IdempotencyRecord) (string, error) { - if strings.TrimSpace(record.TenantID) == "" { - return "", ErrTenantIDRequired + if err := validateRoutingIdentifier("tenant_id", record.TenantID, ErrTenantIDRequired); err != nil { + return "", err } - if strings.TrimSpace(record.Channel) == "" { - return "", ErrChannelRequired + if err := validateRoutingIdentifier("channel", record.Channel, ErrChannelRequired); err != nil { + return "", err } - if strings.TrimSpace(record.AccountID) == "" { - return "", ErrAccountIDRequired + if err := validateRoutingIdentifier("account_id", record.AccountID, ErrAccountIDRequired); err != nil { + return "", err } - if strings.TrimSpace(record.PlatformMessageID) == "" { - return "", ErrPlatformMessageIDRequired + if err := validateRoutingIdentifier("platform_message_id", record.PlatformMessageID, ErrPlatformMessageIDRequired); err != nil { + return "", err } return IdempotencyKey( record.TenantID, diff --git a/platform/identity.go b/platform/identity.go index 6195d99ec2..cf92feb558 100644 --- a/platform/identity.go +++ b/platform/identity.go @@ -12,7 +12,6 @@ import ( "crypto/sha256" "encoding/hex" "fmt" - "strings" ) type stableIDPart struct { @@ -22,12 +21,22 @@ type stableIDPart struct { // InternalUserID returns a stable tenant-scoped user identifier. func InternalUserID(tenantID, channel, externalUserID string) string { - return "usr_" + shortHash(tenantID, channel, externalUserID) + return stableID( + "usr", + stableIDPart{"tenant", tenantID}, + stableIDPart{"channel", channel}, + stableIDPart{"user", externalUserID}, + ) } // UserIDHash returns a low-sensitivity hash for logs and trace attributes. func UserIDHash(tenantID, channel, userID string) string { - return "user_hash_" + shortHash(tenantID, channel, userID) + return stableID( + "user_hash", + stableIDPart{"tenant", tenantID}, + stableIDPart{"channel", channel}, + stableIDPart{"user", userID}, + ) } // AuditID returns a stable audit identifier for one audit event boundary. @@ -79,29 +88,39 @@ func SessionIDForInbound(msg InboundMessage) (string, error) { return stableID("ses", parts...), nil } -// SessionID returns the stable tenant/app/channel-scoped session id. +// SessionID returns the stable tenant/app/binding/channel/account-scoped session id. func SessionID( tenantID string, appID string, + bindingID string, channel string, + channelAccountID string, conversationType ConversationType, externalUserID string, externalGroupID string, threadID string, ) (string, error) { - if strings.TrimSpace(tenantID) == "" { - return "", ErrTenantIDRequired + if err := validateRoutingIdentifier("tenant_id", tenantID, ErrTenantIDRequired); err != nil { + return "", err } - if strings.TrimSpace(appID) == "" { - return "", ErrAppIDRequired + if err := validateRoutingIdentifier("app_id", appID, ErrAppIDRequired); err != nil { + return "", err + } + if err := validateRoutingIdentifier("binding_id", bindingID, ErrBindingIDRequired); err != nil { + return "", err + } + if err := validateRoutingIdentifier("channel", channel, ErrChannelRequired); err != nil { + return "", err } - if strings.TrimSpace(channel) == "" { - return "", ErrChannelRequired + if err := validateRoutingIdentifier("channel_account_id", channelAccountID, ErrAccountIDRequired); err != nil { + return "", err } parts := []stableIDPart{ {"tenant", tenantID}, {"app", appID}, + {"binding", bindingID}, {"channel", channel}, + {"account", channelAccountID}, } conversationParts, err := sessionConversationParts( conversationType, @@ -124,27 +143,27 @@ func sessionConversationParts( ) ([]stableIDPart, error) { switch conversationType { case ConversationTypeDM: - if strings.TrimSpace(externalUserID) == "" { - return nil, ErrExternalUserIDRequired + if err := validateRoutingIdentifier("external_user_id", externalUserID, ErrExternalUserIDRequired); err != nil { + return nil, err } return []stableIDPart{ {"conversation_type", string(ConversationTypeDM)}, {"user", externalUserID}, }, nil case ConversationTypeGroup: - if strings.TrimSpace(externalGroupID) == "" { - return nil, ErrExternalGroupIDRequired + if err := validateRoutingIdentifier("external_group_id", externalGroupID, ErrExternalGroupIDRequired); err != nil { + return nil, err } return []stableIDPart{ {"conversation_type", string(ConversationTypeGroup)}, {"group", externalGroupID}, }, nil case ConversationTypeThread: - if strings.TrimSpace(externalGroupID) == "" { - return nil, ErrExternalGroupIDRequired + if err := validateRoutingIdentifier("external_group_id", externalGroupID, ErrExternalGroupIDRequired); err != nil { + return nil, err } - if strings.TrimSpace(threadID) == "" { - return nil, fmt.Errorf("thread_id is required") + if err := validateRoutingIdentifier("thread_id", threadID, fmt.Errorf("thread_id is required")); err != nil { + return nil, err } return []stableIDPart{ {"conversation_type", string(ConversationTypeThread)}, @@ -158,16 +177,11 @@ func sessionConversationParts( } } -func shortHash(parts ...string) string { - sum := sha256.Sum256([]byte(strings.Join(parts, "\x00"))) - return hex.EncodeToString(sum[:])[:24] -} - func stableID(prefix string, parts ...stableIDPart) string { hash := sha256.New() writeStablePart(hash, "prefix", prefix) for _, part := range parts { - writeStablePart(hash, strings.TrimSpace(part.name), strings.TrimSpace(part.value)) + writeStablePart(hash, part.name, part.value) } return prefix + "_" + hex.EncodeToString(hash.Sum(nil))[:32] } diff --git a/platform/redaction.go b/platform/redaction.go index 4e87412614..569bc53eaf 100644 --- a/platform/redaction.go +++ b/platform/redaction.go @@ -22,7 +22,7 @@ var defaultRedactionPatterns = []*regexp.Regexp{ regexp.MustCompile(`(?i)(api[_-]?key|token|secret|password|passwd|authorization|cookie):\s*([^,\s]+)`), regexp.MustCompile(`(?i)("(?:api[_-]?key|token|secret|password|passwd|authorization|cookie)"\s*:\s*")([^"]+)(")`), regexp.MustCompile(`(?i)(sk-[A-Za-z0-9._~+/\-]{8,})`), - regexp.MustCompile(`(?i)(://[^\s/?#@]+:)([^\s/?#]+)(@[^\s/?#@]+)`), + regexp.MustCompile(`(?i)[a-z][a-z0-9+.-]*://[^\s/?#]*@[^\s/?#]+`), regexp.MustCompile(`(?s)-----BEGIN [A-Z ]*PRIVATE KEY-----.*?-----END [A-Z ]*PRIVATE KEY-----`), } @@ -69,12 +69,8 @@ func redactMatch(match string) string { return match[:strings.Index(lower, "bearer ")+7] + "****" } if strings.Contains(match, "://") && strings.Contains(match, "@") { - start := strings.Index(match, "://") - at := strings.LastIndex(match, "@") - credential := match[start+3 : at] - colon := strings.LastIndex(credential, ":") - if colon >= 0 { - return match[:start+3+colon+1] + "****" + match[at:] + if redacted, ok := redactURLUserinfo(match); ok { + return redacted } } if strings.HasPrefix(match, "-----BEGIN ") { @@ -96,3 +92,21 @@ func redactMatch(match string) string { } return match[:4] + "****" + match[len(match)-4:] } + +func redactURLUserinfo(match string) (string, bool) { + scheme := strings.Index(match, "://") + if scheme < 0 { + return match, false + } + authorityStart := scheme + len("://") + authorityEnd := len(match) + if end := strings.IndexAny(match[authorityStart:], "/?# \t\r\n"); end >= 0 { + authorityEnd = authorityStart + end + } + at := strings.LastIndex(match[authorityStart:authorityEnd], "@") + if at <= 0 { + return match, false + } + at += authorityStart + return match[:authorityStart] + "****" + match[at:], true +} diff --git a/platform/types_test.go b/platform/types_test.go index 557a6e96a3..d7323772a8 100644 --- a/platform/types_test.go +++ b/platform/types_test.go @@ -48,11 +48,11 @@ func TestSessionIDForInboundIsTenantScoped(t *testing.T) { } func TestSessionIDForInboundSupportsGroupAndThread(t *testing.T) { - groupID, err := SessionID("tenant", "app", "wecom", ConversationTypeGroup, "user", "room 1", "") + groupID, err := SessionID("tenant", "app", "binding", "wecom", "bot", ConversationTypeGroup, "user", "room 1", "") if err != nil { t.Fatalf("group session: %v", err) } - groupIDAgain, err := SessionID("tenant", "app", "wecom", ConversationTypeGroup, "user", "room 1", "") + groupIDAgain, err := SessionID("tenant", "app", "binding", "wecom", "bot", ConversationTypeGroup, "user", "room 1", "") if err != nil { t.Fatalf("group session again: %v", err) } @@ -63,7 +63,7 @@ func TestSessionIDForInboundSupportsGroupAndThread(t *testing.T) { t.Fatalf("group session id should be opaque and delimiter-safe, got %q", groupID) } - threadID, err := SessionID("tenant", "app", "telegram", ConversationTypeThread, "user", "chat", "topic/7") + threadID, err := SessionID("tenant", "app", "binding", "telegram", "bot", ConversationTypeThread, "user", "chat", "topic/7") if err != nil { t.Fatalf("thread session: %v", err) } @@ -75,12 +75,33 @@ func TestSessionIDForInboundSupportsGroupAndThread(t *testing.T) { } } +func TestSessionIDIncludesBindingAndAccountScope(t *testing.T) { + bindingA, err := SessionID("tenant", "app", "binding-a", "telegram", "bot-a", ConversationTypeDM, "user", "", "") + if err != nil { + t.Fatalf("binding A: %v", err) + } + bindingB, err := SessionID("tenant", "app", "binding-b", "telegram", "bot-a", ConversationTypeDM, "user", "", "") + if err != nil { + t.Fatalf("binding B: %v", err) + } + if bindingA == bindingB { + t.Fatalf("sessions should differ across bindings: %q", bindingA) + } + accountB, err := SessionID("tenant", "app", "binding-a", "telegram", "bot-b", ConversationTypeDM, "user", "", "") + if err != nil { + t.Fatalf("account B: %v", err) + } + if bindingA == accountB { + t.Fatalf("sessions should differ across channel accounts: %q", bindingA) + } +} + func TestStableIDsAreDelimiterSafe(t *testing.T) { - sessionA, err := SessionID("tenant", "app", "chat:dm:user", ConversationTypeDM, "leaf", "", "") + sessionA, err := SessionID("tenant", "app", "binding", "chat:dm:user", "bot", ConversationTypeDM, "leaf", "", "") if err != nil { t.Fatalf("session A: %v", err) } - sessionB, err := SessionID("tenant", "app", "chat", ConversationTypeDM, "user:dm:leaf", "", "") + sessionB, err := SessionID("tenant", "app", "binding", "chat", "bot", ConversationTypeDM, "user:dm:leaf", "", "") if err != nil { t.Fatalf("session B: %v", err) } @@ -110,6 +131,7 @@ func TestValidateInboundRequiresGroupForGroupConversation(t *testing.T) { PlatformMessageID: "msg", ExternalUserID: "user", ConversationType: ConversationTypeGroup, + MessageType: MessageTypeText, } if err := msg.Validate(); !errors.Is(err, ErrExternalGroupIDRequired) { t.Fatalf("expected ErrExternalGroupIDRequired, got %v", err) @@ -148,6 +170,240 @@ func TestValidateInboundEventDoesNotRequireConversationIdentity(t *testing.T) { } } +func TestValidateRejectsNonNormalizedRoutingIdentifiers(t *testing.T) { + badValues := []struct { + name string + value string + }{ + {name: "leading_space", value: " value"}, + {name: "trailing_space", value: "value "}, + {name: "control_nul", value: "value\x00x"}, + {name: "control_newline", value: "value\nx"}, + } + tests := []struct { + name string + validate func(string) error + }{ + { + name: "TenantID", + validate: func(value string) error { + return Tenant{TenantID: value}.Validate() + }, + }, + { + name: "AppID", + validate: func(value string) error { + return AgentApp{TenantID: "tenant", AppID: value}.Validate() + }, + }, + { + name: "BindingID", + validate: func(value string) error { + return ChannelBinding{ + TenantID: "tenant", + AppID: "app", + BindingID: value, + Channel: "telegram", + AccountID: "bot", + WebhookPath: "/channels/telegram/binding/callback", + }.Validate() + }, + }, + { + name: "Channel", + validate: func(value string) error { + return ChannelBinding{ + TenantID: "tenant", + AppID: "app", + BindingID: "binding", + Channel: value, + AccountID: "bot", + WebhookPath: "/channels/telegram/binding/callback", + }.Validate() + }, + }, + { + name: "ChannelAccountID", + validate: func(value string) error { + return InboundMessage{ + TenantID: "tenant", + AppID: "app", + BindingID: "binding", + Channel: "telegram", + ChannelAccountID: value, + PlatformMessageID: "msg", + ExternalUserID: "user", + ConversationType: ConversationTypeDM, + MessageType: MessageTypeText, + }.Validate() + }, + }, + { + name: "PlatformMessageID", + validate: func(value string) error { + return InboundMessage{ + TenantID: "tenant", + AppID: "app", + BindingID: "binding", + Channel: "telegram", + ChannelAccountID: "bot", + PlatformMessageID: value, + ExternalUserID: "user", + ConversationType: ConversationTypeDM, + MessageType: MessageTypeText, + }.Validate() + }, + }, + } + for _, tt := range tests { + for _, bad := range badValues { + t.Run(tt.name+"/"+bad.name, func(t *testing.T) { + if err := tt.validate(bad.value); err == nil { + t.Fatalf("expected %s to reject %q", tt.name, bad.value) + } + }) + } + } +} + +func TestIdempotencyStartRejectsNonNormalizedKeyFields(t *testing.T) { + badValues := []struct { + name string + value string + }{ + {name: "leading_space", value: " value"}, + {name: "trailing_space", value: "value "}, + {name: "control_nul", value: "value\x00x"}, + {name: "control_newline", value: "value\nx"}, + } + tests := []struct { + name string + mutate func(*IdempotencyRecord, string) + }{ + { + name: "TenantID", + mutate: func(record *IdempotencyRecord, value string) { + record.TenantID = value + }, + }, + { + name: "Channel", + mutate: func(record *IdempotencyRecord, value string) { + record.Channel = value + }, + }, + { + name: "AccountID", + mutate: func(record *IdempotencyRecord, value string) { + record.AccountID = value + }, + }, + { + name: "PlatformMessageID", + mutate: func(record *IdempotencyRecord, value string) { + record.PlatformMessageID = value + }, + }, + } + for _, tt := range tests { + for _, bad := range badValues { + t.Run(tt.name+"/"+bad.name, func(t *testing.T) { + record := IdempotencyRecord{ + TenantID: "tenant", + Channel: "telegram", + AccountID: "bot", + PlatformMessageID: "msg", + } + tt.mutate(&record, bad.value) + store := NewInMemoryIdempotencyStore() + _, started, err := store.Start(context.Background(), record) + if err == nil { + t.Fatalf("expected %s to reject %q", tt.name, bad.value) + } + if started { + t.Fatalf("invalid record should not start") + } + }) + } + } +} + +func TestValidateInboundRequiresKnownMessageType(t *testing.T) { + base := InboundMessage{ + TenantID: "tenant", + AppID: "app", + BindingID: "binding", + Channel: "telegram", + ChannelAccountID: "bot", + PlatformMessageID: "msg", + ExternalUserID: "user", + ConversationType: ConversationTypeDM, + } + tests := []struct { + name string + messageType MessageType + }{ + {name: "empty"}, + {name: "unknown", messageType: MessageTypeUnknown}, + {name: "custom", messageType: MessageType("reaction")}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + msg := base + msg.MessageType = tt.messageType + if err := msg.Validate(); err == nil { + t.Fatalf("expected message type %q to be rejected", tt.messageType) + } + }) + } +} + +func TestValidateInboundAllowsKnownConversationalMessageTypes(t *testing.T) { + for _, messageType := range []MessageType{ + MessageTypeText, + MessageTypeImage, + MessageTypeFile, + MessageTypeAudio, + MessageTypeVideo, + } { + t.Run(string(messageType), func(t *testing.T) { + msg := InboundMessage{ + TenantID: "tenant", + AppID: "app", + BindingID: "binding", + Channel: "telegram", + ChannelAccountID: "bot", + PlatformMessageID: "msg", + ExternalUserID: "user", + ConversationType: ConversationTypeDM, + MessageType: messageType, + } + if err := msg.Validate(); err != nil { + t.Fatalf("expected %q to be accepted, got %v", messageType, err) + } + }) + } +} + +func TestValidateInboundEventRequiresRawEventType(t *testing.T) { + msg := InboundMessage{ + TenantID: "tenant", + AppID: "app", + BindingID: "binding", + Channel: "telegram", + ChannelAccountID: "bot", + PlatformMessageID: "event-1", + MessageType: MessageTypeEvent, + } + if err := msg.Validate(); err == nil { + t.Fatalf("expected event without raw_event_type to fail") + } + msg.RawEventType = " " + if err := msg.Validate(); err == nil { + t.Fatalf("expected event with blank raw_event_type to fail") + } +} + func TestSessionIDForInboundIncludesBindingAndAccountScope(t *testing.T) { base := InboundMessage{ TenantID: "tenant", @@ -213,6 +469,20 @@ func TestAuditIDIsStableAndScoped(t *testing.T) { } } +func TestUserIdentifiersAreLengthPrefixed(t *testing.T) { + internalA := InternalUserID("tenant\x00telegram", "user", "42") + internalB := InternalUserID("tenant", "telegram\x00user", "42") + if internalA == internalB { + t.Fatalf("internal user IDs should not collide on NUL-delimited inputs: %q", internalA) + } + + hashA := UserIDHash("tenant\x00telegram", "user", "42") + hashB := UserIDHash("tenant", "telegram\x00user", "42") + if hashA == hashB { + t.Fatalf("user ID hashes should not collide on NUL-delimited inputs: %q", hashA) + } +} + func TestIdempotencyStoreDoesNotRestartCompletedMessage(t *testing.T) { store := NewInMemoryIdempotencyStore() record := IdempotencyRecord{ @@ -534,11 +804,56 @@ func TestRedactorMasksURLUserinfoWithMultipleAtSigns(t *testing.T) { strings.Contains(got, "ss@word@example.com") { t.Fatalf("redacted output leaked URL password fragments: %q", got) } - if !strings.Contains(got, "postgres://user:****@example.com/db") { + if !strings.Contains(got, "postgres://****@example.com/db") { t.Fatalf("expected URL userinfo password mask, got %q", got) } } +func TestRedactorMasksURLUserinfo(t *testing.T) { + redactor, err := NewRedactor() + if err != nil { + t.Fatalf("NewRedactor: %v", err) + } + tests := []struct { + name string + input string + want string + leaks []string + }{ + { + name: "username_only", + input: "https://token@example.com/path", + want: "https://****@example.com/path", + leaks: []string{"token@example.com"}, + }, + { + name: "percent_encoded", + input: "https://tok%40en@example.com/path", + want: "https://****@example.com/path", + leaks: []string{"tok%40en"}, + }, + { + name: "username_with_at", + input: "postgres://svc@example.com:password@db/prod", + want: "postgres://****@db/prod", + leaks: []string{"svc@example.com", "password"}, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := redactor.Redact(tt.input) + if got != tt.want { + t.Fatalf("expected %q, got %q", tt.want, got) + } + for _, leaked := range tt.leaks { + if strings.Contains(got, leaked) { + t.Fatalf("redacted output leaked %q: %q", leaked, got) + } + } + }) + } +} + func TestAuditSinkStoresSnapshot(t *testing.T) { sink := NewInMemoryAuditSink() record := AuditRecord{ diff --git a/platform/validation.go b/platform/validation.go index a5452a1c4f..32224157ae 100644 --- a/platform/validation.go +++ b/platform/validation.go @@ -14,6 +14,7 @@ import ( "net/url" "strconv" "strings" + "unicode" ) var rawSecretPrefixes = []string{ @@ -26,10 +27,26 @@ var rawSecretPrefixes = []string{ "glpat-", } +func validateRoutingIdentifier(field, value string, requiredErr error) error { + trimmed := strings.TrimSpace(value) + if trimmed == "" { + return requiredErr + } + if trimmed != value { + return fmt.Errorf("%s must not contain leading or trailing whitespace", field) + } + for _, r := range value { + if unicode.IsControl(r) { + return fmt.Errorf("%s must not contain control characters", field) + } + } + return nil +} + // Validate checks that the tenant can be used as an isolation boundary. func (t Tenant) Validate() error { - if strings.TrimSpace(t.TenantID) == "" { - return ErrTenantIDRequired + if err := validateRoutingIdentifier("tenant_id", t.TenantID, ErrTenantIDRequired); err != nil { + return err } switch t.Status { case "", TenantStatusActive, TenantStatusSuspended, TenantStatusDeleted: @@ -41,11 +58,11 @@ func (t Tenant) Validate() error { // Validate checks that the app has the identifiers required for routing. func (a AgentApp) Validate() error { - if strings.TrimSpace(a.TenantID) == "" { - return ErrTenantIDRequired + if err := validateRoutingIdentifier("tenant_id", a.TenantID, ErrTenantIDRequired); err != nil { + return err } - if strings.TrimSpace(a.AppID) == "" { - return ErrAppIDRequired + if err := validateRoutingIdentifier("app_id", a.AppID, ErrAppIDRequired); err != nil { + return err } if a.GrayPercent < 0 || a.GrayPercent > 100 { return fmt.Errorf("gray_percent must be between 0 and 100") @@ -60,11 +77,11 @@ func (a AgentApp) Validate() error { // Validate checks that model profile sensitive values are stored by reference. func (p ModelProfile) Validate() error { - if strings.TrimSpace(p.TenantID) == "" { - return ErrTenantIDRequired + if err := validateRoutingIdentifier("tenant_id", p.TenantID, ErrTenantIDRequired); err != nil { + return err } - if strings.TrimSpace(p.ProfileID) == "" { - return fmt.Errorf("profile_id is required") + if err := validateRoutingIdentifier("profile_id", p.ProfileID, fmt.Errorf("profile_id is required")); err != nil { + return err } if err := validateSecretReference("base_url_ref", p.BaseURLRef); err != nil { return err @@ -77,20 +94,20 @@ func (p ModelProfile) Validate() error { // Validate checks that a binding has safe routing and secret references. func (b ChannelBinding) Validate() error { - if strings.TrimSpace(b.TenantID) == "" { - return ErrTenantIDRequired + if err := validateRoutingIdentifier("tenant_id", b.TenantID, ErrTenantIDRequired); err != nil { + return err } - if strings.TrimSpace(b.AppID) == "" { - return ErrAppIDRequired + if err := validateRoutingIdentifier("app_id", b.AppID, ErrAppIDRequired); err != nil { + return err } - if strings.TrimSpace(b.BindingID) == "" { - return ErrBindingIDRequired + if err := validateRoutingIdentifier("binding_id", b.BindingID, ErrBindingIDRequired); err != nil { + return err } - if strings.TrimSpace(b.Channel) == "" { - return ErrChannelRequired + if err := validateRoutingIdentifier("channel", b.Channel, ErrChannelRequired); err != nil { + return err } - if strings.TrimSpace(b.AccountID) == "" { - return ErrAccountIDRequired + if err := validateRoutingIdentifier("account_id", b.AccountID, ErrAccountIDRequired); err != nil { + return err } if strings.TrimSpace(b.WebhookPath) == "" { return ErrWebhookPathRequired @@ -114,44 +131,47 @@ func (b ChannelBinding) Validate() error { // Validate checks that an inbound message has enough identity for routing. func (m InboundMessage) Validate() error { - if strings.TrimSpace(m.TenantID) == "" { - return ErrTenantIDRequired + if err := validateRoutingIdentifier("tenant_id", m.TenantID, ErrTenantIDRequired); err != nil { + return err } - if strings.TrimSpace(m.AppID) == "" { - return ErrAppIDRequired + if err := validateRoutingIdentifier("app_id", m.AppID, ErrAppIDRequired); err != nil { + return err + } + if err := validateRoutingIdentifier("binding_id", m.BindingID, ErrBindingIDRequired); err != nil { + return err } - if strings.TrimSpace(m.BindingID) == "" { - return ErrBindingIDRequired + if err := validateRoutingIdentifier("channel", m.Channel, ErrChannelRequired); err != nil { + return err } - if strings.TrimSpace(m.Channel) == "" { - return ErrChannelRequired + if err := validateRoutingIdentifier("channel_account_id", m.ChannelAccountID, ErrAccountIDRequired); err != nil { + return err } - if strings.TrimSpace(m.ChannelAccountID) == "" { - return ErrAccountIDRequired + if err := validateRoutingIdentifier("platform_message_id", m.PlatformMessageID, ErrPlatformMessageIDRequired); err != nil { + return err } - if strings.TrimSpace(m.PlatformMessageID) == "" { - return ErrPlatformMessageIDRequired + if err := validateInboundMessageType(m); err != nil { + return err } if m.MessageType == MessageTypeEvent { return nil } - if strings.TrimSpace(m.ExternalUserID) == "" { - return ErrExternalUserIDRequired + if err := validateRoutingIdentifier("external_user_id", m.ExternalUserID, ErrExternalUserIDRequired); err != nil { + return err } switch m.ConversationType { case ConversationTypeDM: return nil case ConversationTypeGroup: - if strings.TrimSpace(m.ExternalGroupID) == "" { - return ErrExternalGroupIDRequired + if err := validateRoutingIdentifier("external_group_id", m.ExternalGroupID, ErrExternalGroupIDRequired); err != nil { + return err } return nil case ConversationTypeThread: - if strings.TrimSpace(m.ExternalGroupID) == "" { - return ErrExternalGroupIDRequired + if err := validateRoutingIdentifier("external_group_id", m.ExternalGroupID, ErrExternalGroupIDRequired); err != nil { + return err } - if strings.TrimSpace(m.ThreadID) == "" { - return fmt.Errorf("thread_id is required") + if err := validateRoutingIdentifier("thread_id", m.ThreadID, fmt.Errorf("thread_id is required")); err != nil { + return err } return nil case "": @@ -161,13 +181,26 @@ func (m InboundMessage) Validate() error { } } +func validateInboundMessageType(m InboundMessage) error { + switch m.MessageType { + case MessageTypeText, MessageTypeImage, MessageTypeFile, MessageTypeAudio, MessageTypeVideo: + return nil + case MessageTypeEvent: + return validateRoutingIdentifier("raw_event_type", m.RawEventType, fmt.Errorf("raw_event_type is required")) + case "": + return fmt.Errorf("message_type is required") + default: + return fmt.Errorf("invalid message_type %q", m.MessageType) + } +} + // Validate checks that a storage profile uses references for sensitive values. func (p StorageProfile) Validate() error { - if strings.TrimSpace(p.TenantID) == "" { - return ErrTenantIDRequired + if err := validateRoutingIdentifier("tenant_id", p.TenantID, ErrTenantIDRequired); err != nil { + return err } - if strings.TrimSpace(p.ProfileID) == "" { - return fmt.Errorf("profile_id is required") + if err := validateRoutingIdentifier("profile_id", p.ProfileID, fmt.Errorf("profile_id is required")); err != nil { + return err } if err := validateSecretReference("dsn_ref", p.DSNRef); err != nil { return err From 8c34131fb22b7a1e1e3f21e09aea2235fdaf9dff Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Thu, 9 Jul 2026 14:12:08 +0800 Subject: [PATCH 38/95] fix(platform): reconcile hardened contracts with gateway --- platform/channeladapter/outbox_test.go | 2 +- platform/gateway/service_test.go | 4 +++- platform/idempotency.go | 11 +++++++++-- platform/identity.go | 11 +++++++++++ platform/redaction.go | 22 +++++++++++++++------- platform/types_test.go | 11 +++++++++-- 6 files changed, 48 insertions(+), 13 deletions(-) diff --git a/platform/channeladapter/outbox_test.go b/platform/channeladapter/outbox_test.go index e714a55578..db5d3b7788 100644 --- a/platform/channeladapter/outbox_test.go +++ b/platform/channeladapter/outbox_test.go @@ -190,7 +190,7 @@ func TestOutboxRedactsFailureDetails(t *testing.T) { } } if !strings.Contains(failed.LastError, "Authorization: ****") || - !strings.Contains(failed.LastError, "user:****@") || + !strings.Contains(failed.LastError, "postgres://****@example/db") || !strings.Contains(failed.LastError, "api_key=****") { t.Fatalf("expected redacted diagnostic, got %q", failed.LastError) } diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go index 96daf29088..095df5d4da 100644 --- a/platform/gateway/service_test.go +++ b/platform/gateway/service_test.go @@ -259,7 +259,9 @@ func TestServiceHandleInboundSerializesSameSession(t *testing.T) { assert.False(t, busy.Duplicate) assert.True(t, busy.Processing) assert.Equal(t, platform.IdempotencyStatusProcessing, busy.Status) - assert.Equal(t, "tenant:tenant-a:app:app:channel:wecom:dm:user-1", busy.SessionID) + wantSessionID, err := platform.SessionIDForInbound(second) + require.NoError(t, err) + assert.Equal(t, wantSessionID, busy.SessionID) assert.Len(t, r.calls, 1) record, ok, err := svc.idempotencyStore.Get( ctx, diff --git a/platform/idempotency.go b/platform/idempotency.go index 005b2aae95..5dfee8bdd9 100644 --- a/platform/idempotency.go +++ b/platform/idempotency.go @@ -86,7 +86,7 @@ func (s *InMemoryIdempotencyStore) MarkReplyFailed( key string, resultRef string, ) (IdempotencyRecord, error) { - return s.update(ctx, key, IdempotencyStatusCompleted, IdempotencyStatusReplyFailed, resultRef) + return s.update(ctx, key, "", IdempotencyStatusReplyFailed, resultRef) } // Get returns the record for key. @@ -119,7 +119,14 @@ func (s *InMemoryIdempotencyStore) update( if !ok { return IdempotencyRecord{}, ErrIdempotencyRecordNotFound } - if record.Status != from { + if from != "" && record.Status != from { + return IdempotencyRecord{}, fmt.Errorf( + "invalid idempotency transition from %q to %q", + record.Status, + status, + ) + } + if from == "" && record.Status != IdempotencyStatusCompleted && record.Status != IdempotencyStatusProcessing { return IdempotencyRecord{}, fmt.Errorf( "invalid idempotency transition from %q to %q", record.Status, diff --git a/platform/identity.go b/platform/identity.go index cf92feb558..77939f0b27 100644 --- a/platform/identity.go +++ b/platform/identity.go @@ -12,6 +12,8 @@ import ( "crypto/sha256" "encoding/hex" "fmt" + "net/url" + "strings" ) type stableIDPart struct { @@ -189,3 +191,12 @@ func stableID(prefix string, parts ...stableIDPart) string { func writeStablePart(hash interface{ Write([]byte) (int, error) }, name, value string) { fmt.Fprintf(hash, "%d:%s=%d:%s;", len(name), name, len(value), value) } + +func shortHash(parts ...string) string { + sum := sha256.Sum256([]byte(strings.Join(parts, "\x00"))) + return hex.EncodeToString(sum[:])[:24] +} + +func escapeKeyPart(value string) string { + return url.PathEscape(strings.TrimSpace(value)) +} diff --git a/platform/redaction.go b/platform/redaction.go index 569bc53eaf..46643b4232 100644 --- a/platform/redaction.go +++ b/platform/redaction.go @@ -14,12 +14,13 @@ import ( ) var defaultRedactionPatterns = []*regexp.Regexp{ - regexp.MustCompile(`(?i)(Authorization:\s*Basic\s+)[A-Za-z0-9._~+/\-]+=*`), + regexp.MustCompile(`(?i)(Authorization:\s*(?:Basic|Bearer)\s+)[A-Za-z0-9._~+/\-]+=*`), + regexp.MustCompile(`(?i)(Authorization\s*=\s*Bearer\s+)[^\r\n\s]+`), regexp.MustCompile(`(?i)(Bearer\s+)[A-Za-z0-9._~+/\-]+=*`), - regexp.MustCompile(`(?im)(authorization\s*:\s*)[^\r\n]+`), - regexp.MustCompile(`(?im)(authorization\s*=\s*)[^\r\n]+`), - regexp.MustCompile(`(?i)(api[_-]?key|token|secret|password|passwd|authorization|cookie)=([^&\s]+)`), - regexp.MustCompile(`(?i)(api[_-]?key|token|secret|password|passwd|authorization|cookie):\s*([^,\s]+)`), + regexp.MustCompile(`(?im)(authorization\s*:\s*(?:token|digest)\s+)[^\r\n]+`), + regexp.MustCompile(`(?im)(authorization\s*=\s*(?:token|digest)\s+)[^\r\n]+`), + regexp.MustCompile(`(?i)(api[_-]?key|token|secret|password|passwd|cookie)=([^&\s]+)`), + regexp.MustCompile(`(?i)(api[_-]?key|token|secret|password|passwd|cookie):\s*([^,\s]+)`), regexp.MustCompile(`(?i)("(?:api[_-]?key|token|secret|password|passwd|authorization|cookie)"\s*:\s*")([^"]+)(")`), regexp.MustCompile(`(?i)(sk-[A-Za-z0-9._~+/\-]{8,})`), regexp.MustCompile(`(?i)[a-z][a-z0-9+.-]*://[^\s/?#]*@[^\s/?#]+`), @@ -62,8 +63,15 @@ func (r *Redactor) Redact(text string) string { func redactMatch(match string) string { lower := strings.ToLower(match) - if strings.Contains(lower, "authorization:") && strings.Contains(lower, "basic ") { - return match[:strings.Index(lower, "basic ")+6] + "****" + if strings.Contains(lower, "authorization:") { + if idx := strings.Index(match, ":"); idx >= 0 { + return match[:idx+1] + " ****" + } + } + if strings.Contains(lower, "authorization=") { + if idx := strings.Index(match, "="); idx >= 0 { + return match[:idx+1] + "****" + } } if strings.Contains(lower, "bearer ") { return match[:strings.Index(lower, "bearer ")+7] + "****" diff --git a/platform/types_test.go b/platform/types_test.go index d7323772a8..0e49a84977 100644 --- a/platform/types_test.go +++ b/platform/types_test.go @@ -610,8 +610,15 @@ func TestIdempotencyStoreEnforcesStateTransitions(t *testing.T) { if !started { t.Fatalf("processing record should start") } - if _, err := processingStore.MarkReplyFailed(ctx, processing.IdempotencyKey, "outbound-1"); err == nil { - t.Fatalf("reply failure should only be allowed after completion") + replyFailed, err := processingStore.MarkReplyFailed(ctx, processing.IdempotencyKey, "outbound-1") + if err != nil { + t.Fatalf("mark reply failed from processing: %v", err) + } + if replyFailed.Status != IdempotencyStatusReplyFailed || replyFailed.ResultRef != "outbound-1" { + t.Fatalf("unexpected processing reply-failed record: %#v", replyFailed) + } + if _, err := processingStore.Complete(ctx, processing.IdempotencyKey, "outbound-2"); err == nil { + t.Fatalf("reply-failed processing record should not transition to completed") } completedStore := NewInMemoryIdempotencyStore() From 8c6ffc500d6e0a223df0e4a939b95a6990a086c0 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 17:14:53 +0800 Subject: [PATCH 39/95] platform: add budget decision audit contracts --- platform/budget_audit.go | 374 ++++++++++++++++++++++++++++++++++ platform/budget_audit_test.go | 298 +++++++++++++++++++++++++++ 2 files changed, 672 insertions(+) create mode 100644 platform/budget_audit.go create mode 100644 platform/budget_audit_test.go diff --git a/platform/budget_audit.go b/platform/budget_audit.go new file mode 100644 index 0000000000..77e882da06 --- /dev/null +++ b/platform/budget_audit.go @@ -0,0 +1,374 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "fmt" + "strconv" + "strings" + "time" +) + +// BudgetDecisionOutcome is the externally visible budget gate outcome. +type BudgetDecisionOutcome string + +const ( + // BudgetDecisionOutcomeAllow means the request can continue unchanged. + BudgetDecisionOutcomeAllow BudgetDecisionOutcome = "allow" + // BudgetDecisionOutcomeDeny means the request must be rejected. + BudgetDecisionOutcomeDeny BudgetDecisionOutcome = "deny" + // BudgetDecisionOutcomeDegrade means the request can continue with a lower-cost path. + BudgetDecisionOutcomeDegrade BudgetDecisionOutcome = "degrade" +) + +// BudgetDecisionAuditInput contains safe dimensions for a budget gate decision. +type BudgetDecisionAuditInput struct { + TenantID string + AppID string + RequestID string + TraceID string + Decision BudgetDecision + Estimate UsageEstimate + Quota TenantQuota + Outcome BudgetDecisionOutcome + DegradeStrategy string + CreatedAt time.Time +} + +// BudgetDecisionSummary is a safe, auditable summary of one budget decision. +type BudgetDecisionSummary struct { + TenantID string + AppID string + RequestID string + TraceID string + Outcome BudgetDecisionOutcome + Reason string + DegradeStrategy string + EstimatedPrompt int + EstimatedCompletion int + EstimatedTotalTokens int + EstimatedCost float64 + MaxPromptTokens int + MaxCompletionTokens int + MaxTotalTokens int + MaxCost float64 + RedactionVersion string + CreatedAt time.Time +} + +// NewBudgetDecisionSummary builds a safe summary for budget gate observability. +func NewBudgetDecisionSummary(input BudgetDecisionAuditInput) (BudgetDecisionSummary, error) { + normalized, err := input.normalize() + if err != nil { + return BudgetDecisionSummary{}, err + } + totalTokens, err := normalized.Estimate.effectiveTotalTokens() + if err != nil { + return BudgetDecisionSummary{}, err + } + summary := BudgetDecisionSummary{ + TenantID: normalized.TenantID, + AppID: normalized.AppID, + RequestID: normalized.RequestID, + TraceID: normalized.TraceID, + Outcome: normalized.Outcome, + Reason: strings.TrimSpace(normalized.Decision.Reason), + DegradeStrategy: normalized.DegradeStrategy, + EstimatedPrompt: normalized.Estimate.PromptTokens, + EstimatedCompletion: normalized.Estimate.CompletionTokens, + EstimatedTotalTokens: totalTokens, + EstimatedCost: normalized.Estimate.Cost, + MaxPromptTokens: normalized.Quota.MaxPromptTokens, + MaxCompletionTokens: normalized.Quota.MaxCompletionTokens, + MaxTotalTokens: normalized.Quota.MaxTotalTokens, + MaxCost: normalized.Quota.MaxCost, + RedactionVersion: "platform-budget-decision-v1", + CreatedAt: normalized.CreatedAt, + } + if err := summary.Validate(); err != nil { + return BudgetDecisionSummary{}, err + } + return summary, nil +} + +// NewBudgetDecisionAuditRecord maps one budget decision into an audit record. +func NewBudgetDecisionAuditRecord(input BudgetDecisionAuditInput) (AuditRecord, error) { + summary, err := NewBudgetDecisionSummary(input) + if err != nil { + return AuditRecord{}, err + } + record := AuditRecord{ + TenantID: summary.TenantID, + AppID: summary.AppID, + AuditID: summary.auditID(), + RequestID: summary.RequestID, + TraceID: summary.TraceID, + ToolName: "budget:tenant", + Decision: string(summary.Outcome), + DecisionReason: summary.Reason, + Cost: summary.EstimatedCost, + TokenUsageJSON: summary.tokenUsageRef(), + RedactedDetailRef: summary.DetailRef(), + RedactionVersion: summary.RedactionVersion, + CreatedAt: summary.CreatedAt, + } + if err := record.Validate(); err != nil { + return AuditRecord{}, err + } + return record, nil +} + +// Validate checks that the summary is complete and safe to expose. +func (s BudgetDecisionSummary) Validate() error { + if strings.TrimSpace(s.TenantID) == "" { + return ErrTenantIDRequired + } + if err := validateAuditRedactedText("app_id", s.AppID); err != nil { + return err + } + if err := validateAuditRedactedText("request_id", s.RequestID); err != nil { + return err + } + if err := validateAuditRedactedText("trace_id", s.TraceID); err != nil { + return err + } + if err := validateAuditRedactedText("decision_reason", s.Reason); err != nil { + return err + } + if err := validateAuditRedactedText("degrade_strategy", s.DegradeStrategy); err != nil { + return err + } + if strings.TrimSpace(s.Reason) == "" && s.Outcome != BudgetDecisionOutcomeAllow { + return fmt.Errorf("reason is required for non-allow budget outcomes") + } + estimate := UsageEstimate{ + PromptTokens: s.EstimatedPrompt, + CompletionTokens: s.EstimatedCompletion, + TotalTokens: s.EstimatedTotalTokens, + Cost: s.EstimatedCost, + } + if err := validateUsageEstimate(estimate); err != nil { + return err + } + canonicalTotal, err := estimate.effectiveTotalTokens() + if err != nil { + return err + } + if canonicalTotal != s.EstimatedTotalTokens { + return fmt.Errorf("estimated_total_tokens must match effective total tokens") + } + quota := s.quota() + if err := quota.Validate(); err != nil { + return err + } + expected, err := quota.Check(estimate) + if err != nil { + return err + } + switch s.Outcome { + case BudgetDecisionOutcomeAllow: + if !expected.Allowed { + return fmt.Errorf("allow outcome does not match budget decision") + } + if strings.TrimSpace(s.Reason) != "" { + return fmt.Errorf("reason must be empty for allow budget outcomes") + } + if strings.TrimSpace(s.DegradeStrategy) != "" { + return fmt.Errorf("degrade_strategy must be empty for allow budget outcomes") + } + case BudgetDecisionOutcomeDeny: + if expected.Allowed { + return fmt.Errorf("deny outcome does not match budget decision") + } + if s.Reason != expected.Reason { + return fmt.Errorf("reason must match budget decision") + } + if strings.TrimSpace(s.DegradeStrategy) != "" { + return fmt.Errorf("degrade_strategy must be empty for deny budget outcomes") + } + case BudgetDecisionOutcomeDegrade: + if expected.Allowed { + return fmt.Errorf("degrade outcome does not match budget decision") + } + if s.Reason != expected.Reason { + return fmt.Errorf("reason must match budget decision") + } + if strings.TrimSpace(s.DegradeStrategy) == "" { + return fmt.Errorf("degrade_strategy is required for degrade budget outcomes") + } + case "": + return fmt.Errorf("outcome is required") + default: + return fmt.Errorf("invalid budget decision outcome %q", s.Outcome) + } + if strings.TrimSpace(s.RedactionVersion) == "" { + return fmt.Errorf("redaction_version is required") + } + if s.CreatedAt.IsZero() { + return fmt.Errorf("created_at is required") + } + if err := validateAuditRedactedText("redacted_detail_ref", s.DetailRef()); err != nil { + return err + } + return nil +} + +// DetailRef returns compact non-secret detail for audit logs. +func (s BudgetDecisionSummary) DetailRef() string { + parts := []string{ + "outcome:" + string(s.Outcome), + "estimated_total_tokens:" + fmt.Sprint(s.EstimatedTotalTokens), + "estimated_cost:" + fmt.Sprintf("%.6f", s.EstimatedCost), + } + if s.AppID != "" { + parts = append(parts, "app:"+s.AppID) + } + if s.RequestID != "" { + parts = append(parts, "request:"+s.RequestID) + } + if s.TraceID != "" { + parts = append(parts, "trace:"+s.TraceID) + } + if s.Reason != "" { + parts = append(parts, "reason:"+s.Reason) + } + if s.DegradeStrategy != "" { + parts = append(parts, "degrade:"+s.DegradeStrategy) + } + if s.MaxPromptTokens > 0 { + parts = append(parts, "max_prompt_tokens:"+fmt.Sprint(s.MaxPromptTokens)) + } + if s.MaxCompletionTokens > 0 { + parts = append(parts, "max_completion_tokens:"+fmt.Sprint(s.MaxCompletionTokens)) + } + if s.MaxTotalTokens > 0 { + parts = append(parts, "max_total_tokens:"+fmt.Sprint(s.MaxTotalTokens)) + } + if s.MaxCost > 0 { + parts = append(parts, "max_cost:"+fmt.Sprintf("%.6f", s.MaxCost)) + } + return strings.Join(parts, " ") +} + +func (i BudgetDecisionAuditInput) normalize() (BudgetDecisionAuditInput, error) { + i.TenantID = strings.TrimSpace(i.TenantID) + if i.TenantID == "" { + return BudgetDecisionAuditInput{}, ErrTenantIDRequired + } + i.AppID = strings.TrimSpace(i.AppID) + i.RequestID = strings.TrimSpace(i.RequestID) + i.TraceID = strings.TrimSpace(i.TraceID) + i.DegradeStrategy = strings.TrimSpace(i.DegradeStrategy) + if err := validateUsageEstimate(i.Estimate); err != nil { + return BudgetDecisionAuditInput{}, err + } + if err := i.Quota.Validate(); err != nil { + return BudgetDecisionAuditInput{}, err + } + expected, err := i.Quota.Check(i.Estimate) + if err != nil { + return BudgetDecisionAuditInput{}, err + } + i.Decision.Reason = strings.TrimSpace(i.Decision.Reason) + if i.Decision.Allowed != expected.Allowed { + return BudgetDecisionAuditInput{}, fmt.Errorf("budget decision does not match quota and estimate") + } + if i.Decision.Reason != expected.Reason { + return BudgetDecisionAuditInput{}, fmt.Errorf("budget decision reason does not match quota and estimate") + } + i.Outcome = normalizeBudgetOutcome(i.Decision, i.Outcome, i.DegradeStrategy) + if i.Outcome == BudgetDecisionOutcomeAllow && !i.Decision.Allowed { + return BudgetDecisionAuditInput{}, fmt.Errorf("allow outcome requires an allowed budget decision") + } + if i.Outcome != BudgetDecisionOutcomeAllow && i.Decision.Allowed { + return BudgetDecisionAuditInput{}, fmt.Errorf("non-allow outcome requires a denied budget decision") + } + for field, value := range map[string]string{ + "app_id": i.AppID, + "request_id": i.RequestID, + "trace_id": i.TraceID, + "decision_reason": i.Decision.Reason, + "degrade_strategy": i.DegradeStrategy, + } { + if err := validateAuditRedactedText(field, value); err != nil { + return BudgetDecisionAuditInput{}, err + } + } + return i, nil +} + +func normalizeBudgetOutcome(decision BudgetDecision, outcome BudgetDecisionOutcome, degradeStrategy string) BudgetDecisionOutcome { + outcome = BudgetDecisionOutcome(strings.TrimSpace(string(outcome))) + if outcome != "" { + return outcome + } + if decision.Allowed { + return BudgetDecisionOutcomeAllow + } + if strings.TrimSpace(degradeStrategy) != "" { + return BudgetDecisionOutcomeDegrade + } + return BudgetDecisionOutcomeDeny +} + +func (s BudgetDecisionSummary) auditID() string { + return AuditID( + s.TenantID, + s.AppID, + s.RequestID, + s.TraceID, + string(s.Outcome), + s.Reason, + s.DegradeStrategy, + fmt.Sprint(s.EstimatedPrompt), + fmt.Sprint(s.EstimatedCompletion), + fmt.Sprint(s.EstimatedTotalTokens), + canonicalBudgetCost(s.EstimatedCost), + fmt.Sprint(s.MaxPromptTokens), + fmt.Sprint(s.MaxCompletionTokens), + fmt.Sprint(s.MaxTotalTokens), + canonicalBudgetCost(s.MaxCost), + ) +} + +func (s BudgetDecisionSummary) tokenUsageRef() string { + return strings.Join([]string{ + "prompt_tokens:" + fmt.Sprint(s.EstimatedPrompt), + "completion_tokens:" + fmt.Sprint(s.EstimatedCompletion), + "total_tokens:" + fmt.Sprint(s.EstimatedTotalTokens), + }, " ") +} + +func (s BudgetDecisionSummary) quota() TenantQuota { + return TenantQuota{ + MaxPromptTokens: s.MaxPromptTokens, + MaxCompletionTokens: s.MaxCompletionTokens, + MaxTotalTokens: s.MaxTotalTokens, + MaxCost: s.MaxCost, + } +} + +func validateUsageEstimate(estimate UsageEstimate) error { + if estimate.PromptTokens < 0 || + estimate.CompletionTokens < 0 || + estimate.TotalTokens < 0 { + return fmt.Errorf("usage estimate values must be non-negative") + } + if !isFiniteNonNegative(estimate.Cost) { + return fmt.Errorf("usage estimate cost must be finite and non-negative") + } + if _, err := estimate.effectiveTotalTokens(); err != nil { + return err + } + return nil +} + +func canonicalBudgetCost(value float64) string { + return strconv.FormatFloat(value, 'g', -1, 64) +} diff --git a/platform/budget_audit_test.go b/platform/budget_audit_test.go new file mode 100644 index 0000000000..53ae7ca1e8 --- /dev/null +++ b/platform/budget_audit_test.go @@ -0,0 +1,298 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "errors" + "strings" + "testing" + "time" +) + +func TestNewBudgetDecisionSummaryBuildsDenyAuditSummary(t *testing.T) { + now := time.Unix(100, 0) + quota := TenantQuota{MaxTotalTokens: 100, MaxCost: 1.25} + estimate := UsageEstimate{PromptTokens: 80, CompletionTokens: 30, Cost: 0.50} + decision, err := quota.Check(estimate) + if err != nil { + t.Fatalf("quota check: %v", err) + } + + summary, err := NewBudgetDecisionSummary(BudgetDecisionAuditInput{ + TenantID: " tenant ", + AppID: " app ", + RequestID: " request-1 ", + TraceID: " trace-1 ", + Decision: decision, + Estimate: estimate, + Quota: quota, + CreatedAt: now, + }) + if err != nil { + t.Fatalf("NewBudgetDecisionSummary: %v", err) + } + if summary.TenantID != "tenant" || + summary.AppID != "app" || + summary.RequestID != "request-1" || + summary.TraceID != "trace-1" || + summary.Outcome != BudgetDecisionOutcomeDeny || + summary.Reason != "total_tokens_exceeded" || + summary.EstimatedTotalTokens != 110 || + summary.MaxTotalTokens != 100 || + !summary.CreatedAt.Equal(now) { + t.Fatalf("unexpected summary: %+v", summary) + } + detail := summary.DetailRef() + if !strings.Contains(detail, "outcome:deny") || + !strings.Contains(detail, "reason:total_tokens_exceeded") || + !strings.Contains(detail, "estimated_total_tokens:110") { + t.Fatalf("detail ref missing decision fields: %q", detail) + } +} + +func TestNewBudgetDecisionAuditRecordBuildsStableRedactedRecord(t *testing.T) { + now := time.Unix(200, 0) + input := BudgetDecisionAuditInput{ + TenantID: "tenant", + AppID: "app", + RequestID: "request-1", + TraceID: "trace-1", + Decision: BudgetDecision{Reason: "cost_exceeded"}, + Estimate: UsageEstimate{PromptTokens: 10, CompletionTokens: 5, Cost: 2.00}, + Quota: TenantQuota{MaxCost: 1.00}, + CreatedAt: now, + } + record, err := NewBudgetDecisionAuditRecord(input) + if err != nil { + t.Fatalf("NewBudgetDecisionAuditRecord: %v", err) + } + if record.TenantID != "tenant" || + record.AppID != "app" || + record.ToolName != "budget:tenant" || + record.Decision != "deny" || + record.DecisionReason != "cost_exceeded" || + record.RequestID != "request-1" || + record.TraceID != "trace-1" || + record.RedactionVersion != "platform-budget-decision-v1" || + !record.CreatedAt.Equal(now) { + t.Fatalf("unexpected audit record: %+v", record) + } + if !strings.Contains(record.TokenUsageJSON, "prompt_tokens:10") || + !strings.Contains(record.TokenUsageJSON, "completion_tokens:5") || + !strings.Contains(record.TokenUsageJSON, "total_tokens:15") { + t.Fatalf("unexpected token usage ref: %q", record.TokenUsageJSON) + } + if strings.Contains(record.RedactedDetailRef, "sk-secret") || + strings.Contains(record.RedactedDetailRef, "password") { + t.Fatalf("audit detail leaked secret content: %q", record.RedactedDetailRef) + } + + again, err := NewBudgetDecisionAuditRecord(input) + if err != nil { + t.Fatalf("NewBudgetDecisionAuditRecord again: %v", err) + } + if record.AuditID != again.AuditID { + t.Fatalf("expected stable audit id, got %q and %q", record.AuditID, again.AuditID) + } +} + +func TestNewBudgetDecisionSummaryBuildsDegradeOutcome(t *testing.T) { + summary, err := NewBudgetDecisionSummary(BudgetDecisionAuditInput{ + TenantID: "tenant", + AppID: "app", + RequestID: "request-1", + TraceID: "trace-1", + Decision: BudgetDecision{Reason: "cost_exceeded"}, + Estimate: UsageEstimate{PromptTokens: 10, CompletionTokens: 5, Cost: 2.00}, + Quota: TenantQuota{MaxCost: 1.00}, + DegradeStrategy: "fallback_model", + CreatedAt: time.Unix(300, 0), + }) + if err != nil { + t.Fatalf("NewBudgetDecisionSummary: %v", err) + } + if summary.Outcome != BudgetDecisionOutcomeDegrade || + summary.DegradeStrategy != "fallback_model" || + !strings.Contains(summary.DetailRef(), "degrade:fallback_model") { + t.Fatalf("unexpected degrade summary: %+v", summary) + } +} + +func TestBudgetDecisionSummaryValidationRejectsUnsafeOrInconsistentFields(t *testing.T) { + valid := BudgetDecisionSummary{ + TenantID: "tenant", + AppID: "app", + RequestID: "request-1", + TraceID: "trace-1", + Outcome: BudgetDecisionOutcomeDeny, + Reason: "cost_exceeded", + EstimatedPrompt: 10, + EstimatedCompletion: 5, + EstimatedTotalTokens: 15, + EstimatedCost: 2.00, + MaxCost: 1.00, + RedactionVersion: "platform-budget-decision-v1", + CreatedAt: time.Unix(400, 0), + } + if err := valid.Validate(); err != nil { + t.Fatalf("Validate valid summary: %v", err) + } + + unsafeTrace := valid + unsafeTrace.TraceID = "token=sk-secret-token" + if err := unsafeTrace.Validate(); err == nil || !strings.Contains(err.Error(), "trace_id") { + t.Fatalf("expected unsafe trace rejection, got %v", err) + } + + unsafeDetail := valid + unsafeDetail.DegradeStrategy = "api_key: sk-secret-token" + unsafeDetail.Outcome = BudgetDecisionOutcomeDegrade + if err := unsafeDetail.Validate(); err == nil || !strings.Contains(err.Error(), "degrade_strategy") { + t.Fatalf("expected unsafe degrade strategy rejection, got %v", err) + } + + allowWithReason := valid + allowWithReason.Outcome = BudgetDecisionOutcomeAllow + allowWithReason.Reason = "should_be_empty" + allowWithReason.EstimatedCost = 0.50 + allowWithReason.MaxCost = 1.00 + if err := allowWithReason.Validate(); err == nil || !strings.Contains(err.Error(), "reason") { + t.Fatalf("expected allow reason rejection, got %v", err) + } + + degradeWithoutStrategy := valid + degradeWithoutStrategy.Outcome = BudgetDecisionOutcomeDegrade + if err := degradeWithoutStrategy.Validate(); err == nil || !strings.Contains(err.Error(), "degrade_strategy") { + t.Fatalf("expected missing degrade strategy rejection, got %v", err) + } + + invalidEstimate := valid + invalidEstimate.EstimatedCost = -0.01 + if err := invalidEstimate.Validate(); err == nil || !strings.Contains(err.Error(), "usage estimate cost") { + t.Fatalf("expected invalid estimate rejection, got %v", err) + } + + underReportedTotal := valid + underReportedTotal.EstimatedPrompt = 80 + underReportedTotal.EstimatedCompletion = 30 + underReportedTotal.EstimatedTotalTokens = 1 + underReportedTotal.MaxTotalTokens = 100 + underReportedTotal.MaxCost = 0 + underReportedTotal.Reason = "total_tokens_exceeded" + if err := underReportedTotal.Validate(); err == nil || !strings.Contains(err.Error(), "effective total tokens") { + t.Fatalf("expected under-reported total rejection, got %v", err) + } +} + +func TestBudgetDecisionAuditInputRejectsMismatchedDecisionAndOutcome(t *testing.T) { + _, err := NewBudgetDecisionSummary(BudgetDecisionAuditInput{ + TenantID: "tenant", + Decision: BudgetDecision{Allowed: true}, + Outcome: BudgetDecisionOutcomeDeny, + Estimate: UsageEstimate{}, + CreatedAt: time.Unix(500, 0), + }) + if err == nil || !strings.Contains(err.Error(), "non-allow outcome") { + t.Fatalf("expected mismatched denied outcome rejection, got %v", err) + } + + _, err = NewBudgetDecisionSummary(BudgetDecisionAuditInput{ + TenantID: "tenant", + Decision: BudgetDecision{Reason: "cost_exceeded"}, + Outcome: BudgetDecisionOutcomeAllow, + Estimate: UsageEstimate{PromptTokens: 10, CompletionTokens: 5, Cost: 2.00}, + Quota: TenantQuota{MaxCost: 1.00}, + CreatedAt: time.Unix(500, 0), + }) + if err == nil || !strings.Contains(err.Error(), "allow outcome") { + t.Fatalf("expected mismatched allow outcome rejection, got %v", err) + } + + _, err = NewBudgetDecisionSummary(BudgetDecisionAuditInput{ + TenantID: "tenant", + Decision: BudgetDecision{Reason: "cost_exceeded"}, + Estimate: UsageEstimate{PromptTokens: 1, CompletionTokens: 1, Cost: 0.01}, + Quota: TenantQuota{MaxCost: 1.00}, + CreatedAt: time.Unix(500, 0), + }) + if err == nil || !strings.Contains(err.Error(), "budget decision") { + t.Fatalf("expected decision/quota mismatch rejection, got %v", err) + } + + _, err = NewBudgetDecisionSummary(BudgetDecisionAuditInput{ + TenantID: "tenant", + Decision: BudgetDecision{Reason: "cost_exceeded"}, + Estimate: UsageEstimate{PromptTokens: 10, CompletionTokens: 5, Cost: 2.00}, + Quota: TenantQuota{MaxCost: 1.00}, + Outcome: BudgetDecisionOutcomeDegrade, + CreatedAt: time.Unix(500, 0), + }) + if err == nil || !strings.Contains(err.Error(), "degrade_strategy") { + t.Fatalf("expected explicit degrade strategy requirement, got %v", err) + } +} + +func TestNewBudgetDecisionSummaryRequiresTenant(t *testing.T) { + _, err := NewBudgetDecisionSummary(BudgetDecisionAuditInput{ + Decision: BudgetDecision{Allowed: true}, + CreatedAt: time.Unix(600, 0), + }) + if !errors.Is(err, ErrTenantIDRequired) { + t.Fatalf("expected tenant requirement, got %v", err) + } +} + +func TestBudgetDecisionAuditIDIncludesEstimateAndQuotaBoundary(t *testing.T) { + now := time.Unix(700, 0) + base := BudgetDecisionAuditInput{ + TenantID: "tenant", + AppID: "app", + RequestID: "request-1", + TraceID: "trace-1", + Decision: BudgetDecision{Reason: "cost_exceeded"}, + Estimate: UsageEstimate{PromptTokens: 10, CompletionTokens: 5, Cost: 2.00}, + Quota: TenantQuota{MaxCost: 1.00}, + CreatedAt: now, + } + record, err := NewBudgetDecisionAuditRecord(base) + if err != nil { + t.Fatalf("NewBudgetDecisionAuditRecord base: %v", err) + } + + differentEstimate := base + differentEstimate.Estimate = UsageEstimate{PromptTokens: 10, CompletionTokens: 5, Cost: 3.00} + differentEstimate.Quota = TenantQuota{MaxCost: 1.00} + estimateRecord, err := NewBudgetDecisionAuditRecord(differentEstimate) + if err != nil { + t.Fatalf("NewBudgetDecisionAuditRecord different estimate: %v", err) + } + if record.AuditID == estimateRecord.AuditID { + t.Fatalf("expected different audit id for changed estimate, got %q", record.AuditID) + } + + differentQuota := base + differentQuota.Quota = TenantQuota{MaxCost: 1.50} + quotaRecord, err := NewBudgetDecisionAuditRecord(differentQuota) + if err != nil { + t.Fatalf("NewBudgetDecisionAuditRecord different quota: %v", err) + } + if record.AuditID == quotaRecord.AuditID { + t.Fatalf("expected different audit id for changed quota, got %q", record.AuditID) + } + + closeEstimate := base + closeEstimate.Estimate = UsageEstimate{PromptTokens: 10, CompletionTokens: 5, Cost: 2.0000001} + closeEstimateRecord, err := NewBudgetDecisionAuditRecord(closeEstimate) + if err != nil { + t.Fatalf("NewBudgetDecisionAuditRecord close estimate: %v", err) + } + if record.AuditID == closeEstimateRecord.AuditID { + t.Fatalf("expected different audit id for close cost estimate, got %q", record.AuditID) + } +} From e29fb65810bba0fa6062483859668c2ab5afc448 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Wed, 8 Jul 2026 17:32:51 +0800 Subject: [PATCH 40/95] platform/gateway: add minimum loop acceptance test --- platform/gateway/service_test.go | 246 +++++++++++++++++++++++++++++++ 1 file changed, 246 insertions(+) diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go index 095df5d4da..9eec9126ce 100644 --- a/platform/gateway/service_test.go +++ b/platform/gateway/service_test.go @@ -99,6 +99,159 @@ func TestServiceHandleInboundEnqueuesChannelOutbox(t *testing.T) { assert.Equal(t, 5, record.MaxAttempts) } +func TestServiceHandleInboundCoversMinimumLoopAcceptance(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + wecomRunner := &recordingRunner{response: "wecom reply"} + telegramRunner := &recordingRunner{response: "telegram reply"} + require.NoError(t, registry.Register(validRuntimeForBinding( + "tenant-a", + "app-wecom", + "binding-wecom", + "wecom", + "acct-wecom", + wecomRunner, + ))) + require.NoError(t, registry.Register(validRuntimeForBinding( + "tenant-b", + "app-telegram", + "binding-telegram", + "telegram", + "acct-telegram", + telegramRunner, + ))) + require.NoError(t, registry.Register(validRuntimeForBinding( + "tenant-b", + "app-wecom", + "binding-wecom-tenant-b", + "wecom", + "acct-wecom", + wecomRunner, + ))) + idempotency := platform.NewInMemoryIdempotencyStore() + outbox := channeladapter.NewInMemoryOutboxStore() + audit := platform.NewInMemoryAuditSink() + svc := NewService( + registry, + idempotency, + NewOutboxBackedOutboundStore(outbox), + WithAuditSink(audit), + ) + wecomDM := inboundForRuntime( + "tenant-a", + "app-wecom", + "binding-wecom", + "wecom", + "acct-wecom", + "msg-wecom-dm", + "user-shared", + "hello wecom", + ) + telegramDM := inboundForRuntime( + "tenant-b", + "app-telegram", + "binding-telegram", + "telegram", + "acct-telegram", + "msg-telegram-dm", + "user-shared", + "hello telegram", + ) + tenantBWeComDM := inboundForRuntime( + "tenant-b", + "app-wecom", + "binding-wecom-tenant-b", + "wecom", + "acct-wecom", + "msg-tenant-b-wecom-dm", + "user-shared", + "hello same channel", + ) + wecomGroup := inboundForRuntime( + "tenant-a", + "app-wecom", + "binding-wecom", + "wecom", + "acct-wecom", + "msg-wecom-group", + "user-shared", + "hello group", + ) + wecomGroup.ConversationType = platform.ConversationTypeGroup + wecomGroup.ExternalGroupID = "room-1" + + wecomResult, err := svc.HandleInbound(ctx, wecomDM) + require.NoError(t, err) + telegramResult, err := svc.HandleInbound(ctx, telegramDM) + require.NoError(t, err) + tenantBWeComResult, err := svc.HandleInbound(ctx, tenantBWeComDM) + require.NoError(t, err) + groupResult, err := svc.HandleInbound(ctx, wecomGroup) + require.NoError(t, err) + duplicate, err := svc.HandleInbound(ctx, wecomDM) + require.NoError(t, err) + + assert.False(t, wecomResult.Duplicate) + assert.True(t, duplicate.Duplicate) + assert.Equal(t, wecomResult.Outbound, duplicate.Outbound) + assert.Len(t, wecomRunner.calls, 3) + assert.Len(t, telegramRunner.calls, 1) + assert.Equal(t, "wecom reply", wecomResult.Outbound.Content) + assert.Equal(t, "telegram reply", telegramResult.Outbound.Content) + assert.Equal(t, "wecom reply", tenantBWeComResult.Outbound.Content) + assert.Equal(t, "wecom reply", groupResult.Outbound.Content) + wantWeComSessionID, err := platform.SessionIDForInbound(wecomDM) + require.NoError(t, err) + wantTelegramSessionID, err := platform.SessionIDForInbound(telegramDM) + require.NoError(t, err) + wantTenantBWeComSessionID, err := platform.SessionIDForInbound(tenantBWeComDM) + require.NoError(t, err) + wantGroupSessionID, err := platform.SessionIDForInbound(wecomGroup) + require.NoError(t, err) + assert.Equal(t, wantWeComSessionID, wecomResult.SessionID) + assert.Equal(t, wantTelegramSessionID, telegramResult.SessionID) + assert.Equal(t, wantTenantBWeComSessionID, tenantBWeComResult.SessionID) + assert.Equal(t, wantGroupSessionID, groupResult.SessionID) + assert.NotContains(t, wecomResult.SessionID, "user-shared") + assert.NotContains(t, wecomResult.SessionID, ":") + assert.NotEqual(t, wecomResult.SessionID, telegramResult.SessionID) + assert.NotEqual(t, wecomResult.SessionID, tenantBWeComResult.SessionID) + assert.NotEqual(t, wecomResult.SessionID, groupResult.SessionID) + assert.NotEqual(t, wecomRunner.calls[0].userID, telegramRunner.calls[0].userID) + assert.NotEqual(t, wecomRunner.calls[0].userID, wecomRunner.calls[1].userID) + assertRunnerCall(t, wecomRunner.calls[0], wecomResult, wecomDM, "hello wecom") + assertRunnerCall(t, telegramRunner.calls[0], telegramResult, telegramDM, "hello telegram") + assertRunnerCall(t, wecomRunner.calls[1], tenantBWeComResult, tenantBWeComDM, "hello same channel") + assertRunnerCall(t, wecomRunner.calls[2], groupResult, wecomGroup, "hello group") + + due, err := outbox.ListDue(ctx, time.Now().Add(time.Hour), 10) + require.NoError(t, err) + require.Len(t, due, 4) + assertOutboundQueued(t, due, wecomResult.Outbound) + assertOutboundQueued(t, due, telegramResult.Outbound) + assertOutboundQueued(t, due, tenantBWeComResult.Outbound) + assertOutboundQueued(t, due, groupResult.Outbound) + records := audit.Records() + require.Len(t, records, 4) + for _, record := range records { + expectedUserHash := platform.UserIDHash(record.TenantID, record.Channel, "user-shared") + expectedInternalUserID := platform.InternalUserID(record.TenantID, record.Channel, "user-shared") + assert.Equal(t, "completed", record.Decision) + assert.NotEmpty(t, record.AuditID) + assert.NotEmpty(t, record.SessionID) + assert.Equal(t, expectedUserHash, record.UserID) + assert.Equal(t, expectedUserHash, record.UserIDHash) + assert.Equal(t, expectedInternalUserID, record.InternalUserID) + assert.NotContains(t, record.UserID, "user-shared") + assert.NotContains(t, record.UserIDHash, "user-shared") + assert.NotContains(t, record.InternalUserID, "user-shared") + } + assert.Equal(t, platform.IdempotencyStatusCompleted, wecomResult.Status) + assert.Equal(t, platform.IdempotencyStatusCompleted, telegramResult.Status) + assert.Equal(t, platform.IdempotencyStatusCompleted, tenantBWeComResult.Status) + assert.Equal(t, platform.IdempotencyStatusCompleted, groupResult.Status) +} + func TestServiceHandleInboundDuplicateReusesOutboxBackedResult(t *testing.T) { ctx := context.Background() registry := NewInMemoryRegistry() @@ -661,6 +814,41 @@ func validRuntime(tenantID string, r runnerStub) Runtime { } } +func validRuntimeForBinding( + tenantID string, + appID string, + bindingID string, + channel string, + accountID string, + r runnerStub, +) Runtime { + return Runtime{ + Tenant: platform.Tenant{ + TenantID: tenantID, + Status: platform.TenantStatusActive, + }, + App: platform.AgentApp{ + TenantID: tenantID, + AppID: appID, + AppName: appID, + Status: platform.AppStatusActive, + }, + Binding: platform.ChannelBinding{ + TenantID: tenantID, + AppID: appID, + BindingID: bindingID, + Channel: channel, + AccountID: accountID, + WebhookPath: "/webhook/" + bindingID, + TokenRef: "secret://token/" + bindingID, + SecretRef: "secret://secret/" + bindingID, + Status: platform.BindingStatusActive, + ChannelLimits: platform.ChannelLimits{MaxTextLength: 4096}, + }, + Runner: r, + } +} + func inbound(tenantID, messageID, userID, text string) platform.InboundMessage { return platform.InboundMessage{ TenantID: tenantID, @@ -679,6 +867,64 @@ func inbound(tenantID, messageID, userID, text string) platform.InboundMessage { } } +func inboundForRuntime( + tenantID string, + appID string, + bindingID string, + channel string, + accountID string, + messageID string, + userID string, + text string, +) platform.InboundMessage { + return platform.InboundMessage{ + TenantID: tenantID, + AppID: appID, + BindingID: bindingID, + Channel: channel, + ChannelAccountID: accountID, + PlatformMessageID: messageID, + ExternalUserID: userID, + ConversationType: platform.ConversationTypeDM, + MessageType: platform.MessageTypeText, + ContentParts: []platform.ContentPart{ + {Type: platform.ContentPartTypeText, Text: text}, + }, + ReceivedAt: time.Unix(100, 0), + } +} + +func assertOutboundQueued( + t *testing.T, + records []channeladapter.OutboxRecord, + outbound platform.OutboundMessage, +) { + t.Helper() + for _, record := range records { + if record.Message.DedupKey == outbound.DedupKey { + assert.Equal(t, platform.OutboundStatusPending, record.Status) + assert.Equal(t, outbound, record.Message) + return + } + } + t.Fatalf("outbound %q was not queued", outbound.DedupKey) +} + +func assertRunnerCall( + t *testing.T, + call runnerCall, + result Result, + msg platform.InboundMessage, + content string, +) { + t.Helper() + assert.Equal(t, result.SessionID, call.sessionID) + assert.Equal(t, result.RequestID, call.requestID) + assert.Equal(t, platform.InternalUserID(msg.TenantID, msg.Channel, msg.ExternalUserID), call.userID) + assert.Equal(t, model.RoleUser, call.message.Role) + assert.Equal(t, content, call.message.Content) +} + type runnerStub interface { Run( ctx context.Context, From b9bc7e37c1c2f66bdfe5b6c6e840c48eec1508e2 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Thu, 9 Jul 2026 10:18:17 +0800 Subject: [PATCH 41/95] platform/gateway: add outbound dispatch acceptance test --- platform/gateway/service_test.go | 100 +++++++++++++++++++++++++++++++ 1 file changed, 100 insertions(+) diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go index 9eec9126ce..7860b209b8 100644 --- a/platform/gateway/service_test.go +++ b/platform/gateway/service_test.go @@ -99,6 +99,86 @@ func TestServiceHandleInboundEnqueuesChannelOutbox(t *testing.T) { assert.Equal(t, 5, record.MaxAttempts) } +func TestServiceHandleInboundDispatchesOutboundToProvider(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "telegram reply"} + require.NoError(t, registry.Register(validRuntimeForBinding( + "tenant-a", + "app-telegram", + "binding-telegram", + "telegram", + "acct-telegram", + r, + ))) + outbox := channeladapter.NewInMemoryOutboxStore() + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewOutboxBackedOutboundStore(outbox), + ) + msg := inboundForRuntime( + "tenant-a", + "app-telegram", + "binding-telegram", + "telegram", + "acct-telegram", + "msg-telegram-dm", + "user-1", + "hello telegram", + ) + + result, err := svc.HandleInbound(ctx, msg) + require.NoError(t, err) + provider := &recordingOutboundProvider{ + status: platform.OutboundStatusSent, + providerMessageID: "telegram-provider-msg-1", + } + dispatcher := channeladapter.NewDispatcher( + outbox, + channeladapter.ProviderRegistryFunc(func(channel string) (channeladapter.OutboundProvider, bool) { + if channel != "telegram" { + return nil, false + } + return provider, true + }), + ) + + dispatchResults, err := dispatcher.DispatchDue(ctx, 10) + require.NoError(t, err) + + require.Len(t, dispatchResults, 1) + assert.Equal(t, result.Outbound.DedupKey, dispatchResults[0].DedupKey) + assert.Equal(t, platform.OutboundStatusSent, dispatchResults[0].Status) + assert.NoError(t, dispatchResults[0].Error) + require.Len(t, provider.delivered, 1) + delivered := provider.delivered[0] + expectedDedupKey := platform.IdempotencyKey( + "tenant-a", + "telegram", + "acct-telegram", + "msg-telegram-dm", + ) + ":outbound:1" + assert.Equal(t, "tenant-a", delivered.TenantID) + assert.Equal(t, "binding-telegram", delivered.BindingID) + assert.Equal(t, "telegram", delivered.Channel) + assert.Equal(t, result.SessionID, delivered.SessionID) + assert.Equal(t, "msg-telegram-dm", delivered.ReplyToPlatformMessageID) + assert.Equal(t, platform.OutboundMessageKindText, delivered.Kind) + assert.Equal(t, "telegram reply", delivered.Content) + assert.Equal(t, 1, delivered.Sequence) + assert.Equal(t, expectedDedupKey, delivered.DedupKey) + assert.Equal(t, result.RequestID, delivered.TraceID) + assert.Equal(t, result.Outbound, delivered) + record, ok, err := outbox.Get(ctx, result.Outbound.DedupKey) + require.NoError(t, err) + require.True(t, ok) + assert.Equal(t, platform.OutboundStatusSent, record.Status) + assert.Equal(t, "telegram-provider-msg-1", record.ProviderMessageID) + assert.NotNil(t, record.SentAt) + assert.Len(t, r.calls, 1) +} + func TestServiceHandleInboundCoversMinimumLoopAcceptance(t *testing.T) { ctx := context.Background() registry := NewInMemoryRegistry() @@ -950,6 +1030,26 @@ type recordingRunner struct { calls []runnerCall } +type recordingOutboundProvider struct { + status platform.OutboundStatus + providerMessageID string + delivered []platform.OutboundMessage +} + +func (p *recordingOutboundProvider) Deliver( + ctx context.Context, + msg platform.OutboundMessage, +) (channeladapter.DeliveryResult, error) { + if err := ctx.Err(); err != nil { + return channeladapter.DeliveryResult{}, err + } + p.delivered = append(p.delivered, msg) + return channeladapter.DeliveryResult{ + Status: p.status, + ProviderMessageID: p.providerMessageID, + }, nil +} + func (r *recordingRunner) Run( ctx context.Context, userID string, From 232e3bd550ac4ea0f8f290fa546ba27376b5b418 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Thu, 9 Jul 2026 10:26:01 +0800 Subject: [PATCH 42/95] platform/gateway: correlate audit trace ids --- platform/audit_record_test.go | 33 ++++++++++++++++++++++++++++++++ platform/gateway/service.go | 1 + platform/gateway/service_test.go | 11 ++++++++++- platform/validation.go | 8 ++++++++ 4 files changed, 52 insertions(+), 1 deletion(-) diff --git a/platform/audit_record_test.go b/platform/audit_record_test.go index f31de83dd7..661c971a46 100644 --- a/platform/audit_record_test.go +++ b/platform/audit_record_test.go @@ -80,6 +80,39 @@ func TestAuditRecordValidateRejectsSensitiveDecisionReason(t *testing.T) { } } +func TestAuditRecordValidateRejectsSensitiveRequestAndTraceIDs(t *testing.T) { + tests := []struct { + name string + mut func(*AuditRecord) + field string + }{ + { + name: "request_id", + mut: func(record *AuditRecord) { + record.RequestID = "Authorization: Bearer raw-token" + }, + field: "request_id", + }, + { + name: "trace_id", + mut: func(record *AuditRecord) { + record.TraceID = "password=plain" + }, + field: "trace_id", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + record := validAuditRecord() + tt.mut(&record) + if err := record.Validate(); err == nil || !strings.Contains(err.Error(), tt.field) { + t.Fatalf("expected sensitive %s rejection, got %v", tt.field, err) + } + }) + } +} + func TestAuditRecordValidateRejectsSensitiveErrorType(t *testing.T) { record := validAuditRecord() record.ErrorType = "storage_error password=plain" diff --git a/platform/gateway/service.go b/platform/gateway/service.go index 6dd611ed00..dc9061b143 100644 --- a/platform/gateway/service.go +++ b/platform/gateway/service.go @@ -423,6 +423,7 @@ func auditFromMessage( SessionID: sessionID, MessageID: msg.PlatformMessageID, RequestID: requestIDFor(msg), + TraceID: requestIDFor(msg), Decision: decision, DecisionReason: redactAuditReason(reason), LatencyMS: time.Since(start).Milliseconds(), diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go index 7860b209b8..18f35e40b8 100644 --- a/platform/gateway/service_test.go +++ b/platform/gateway/service_test.go @@ -775,7 +775,13 @@ func TestServiceHandleInboundUsesRequestIDAndStreamsText(t *testing.T) { registry := NewInMemoryRegistry() r := &recordingRunner{chunks: []string{"hel", "lo"}} registerRuntime(t, registry, "tenant-a", r) - svc := NewService(registry, platform.NewInMemoryIdempotencyStore(), NewInMemoryOutboundStore()) + audit := platform.NewInMemoryAuditSink() + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + WithAuditSink(audit), + ) msg := inbound("tenant-a", "msg-1", "user-1", "hello") msg.TraceContext = map[string]string{"request_id": "req-123"} @@ -786,6 +792,9 @@ func TestServiceHandleInboundUsesRequestIDAndStreamsText(t *testing.T) { assert.Equal(t, "req-123", r.calls[0].requestID) assert.Equal(t, "hello", result.Outbound.Content) assert.Equal(t, "req-123", result.Outbound.TraceID) + require.Len(t, audit.Records(), 1) + assert.Equal(t, "req-123", audit.Records()[0].RequestID) + assert.Equal(t, "req-123", audit.Records()[0].TraceID) } func TestRuntimeValidateRejectsIdentifierMismatch(t *testing.T) { diff --git a/platform/validation.go b/platform/validation.go index fcac81eb5d..1848b04bc7 100644 --- a/platform/validation.go +++ b/platform/validation.go @@ -327,6 +327,14 @@ func (r AuditRecord) Validate() error { if math.IsNaN(r.Cost) || math.IsInf(r.Cost, 0) || r.Cost < 0 { return fmt.Errorf("cost must be greater than or equal to 0") } + for field, value := range map[string]string{ + "request_id": r.RequestID, + "trace_id": r.TraceID, + } { + if err := validateAuditRedactedText(field, value); err != nil { + return err + } + } if err := validateAuditRedactedText("decision_reason", r.DecisionReason); err != nil { return err } From b2d5affc487a988ea6147be28ac99a695cfc0baa Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Thu, 9 Jul 2026 10:45:04 +0800 Subject: [PATCH 43/95] platform/gateway: add trace skeleton spans --- platform/gateway/service.go | 171 ++++++++++++++++++++++++++-- platform/gateway/service_test.go | 190 +++++++++++++++++++++++++++++++ 2 files changed, 351 insertions(+), 10 deletions(-) diff --git a/platform/gateway/service.go b/platform/gateway/service.go index dc9061b143..54e7e583a5 100644 --- a/platform/gateway/service.go +++ b/platform/gateway/service.go @@ -10,15 +10,23 @@ package gateway import ( "context" + "crypto/sha256" + "encoding/hex" + "errors" "fmt" "strings" "time" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/codes" + oteltrace "go.opentelemetry.io/otel/trace" + "trpc.group/trpc-go/trpc-agent-go/agent" "trpc.group/trpc-go/trpc-agent-go/event" "trpc.group/trpc-go/trpc-agent-go/model" "trpc.group/trpc-go/trpc-agent-go/platform" "trpc.group/trpc-go/trpc-agent-go/platform/channeladapter" + telemetrytrace "trpc.group/trpc-go/trpc-agent-go/telemetry/trace" ) // Service handles normalized inbound platform messages. @@ -97,29 +105,43 @@ func (s *Service) HandleInbound( msg platform.InboundMessage, ) (Result, error) { start := s.now() + ctx, callbackSpan := telemetrytrace.Tracer.Start(ctx, "im.callback") + defer callbackSpan.End() + setInboundTraceAttributes(callbackSpan, msg, "", "", "") if err := s.validateService(); err != nil { + recordSpanError(callbackSpan, err) return Result{}, err } if err := msg.Validate(); err != nil { s.writeAudit(ctx, auditFromMessage(msg, "", "", "reject", err.Error(), start, err)) + recordSpanError(callbackSpan, err) return Result{}, err } - runtime, ok, err := s.registry.Lookup(ctx, msg) + requestID := requestIDFor(msg) + setInboundTraceAttributes(callbackSpan, msg, "", requestID, "") + routeCtx, routeSpan := telemetrytrace.Tracer.Start(ctx, "gateway.route") + defer routeSpan.End() + setInboundTraceAttributes(routeSpan, msg, "", requestID, "") + runtime, ok, err := s.registry.Lookup(routeCtx, msg) if err != nil { + recordSpanError(routeSpan, err) return Result{}, err } if !ok { err := ErrRuntimeNotFound s.writeAudit(ctx, auditFromMessage(msg, "", "", "reject", err.Error(), start, err)) + recordSpanError(routeSpan, err) return Result{}, err } if err := runtime.Validate(); err != nil { s.writeAudit(ctx, auditFromMessage(msg, "", "", "reject", err.Error(), start, err)) + recordSpanError(routeSpan, err) return Result{}, err } if !runtime.matchesInbound(msg) { err := ErrRuntimeMismatch s.writeAudit(ctx, auditFromMessage(msg, "", "", "reject", err.Error(), start, err)) + recordSpanError(routeSpan, err) return Result{}, err } if err := authorizeBinding(runtime.Binding, msg); err != nil { @@ -129,36 +151,49 @@ func (s *Service) HandleInbound( text, err := inboundText(msg) if err != nil { s.writeAudit(ctx, auditFromMessage(msg, "", "", "reject", err.Error(), start, err)) + recordSpanError(routeSpan, err) return Result{}, err } sessionID, err := platform.SessionIDForInbound(msg) if err != nil { + recordSpanError(routeSpan, err) return Result{}, err } internalUserID := platform.InternalUserID(msg.TenantID, msg.Channel, msg.ExternalUserID) - requestID := requestIDFor(msg) + setInboundTraceAttributes(callbackSpan, msg, sessionID, requestID, internalUserID) + setInboundTraceAttributes(routeSpan, msg, sessionID, requestID, internalUserID) key := platform.IdempotencyKey( msg.TenantID, msg.Channel, msg.ChannelAccountID, msg.PlatformMessageID, ) - existing, ok, err := s.idempotencyStore.Get(ctx, key) + idempotencyCtx, idempotencySpan := telemetrytrace.Tracer.Start(routeCtx, "gateway.idempotency") + setInboundTraceAttributes(idempotencySpan, msg, sessionID, requestID, internalUserID) + existing, ok, err := s.idempotencyStore.Get(idempotencyCtx, key) if err != nil { + recordSpanError(idempotencySpan, err) + idempotencySpan.End() return Result{}, err } if ok { + idempotencySpan.End() return s.duplicateResult(ctx, existing) } - lease, acquired, err := s.leaseStore.Acquire(ctx, SessionLeaseKey{ + leaseCtx, leaseSpan := telemetrytrace.Tracer.Start(routeCtx, "gateway.session_lock") + setInboundTraceAttributes(leaseSpan, msg, sessionID, requestID, internalUserID) + lease, acquired, err := s.leaseStore.Acquire(leaseCtx, SessionLeaseKey{ TenantID: msg.TenantID, AppID: msg.AppID, SessionID: sessionID, }) if err != nil { + recordSpanError(leaseSpan, err) + leaseSpan.End() return Result{}, err } if !acquired { + leaseSpan.End() return Result{ RequestID: requestID, SessionID: sessionID, @@ -171,7 +206,8 @@ func (s *Service) HandleInbound( defer cancel() _ = lease.Release(cleanupCtx) }() - record, started, err := s.idempotencyStore.Start(ctx, platform.IdempotencyRecord{ + leaseSpan.End() + record, started, err := s.idempotencyStore.Start(idempotencyCtx, platform.IdempotencyRecord{ TenantID: msg.TenantID, Channel: msg.Channel, AccountID: msg.ChannelAccountID, @@ -181,14 +217,20 @@ func (s *Service) HandleInbound( SessionID: sessionID, }) if err != nil { + recordSpanError(idempotencySpan, err) + idempotencySpan.End() return Result{}, err } if !started { + idempotencySpan.End() return s.duplicateResult(ctx, record) } + idempotencySpan.End() + runnerCtx, runnerSpan := telemetrytrace.Tracer.Start(routeCtx, "runner.run") + setInboundTraceAttributes(runnerSpan, msg, sessionID, requestID, internalUserID) ch, err := runtime.Runner.Run( - ctx, + runnerCtx, internalUserID, sessionID, model.NewUserMessage(text), @@ -196,13 +238,18 @@ func (s *Service) HandleInbound( ) if err != nil { s.writeAudit(ctx, auditFromMessage(msg, sessionID, internalUserID, "runner_error", err.Error(), start, err)) + recordSpanError(runnerSpan, err) + runnerSpan.End() return Result{}, err } content, err := collectAssistantText(ctx, ch) if err != nil { s.writeAudit(ctx, auditFromMessage(msg, sessionID, internalUserID, "runner_error", err.Error(), start, err)) + recordSpanError(runnerSpan, err) + runnerSpan.End() return Result{}, err } + runnerSpan.End() resultRef := key + ":outbound:1" outbound := platform.OutboundMessage{ TenantID: msg.TenantID, @@ -216,25 +263,36 @@ func (s *Service) HandleInbound( DedupKey: resultRef, TraceID: requestID, } - if err := s.outboundStore.Save(ctx, resultRef, outbound); err != nil { + replyCtx, replySpan := telemetrytrace.Tracer.Start(routeCtx, "im.reply") + setInboundTraceAttributes(replySpan, msg, sessionID, requestID, internalUserID) + if err := s.outboundStore.Save(replyCtx, resultRef, outbound); err != nil { s.writeAudit(ctx, auditFromMessage(msg, sessionID, internalUserID, "outbound_error", err.Error(), start, err)) + recordSpanError(replySpan, err) + replySpan.End() return Result{}, err } if err := s.outboundStore.Enqueue( - ctx, + replyCtx, outbound, channeladapter.RetryPolicyForBinding(runtime.Binding), ); err != nil { - if _, markErr := s.idempotencyStore.MarkReplyFailed(ctx, key, resultRef); markErr != nil { + if _, markErr := s.idempotencyStore.MarkReplyFailed(replyCtx, key, resultRef); markErr != nil { + recordSpanError(replySpan, markErr) + replySpan.End() return Result{}, markErr } s.writeAudit(ctx, auditFromMessage(msg, sessionID, internalUserID, "outbound_error", err.Error(), start, err)) + recordSpanError(replySpan, err) + replySpan.End() return Result{}, err } - record, err = s.idempotencyStore.Complete(ctx, key, resultRef) + record, err = s.idempotencyStore.Complete(replyCtx, key, resultRef) if err != nil { + recordSpanError(replySpan, err) + replySpan.End() return Result{}, err } + replySpan.End() s.writeAudit(ctx, auditFromMessage(msg, sessionID, internalUserID, "completed", "", start, nil)) return Result{ RequestID: requestID, @@ -452,3 +510,96 @@ func (s *Service) writeAudit(ctx context.Context, record platform.AuditRecord) { } _ = s.auditSink.WriteAudit(ctx, record) } + +func setInboundTraceAttributes( + span interface{ SetAttributes(...attribute.KeyValue) }, + msg platform.InboundMessage, + sessionID string, + requestID string, + internalUserID string, +) { + attrs := []attribute.KeyValue{ + attribute.String("tenant_id", msg.TenantID), + attribute.String("app_id", msg.AppID), + attribute.String("channel", msg.Channel), + attribute.String("binding_id", msg.BindingID), + attribute.String("request_id_hash", traceSafeHash("request", requestID)), + attribute.String("user_id", platform.UserIDHash(msg.TenantID, msg.Channel, msg.ExternalUserID)), + attribute.String("user_id_hash", platform.UserIDHash(msg.TenantID, msg.Channel, msg.ExternalUserID)), + } + if sessionID != "" { + attrs = append(attrs, attribute.String("session_id_hash", traceSafeHash("session", sessionID))) + } + if internalUserID != "" { + attrs = append(attrs, attribute.String("internal_user_id_hash", traceSafeHash("internal_user", internalUserID))) + } + span.SetAttributes(attrs...) +} + +func traceSafeHash(scope string, value string) string { + value = strings.TrimSpace(value) + if value == "" { + return "" + } + sum := sha256.Sum256([]byte(scope + "\x00" + value)) + return scope + "_hash_" + hex.EncodeToString(sum[:])[:24] +} + +func recordSpanError(span oteltrace.Span, err error) { + if err == nil { + return + } + errType := traceErrorType(err) + span.RecordError(errors.New(errType)) + span.SetAttributes(attribute.String("error.type", errType)) + span.SetStatus(codes.Error, errType) +} + +func traceErrorType(err error) string { + switch { + case err == nil: + return "" + case errors.Is(err, context.Canceled): + return "context_canceled" + case errors.Is(err, context.DeadlineExceeded): + return "context_deadline_exceeded" + case errors.Is(err, platform.ErrTenantIDRequired): + return "tenant_id_required" + case errors.Is(err, platform.ErrAppIDRequired): + return "app_id_required" + case errors.Is(err, platform.ErrBindingIDRequired): + return "binding_id_required" + case errors.Is(err, platform.ErrChannelRequired): + return "channel_required" + case errors.Is(err, platform.ErrAccountIDRequired): + return "account_id_required" + case errors.Is(err, platform.ErrPlatformMessageIDRequired): + return "platform_message_id_required" + case errors.Is(err, platform.ErrExternalUserIDRequired): + return "external_user_id_required" + case errors.Is(err, platform.ErrExternalGroupIDRequired): + return "external_group_id_required" + case errors.Is(err, platform.ErrConversationTypeRequired): + return "conversation_type_required" + case errors.Is(err, platform.ErrInvalidConversationType): + return "invalid_conversation_type" + case errors.Is(err, ErrRuntimeNotFound): + return "runtime_not_found" + case errors.Is(err, ErrRuntimeInactive): + return "runtime_inactive" + case errors.Is(err, ErrRuntimeMismatch): + return "runtime_mismatch" + case errors.Is(err, ErrBindingAccessDenied): + return "binding_access_denied" + case errors.Is(err, ErrBindingMentionRequired): + return "binding_mention_required" + case errors.Is(err, ErrUnsupportedMessageType): + return "unsupported_message_type" + case errors.Is(err, ErrEmptyText): + return "empty_text" + case errors.Is(err, ErrRunnerResponseEmpty): + return "runner_response_empty" + default: + return "gateway_error" + } +} diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go index 18f35e40b8..045c51b1ba 100644 --- a/platform/gateway/service_test.go +++ b/platform/gateway/service_test.go @@ -11,18 +11,23 @@ package gateway import ( "context" "errors" + "slices" + "strings" "sync" "testing" "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + "go.opentelemetry.io/otel/sdk/trace/tracetest" "trpc.group/trpc-go/trpc-agent-go/agent" "trpc.group/trpc-go/trpc-agent-go/event" "trpc.group/trpc-go/trpc-agent-go/model" "trpc.group/trpc-go/trpc-agent-go/platform" "trpc.group/trpc-go/trpc-agent-go/platform/channeladapter" + telemetrytrace "trpc.group/trpc-go/trpc-agent-go/telemetry/trace" ) func TestServiceHandleInboundIsolatesTenants(t *testing.T) { @@ -797,6 +802,124 @@ func TestServiceHandleInboundUsesRequestIDAndStreamsText(t *testing.T) { assert.Equal(t, "req-123", audit.Records()[0].TraceID) } +func TestServiceHandleInboundEmitsTraceSkeleton(t *testing.T) { + recorder := useGatewaySpanRecorder(t) + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "trace reply"} + registerRuntime(t, registry, "tenant-a", r) + audit := platform.NewInMemoryAuditSink() + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + WithAuditSink(audit), + ) + msg := inbound("tenant-a", "msg-1", "external-user-raw", "hello secret-free trace") + msg.TraceContext = map[string]string{"request_id": "req-123"} + + result, err := svc.HandleInbound(ctx, msg) + require.NoError(t, err) + + spans := recorder.Ended() + require.Len(t, spans, 6) + assertSpanNames(t, spans, + "gateway.route", + "gateway.idempotency", + "gateway.session_lock", + "runner.run", + "im.reply", + "im.callback", + ) + callback := spanByName(t, spans, "im.callback") + route := spanByName(t, spans, "gateway.route") + expectedTraceID := callback.SpanContext().TraceID() + assert.Equal(t, expectedTraceID, route.SpanContext().TraceID()) + assert.Equal(t, callback.SpanContext().SpanID(), route.Parent().SpanID()) + for _, name := range []string{ + "gateway.idempotency", + "gateway.session_lock", + "runner.run", + "im.reply", + } { + span := spanByName(t, spans, name) + assert.Equal(t, expectedTraceID, span.SpanContext().TraceID(), name) + assert.Equal(t, route.SpanContext().SpanID(), span.Parent().SpanID(), name) + assert.Equal(t, "tenant-a", spanAttribute(t, span, "tenant_id"), name) + assert.Equal(t, "app", spanAttribute(t, span, "app_id"), name) + assert.Equal(t, "wecom", spanAttribute(t, span, "channel"), name) + assert.Equal(t, "binding", spanAttribute(t, span, "binding_id"), name) + assert.Equal(t, traceSafeHash("request", "req-123"), spanAttribute(t, span, "request_id_hash"), name) + assert.Equal(t, traceSafeHash("session", result.SessionID), spanAttribute(t, span, "session_id_hash"), name) + assert.Equal(t, platform.UserIDHash("tenant-a", "wecom", "external-user-raw"), spanAttribute(t, span, "user_id"), name) + assert.Empty(t, spanAttribute(t, span, "message")) + assert.Empty(t, spanAttribute(t, span, "content")) + assert.Empty(t, spanAttribute(t, span, "request_id")) + assert.Empty(t, spanAttribute(t, span, "session_id")) + assert.Empty(t, spanAttribute(t, span, "internal_user_id")) + assert.NotContains(t, spanAttributesText(span), "raw-token") + assert.NotContains(t, spanAttributesText(span), "external-user-raw") + assert.NotContains(t, spanAttributesText(span), result.SessionID) + } + assert.Equal(t, "completed", audit.Records()[0].Decision) + assert.Equal(t, "msg-1", audit.Records()[0].MessageID) +} + +func TestSetInboundTraceAttributesDoesNotExposeSensitiveIdentifiers(t *testing.T) { + recorder := useGatewaySpanRecorder(t) + msg := inbound("tenant-a", "msg-1", "external-user-raw", "hello") + ctx, span := telemetrytrace.Tracer.Start(context.Background(), "gateway.route") + setInboundTraceAttributes( + span, + msg, + "tenant:tenant-a:app:app:channel:wecom:dm:external-user-raw", + "Authorization: Bearer raw-token", + "usr_raw-internal", + ) + span.End() + _ = ctx + + ended := recorder.Ended() + require.Len(t, ended, 1) + attrs := spanAttributesText(ended[0]) + assert.Equal(t, traceSafeHash("request", "Authorization: Bearer raw-token"), spanAttribute(t, ended[0], "request_id_hash")) + assert.Equal(t, traceSafeHash("internal_user", "usr_raw-internal"), spanAttribute(t, ended[0], "internal_user_id_hash")) + assert.Empty(t, spanAttribute(t, ended[0], "request_id")) + assert.Empty(t, spanAttribute(t, ended[0], "session_id")) + assert.Empty(t, spanAttribute(t, ended[0], "internal_user_id")) + assert.NotContains(t, attrs, "raw-token") + assert.NotContains(t, attrs, "external-user-raw") + assert.NotContains(t, attrs, "usr_raw-internal") +} + +func TestServiceHandleInboundTraceErrorDoesNotExposeSensitiveError(t *testing.T) { + recorder := useGatewaySpanRecorder(t) + rawErr := errors.New("runner failed Authorization: Bearer raw-token api_key=sk-secret") + registry := NewInMemoryRegistry() + registerRuntime(t, registry, "tenant-a", &recordingRunner{runErr: rawErr}) + svc := NewService(registry, platform.NewInMemoryIdempotencyStore(), NewInMemoryOutboundStore()) + msg := inbound("tenant-a", "msg-1", "external-user-raw", "hello") + + _, err := svc.HandleInbound(context.Background(), msg) + require.Error(t, err) + assert.Contains(t, err.Error(), "raw-token") + + spans := recorder.Ended() + runnerSpan := spanByName(t, spans, "runner.run") + status := runnerSpan.Status() + assert.Equal(t, "gateway_error", status.Description) + assert.NotContains(t, status.Description, "raw-token") + assert.NotContains(t, status.Description, "sk-secret") + assert.Equal(t, "gateway_error", spanAttribute(t, runnerSpan, "error.type")) + + traceText := spanAttributesText(runnerSpan) + "\n" + spanEventsText(runnerSpan) + assert.NotContains(t, traceText, "raw-token") + assert.NotContains(t, traceText, "sk-secret") + assert.NotContains(t, traceText, "Authorization") + assert.NotContains(t, traceText, "api_key") + assert.Contains(t, traceText, "gateway_error") +} + func TestRuntimeValidateRejectsIdentifierMismatch(t *testing.T) { runtime := validRuntime("tenant-a", &recordingRunner{response: "unused"}) runtime.Binding.TenantID = "tenant-b" @@ -1291,3 +1414,70 @@ func requestIDFromOptions(opts ...agent.RunOption) string { } return runOptions.RequestID } + +func useGatewaySpanRecorder(t *testing.T) *tracetest.SpanRecorder { + t.Helper() + recorder := tracetest.NewSpanRecorder() + provider := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(recorder)) + originalProvider := telemetrytrace.TracerProvider + originalTracer := telemetrytrace.Tracer + telemetrytrace.TracerProvider = provider + telemetrytrace.Tracer = provider.Tracer("platform-gateway-test") + t.Cleanup(func() { + _ = provider.Shutdown(context.Background()) + telemetrytrace.TracerProvider = originalProvider + telemetrytrace.Tracer = originalTracer + }) + return recorder +} + +func assertSpanNames(t *testing.T, spans []sdktrace.ReadOnlySpan, want ...string) { + t.Helper() + got := make([]string, 0, len(spans)) + for _, span := range spans { + got = append(got, span.Name()) + } + for _, name := range want { + assert.True(t, slices.Contains(got, name), "missing span %q in %v", name, got) + } +} + +func spanByName(t *testing.T, spans []sdktrace.ReadOnlySpan, name string) sdktrace.ReadOnlySpan { + t.Helper() + for _, span := range spans { + if span.Name() == name { + return span + } + } + t.Fatalf("span %q not found", name) + return nil +} + +func spanAttribute(t *testing.T, span sdktrace.ReadOnlySpan, key string) string { + t.Helper() + for _, attr := range span.Attributes() { + if string(attr.Key) == key { + return attr.Value.AsString() + } + } + return "" +} + +func spanAttributesText(span sdktrace.ReadOnlySpan) string { + var values []string + for _, attr := range span.Attributes() { + values = append(values, string(attr.Key), attr.Value.AsString()) + } + return strings.Join(values, "\n") +} + +func spanEventsText(span sdktrace.ReadOnlySpan) string { + var values []string + for _, event := range span.Events() { + values = append(values, event.Name) + for _, attr := range event.Attributes { + values = append(values, string(attr.Key), attr.Value.AsString()) + } + } + return strings.Join(values, "\n") +} From 6d42695bf8015b3a6c9cd8a2008c17b04f81d238 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Thu, 9 Jul 2026 11:05:08 +0800 Subject: [PATCH 44/95] platform/gateway: enable runner session trace --- platform/gateway/service.go | 2 + platform/gateway/service_test.go | 57 ++++++++++++------- runner/diagnostics.go | 38 ++++++++++--- runner/runner_test.go | 96 +++++++++++++++++++++++++++++--- 4 files changed, 156 insertions(+), 37 deletions(-) diff --git a/platform/gateway/service.go b/platform/gateway/service.go index 54e7e583a5..1ea5b5ce88 100644 --- a/platform/gateway/service.go +++ b/platform/gateway/service.go @@ -235,6 +235,8 @@ func (s *Service) HandleInbound( sessionID, model.NewUserMessage(text), agent.WithRequestID(requestID), + agent.WithLatencyDiagnostics(true), + agent.WithLatencyDiagnosticsEvents(false), ) if err != nil { s.writeAudit(ctx, auditFromMessage(msg, sessionID, internalUserID, "runner_error", err.Error(), start, err)) diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go index 045c51b1ba..86e36cb122 100644 --- a/platform/gateway/service_test.go +++ b/platform/gateway/service_test.go @@ -795,6 +795,8 @@ func TestServiceHandleInboundUsesRequestIDAndStreamsText(t *testing.T) { require.Len(t, r.calls, 1) assert.Equal(t, "req-123", r.calls[0].requestID) + assert.True(t, r.calls[0].runOptions.LatencyDiagnosticsEnabled) + assert.False(t, r.calls[0].runOptions.LatencyDiagnosticsEmitEvents) assert.Equal(t, "hello", result.Outbound.Content) assert.Equal(t, "req-123", result.Outbound.TraceID) require.Len(t, audit.Records(), 1) @@ -1149,10 +1151,11 @@ type runnerStub interface { } type runnerCall struct { - userID string - sessionID string - message model.Message - requestID string + userID string + sessionID string + message model.Message + requestID string + runOptions agent.RunOptions } type recordingRunner struct { @@ -1192,11 +1195,13 @@ func (r *recordingRunner) Run( if r.runErr != nil { return nil, r.runErr } + runOptions := runOptionsFromOptions(runOpts...) r.calls = append(r.calls, runnerCall{ - userID: userID, - sessionID: sessionID, - message: message, - requestID: requestIDFromOptions(runOpts...), + userID: userID, + sessionID: sessionID, + message: message, + requestID: runOptions.RequestID, + runOptions: runOptions, }) out := make(chan *event.Event, 2) go func() { @@ -1255,11 +1260,13 @@ func (r *blockingRunner) Run( if r.done == nil { r.done = make(chan string, 1) } + runOptions := runOptionsFromOptions(runOpts...) r.calls = append(r.calls, runnerCall{ - userID: userID, - sessionID: sessionID, - message: message, - requestID: requestIDFromOptions(runOpts...), + userID: userID, + sessionID: sessionID, + message: message, + requestID: runOptions.RequestID, + runOptions: runOptions, }) r.startedOnce.Do(func() { close(r.started) @@ -1293,11 +1300,13 @@ func (r *cancelingRunner) Run( message model.Message, runOpts ...agent.RunOption, ) (<-chan *event.Event, error) { + runOptions := runOptionsFromOptions(runOpts...) r.calls = append(r.calls, runnerCall{ - userID: userID, - sessionID: sessionID, - message: message, - requestID: requestIDFromOptions(runOpts...), + userID: userID, + sessionID: sessionID, + message: message, + requestID: runOptions.RequestID, + runOptions: runOptions, }) r.cancel() return nil, r.runErr @@ -1324,11 +1333,13 @@ func (r *hangingFirstRunner) Run( ) (<-chan *event.Event, error) { r.mu.Lock() callIndex := len(r.calls) + runOptions := runOptionsFromOptions(runOpts...) r.calls = append(r.calls, runnerCall{ - userID: userID, - sessionID: sessionID, - message: message, - requestID: requestIDFromOptions(runOpts...), + userID: userID, + sessionID: sessionID, + message: message, + requestID: runOptions.RequestID, + runOptions: runOptions, }) if callIndex == 0 { r.startedOnce.Do(func() { @@ -1406,13 +1417,17 @@ func chunkEvent(content string, partial bool) *event.Event { } func requestIDFromOptions(opts ...agent.RunOption) string { + return runOptionsFromOptions(opts...).RequestID +} + +func runOptionsFromOptions(opts ...agent.RunOption) agent.RunOptions { var runOptions agent.RunOptions for _, opt := range opts { if opt != nil { opt(&runOptions) } } - return runOptions.RequestID + return runOptions } func useGatewaySpanRecorder(t *testing.T) *tracetest.SpanRecorder { diff --git a/runner/diagnostics.go b/runner/diagnostics.go index d3e31637e4..a697baad1e 100644 --- a/runner/diagnostics.go +++ b/runner/diagnostics.go @@ -10,6 +10,11 @@ package runner import ( "context" + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + "strings" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/codes" @@ -95,8 +100,9 @@ func finishRunnerLatencySpan(span oteltrace.Span, started bool, err error) { return } if err != nil { - span.RecordError(err) - span.SetStatus(codes.Error, err.Error()) + description := runnerTraceErrorDescription(err) + span.RecordError(errors.New(description)) + span.SetStatus(codes.Error, description) } span.End() } @@ -110,9 +116,9 @@ func runnerRunAttrs( ) []attribute.KeyValue { return []attribute.KeyValue{ attribute.String("runner.app", appName), - attribute.String("runner.user_id", userID), - attribute.String("runner.session_id", sessionID), - attribute.String("runner.request_id", ro.RequestID), + attribute.String("runner.user_id_hash", runnerTraceSafeHash("user", userID)), + attribute.String("runner.session_id_hash", runnerTraceSafeHash("session", sessionID)), + attribute.String("runner.request_id_hash", runnerTraceSafeHash("request", ro.RequestID)), attribute.String("runner.message.role", string(message.Role)), attribute.Bool("runner.message.has_payload", model.HasPayload(message)), attribute.Int("runner.options.seed_messages", len(ro.Messages)), @@ -126,15 +132,15 @@ func runnerInvocationAttrs(inv *agent.Invocation) []attribute.KeyValue { return []attribute.KeyValue{ attribute.String("runner.invocation_id", inv.InvocationID), attribute.String("runner.agent", inv.AgentName), - attribute.String("runner.request_id", inv.RunOptions.RequestID), + attribute.String("runner.request_id_hash", runnerTraceSafeHash("request", inv.RunOptions.RequestID)), } } func runnerSessionAttrs(key session.Key, sess *session.Session) []attribute.KeyValue { attrs := []attribute.KeyValue{ attribute.String("runner.session.app", key.AppName), - attribute.String("runner.session.user", key.UserID), - attribute.String("runner.session.id", key.SessionID), + attribute.String("runner.session.user_hash", runnerTraceSafeHash("user", key.UserID)), + attribute.String("runner.session.id_hash", runnerTraceSafeHash("session", key.SessionID)), } if sess != nil { attrs = append( @@ -153,6 +159,22 @@ func runnerSessionAttrs(key session.Key, sess *session.Session) []attribute.KeyV return attrs } +func runnerTraceSafeHash(scope string, value string) string { + value = strings.TrimSpace(value) + if value == "" { + return "" + } + sum := sha256.Sum256([]byte(scope + "\x00" + value)) + return scope + "_hash_" + hex.EncodeToString(sum[:])[:24] +} + +func runnerTraceErrorDescription(err error) string { + if err == nil { + return "" + } + return fmt.Sprintf("%T", err) +} + func runnerSessionSummaryCount(sess *session.Session) int { if sess == nil { return 0 diff --git a/runner/runner_test.go b/runner/runner_test.go index e0735fac55..60d12eaaaa 100644 --- a/runner/runner_test.go +++ b/runner/runner_test.go @@ -26,6 +26,8 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.opentelemetry.io/otel/attribute" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + "go.opentelemetry.io/otel/sdk/trace/tracetest" "trpc.group/trpc-go/trpc-agent-go/agent" "trpc.group/trpc-go/trpc-agent-go/agent/chainagent" @@ -49,6 +51,7 @@ import ( "trpc.group/trpc-go/trpc-agent-go/session" sessioninmemory "trpc.group/trpc-go/trpc-agent-go/session/inmemory" "trpc.group/trpc-go/trpc-agent-go/skill" + telemetrytrace "trpc.group/trpc-go/trpc-agent-go/telemetry/trace" "trpc.group/trpc-go/trpc-agent-go/tool" "trpc.group/trpc-go/trpc-agent-go/tool/function" ) @@ -4432,6 +4435,7 @@ func TestRunner_Run_AgentRunError(t *testing.T) { } func TestRunnerLatencyDiagnosticHelpers(t *testing.T) { + recorder := useRunnerSpanRecorder(t) ctx := context.Background() _, disabledSpan, disabledStarted := startRunnerLatencySpan( ctx, @@ -4466,7 +4470,8 @@ func TestRunnerLatencyDiagnosticHelpers(t *testing.T) { attribute.String("test.attr", "value"), ) require.True(t, started) - finishRunnerLatencySpan(span, started, errors.New("boom")) + rawError := errors.New("raw-user-id raw-session-id raw-request-id") + finishRunnerLatencySpan(span, started, rawError) _, optionSpan, optionStarted := startRunnerRunOptionsLatencySpan( ctx, @@ -4485,36 +4490,50 @@ func TestRunnerLatencyDiagnosticHelpers(t *testing.T) { require.False(t, optionStarted) finishRunnerLatencySpan(optionSpan, optionStarted, nil) + rawUserID := "raw-user-id" + rawSessionID := "raw-session-id" + rawRequestID := "raw-request-id" + runAttrs := runnerRunAttrs( "app", - "user", - "sess", + rawUserID, + rawSessionID, model.NewUserMessage("hello"), agent.RunOptions{ - RequestID: "req-latency", + RequestID: rawRequestID, Messages: []model.Message{model.NewSystemMessage("seed")}, }, ) require.True(t, runnerHasAttr(runAttrs, "runner.app", "app")) - require.True(t, runnerHasAttr(runAttrs, "runner.request_id", "req-latency")) + require.True(t, runnerHasAttr(runAttrs, "runner.user_id_hash", runnerTraceSafeHash("user", rawUserID))) + require.True(t, runnerHasAttr(runAttrs, "runner.session_id_hash", runnerTraceSafeHash("session", rawSessionID))) + require.True(t, runnerHasAttr(runAttrs, "runner.request_id_hash", runnerTraceSafeHash("request", rawRequestID))) require.True(t, runnerHasAttr(runAttrs, "runner.message.role", "user")) require.True(t, runnerHasAttr(runAttrs, "runner.message.has_payload", true)) require.True(t, runnerHasAttr(runAttrs, "runner.options.seed_messages", 1)) + require.NotContains(t, runnerAttrsText(runAttrs), rawUserID) + require.NotContains(t, runnerAttrsText(runAttrs), rawSessionID) + require.NotContains(t, runnerAttrsText(runAttrs), rawRequestID) invAttrs := runnerInvocationAttrs(inv) require.True(t, runnerHasAttr(invAttrs, "runner.agent", "a")) - require.True(t, runnerHasAttr(invAttrs, "runner.request_id", "req-latency")) + require.True(t, runnerHasAttr(invAttrs, "runner.request_id_hash", runnerTraceSafeHash("request", "req-latency"))) + require.NotContains(t, runnerAttrsText(invAttrs), "req-latency") require.Nil(t, runnerInvocationAttrs(nil)) - sess := session.NewSession("app", "user", "sess") + sess := session.NewSession("app", rawUserID, rawSessionID) sess.SetState("state-key", []byte("value")) sess.Summaries = map[string]*session.Summary{"default": {}} - key := session.Key{AppName: "app", UserID: "user", SessionID: "sess"} + key := session.Key{AppName: "app", UserID: rawUserID, SessionID: rawSessionID} sessionAttrs := runnerSessionAttrs(key, sess) require.True(t, runnerHasAttr(sessionAttrs, "runner.session.app", "app")) + require.True(t, runnerHasAttr(sessionAttrs, "runner.session.user_hash", runnerTraceSafeHash("user", rawUserID))) + require.True(t, runnerHasAttr(sessionAttrs, "runner.session.id_hash", runnerTraceSafeHash("session", rawSessionID))) require.True(t, runnerHasAttr(sessionAttrs, "runner.session.events", 0)) require.True(t, runnerHasAttr(sessionAttrs, runnerAttrSessionStateKeys, 1)) require.True(t, runnerHasAttr(sessionAttrs, runnerAttrSessionSummaryKeys, 1)) + require.NotContains(t, runnerAttrsText(sessionAttrs), rawUserID) + require.NotContains(t, runnerAttrsText(sessionAttrs), rawSessionID) evt := event.New( inv.InvocationID, @@ -4554,6 +4573,13 @@ func TestRunnerLatencyDiagnosticHelpers(t *testing.T) { Error: &model.ResponseError{Type: model.ErrorTypeRunError}, }, })) + + errorSpan := runnerSpanByName(t, recorder.Ended(), runnerLatencySpanProcessEvent) + require.Equal(t, runnerTraceErrorDescription(rawError), errorSpan.Status().Description) + errorTraceText := runnerSpanAttributesText(errorSpan) + "\n" + runnerSpanEventsText(errorSpan) + require.NotContains(t, errorTraceText, rawUserID) + require.NotContains(t, errorTraceText, rawSessionID) + require.NotContains(t, errorTraceText, rawRequestID) } func runnerHasAttr(attrs []attribute.KeyValue, key string, want any) bool { @@ -4566,6 +4592,60 @@ func runnerHasAttr(attrs []attribute.KeyValue, key string, want any) bool { return false } +func runnerAttrsText(attrs []attribute.KeyValue) string { + var values []string + for _, attr := range attrs { + values = append(values, fmt.Sprint(attr.Value.AsInterface())) + } + return strings.Join(values, "\n") +} + +func useRunnerSpanRecorder(t *testing.T) *tracetest.SpanRecorder { + t.Helper() + recorder := tracetest.NewSpanRecorder() + provider := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(recorder)) + originalProvider := telemetrytrace.TracerProvider + originalTracer := telemetrytrace.Tracer + telemetrytrace.TracerProvider = provider + telemetrytrace.Tracer = provider.Tracer("runner-test") + t.Cleanup(func() { + _ = provider.Shutdown(context.Background()) + telemetrytrace.TracerProvider = originalProvider + telemetrytrace.Tracer = originalTracer + }) + return recorder +} + +func runnerSpanByName(t *testing.T, spans []sdktrace.ReadOnlySpan, name string) sdktrace.ReadOnlySpan { + t.Helper() + for _, span := range spans { + if span.Name() == name { + return span + } + } + t.Fatalf("span %q not found", name) + return nil +} + +func runnerSpanAttributesText(span sdktrace.ReadOnlySpan) string { + var values []string + for _, attr := range span.Attributes() { + values = append(values, string(attr.Key), attr.Value.AsString()) + } + return strings.Join(values, "\n") +} + +func runnerSpanEventsText(span sdktrace.ReadOnlySpan) string { + var values []string + for _, event := range span.Events() { + values = append(values, event.Name) + for _, attr := range event.Attributes { + values = append(values, string(attr.Key), attr.Value.AsString()) + } + } + return strings.Join(values, "\n") +} + func TestGetOrCreateSession_Existing(t *testing.T) { // Pre-create a session; getOrCreateSession should return it without creating a new one. svc := sessioninmemory.NewSessionService() From 054d83b7b007e394e90622ed80ed5d0c9b495716 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Thu, 9 Jul 2026 11:15:27 +0800 Subject: [PATCH 45/95] platform/gateway: add message event trace contract --- platform/gateway/service.go | 61 +++++++++++++ platform/gateway/service_test.go | 46 +++++++++- platform/message_event.go | 54 ++++++++++++ platform/message_event_test.go | 145 +++++++++++++++++++++++++++++++ platform/types.go | 49 +++++++++++ platform/validation.go | 55 ++++++++++++ 6 files changed, 408 insertions(+), 2 deletions(-) create mode 100644 platform/message_event.go create mode 100644 platform/message_event_test.go diff --git a/platform/gateway/service.go b/platform/gateway/service.go index 1ea5b5ce88..b0592a14ff 100644 --- a/platform/gateway/service.go +++ b/platform/gateway/service.go @@ -36,6 +36,7 @@ type Service struct { outboundStore OutboundStore leaseStore SessionLeaseStore auditSink platform.AuditSink + messageEventSink platform.MessageEventSink now func() time.Time } @@ -49,6 +50,13 @@ func WithAuditSink(sink platform.AuditSink) Option { } } +// WithMessageEventSink sets the message event sink used by the service. +func WithMessageEventSink(sink platform.MessageEventSink) Option { + return func(s *Service) { + s.messageEventSink = sink + } +} + // WithNow sets the clock used by the service. func WithNow(now func() time.Time) Option { return func(s *Service) { @@ -295,6 +303,8 @@ func (s *Service) HandleInbound( return Result{}, err } replySpan.End() + s.writeMessageEvent(ctx, messageEventFromInbound(msg, sessionID, key, requestID, 1, start)) + s.writeMessageEvent(ctx, messageEventFromAssistant(msg, sessionID, resultRef, requestID, 2, s.now())) s.writeAudit(ctx, auditFromMessage(msg, sessionID, internalUserID, "completed", "", start, nil)) return Result{ RequestID: requestID, @@ -513,6 +523,57 @@ func (s *Service) writeAudit(ctx context.Context, record platform.AuditRecord) { _ = s.auditSink.WriteAudit(ctx, record) } +func (s *Service) writeMessageEvent(ctx context.Context, event platform.MessageEvent) { + if s.messageEventSink == nil { + return + } + _ = s.messageEventSink.WriteMessageEvent(ctx, event) +} + +func messageEventFromInbound( + msg platform.InboundMessage, + sessionID string, + idempotencyKey string, + traceID string, + sequence int64, + createdAt time.Time, +) platform.MessageEvent { + return platform.MessageEvent{ + TenantID: msg.TenantID, + AppID: msg.AppID, + SessionID: sessionID, + EventID: idempotencyKey + ":user", + Sequence: sequence, + IdempotencyKey: idempotencyKey, + Role: platform.MessageEventRoleUser, + EventType: platform.MessageEventTypeMessage, + TraceID: traceID, + CreatedAt: createdAt, + } +} + +func messageEventFromAssistant( + msg platform.InboundMessage, + sessionID string, + resultRef string, + traceID string, + sequence int64, + createdAt time.Time, +) platform.MessageEvent { + return platform.MessageEvent{ + TenantID: msg.TenantID, + AppID: msg.AppID, + SessionID: sessionID, + EventID: resultRef + ":assistant", + Sequence: sequence, + IdempotencyKey: resultRef, + Role: platform.MessageEventRoleAssistant, + EventType: platform.MessageEventTypeMessage, + TraceID: traceID, + CreatedAt: createdAt, + } +} + func setInboundTraceAttributes( span interface{ SetAttributes(...attribute.KeyValue) }, msg platform.InboundMessage, diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go index 86e36cb122..30382f0da4 100644 --- a/platform/gateway/service_test.go +++ b/platform/gateway/service_test.go @@ -216,11 +216,13 @@ func TestServiceHandleInboundCoversMinimumLoopAcceptance(t *testing.T) { idempotency := platform.NewInMemoryIdempotencyStore() outbox := channeladapter.NewInMemoryOutboxStore() audit := platform.NewInMemoryAuditSink() + messageEvents := platform.NewInMemoryMessageEventSink() svc := NewService( registry, idempotency, NewOutboxBackedOutboundStore(outbox), WithAuditSink(audit), + WithMessageEventSink(messageEvents), ) wecomDM := inboundForRuntime( "tenant-a", @@ -318,6 +320,12 @@ func TestServiceHandleInboundCoversMinimumLoopAcceptance(t *testing.T) { assertOutboundQueued(t, due, groupResult.Outbound) records := audit.Records() require.Len(t, records, 4) + events := messageEvents.Events() + require.Len(t, events, 8) + eventsByTraceID := make(map[string][]platform.MessageEvent) + for _, event := range events { + eventsByTraceID[event.TraceID] = append(eventsByTraceID[event.TraceID], event) + } for _, record := range records { expectedUserHash := platform.UserIDHash(record.TenantID, record.Channel, "user-shared") expectedInternalUserID := platform.InternalUserID(record.TenantID, record.Channel, "user-shared") @@ -330,6 +338,12 @@ func TestServiceHandleInboundCoversMinimumLoopAcceptance(t *testing.T) { assert.NotContains(t, record.UserID, "user-shared") assert.NotContains(t, record.UserIDHash, "user-shared") assert.NotContains(t, record.InternalUserID, "user-shared") + traceEvents := eventsByTraceID[record.TraceID] + require.Len(t, traceEvents, 2) + assert.Equal(t, record.SessionID, traceEvents[0].SessionID) + assert.Equal(t, record.SessionID, traceEvents[1].SessionID) + assert.Equal(t, platform.MessageEventRoleUser, traceEvents[0].Role) + assert.Equal(t, platform.MessageEventRoleAssistant, traceEvents[1].Role) } assert.Equal(t, platform.IdempotencyStatusCompleted, wecomResult.Status) assert.Equal(t, platform.IdempotencyStatusCompleted, telegramResult.Status) @@ -371,6 +385,7 @@ func TestServiceHandleInboundOutboxFailureDoesNotCompleteIdempotency(t *testing. idempotency := platform.NewInMemoryIdempotencyStore() outbox := channeladapter.NewInMemoryOutboxStore() store := NewOutboxBackedOutboundStore(outbox) + messageEvents := platform.NewInMemoryMessageEventSink() msg := inbound("tenant-a", "msg-1", "user-1", "hello") resultRef := platform.IdempotencyKey("tenant-a", "wecom", "acct", "msg-1") + ":outbound:1" colliding := platform.OutboundMessage{ @@ -387,7 +402,13 @@ func TestServiceHandleInboundOutboxFailureDoesNotCompleteIdempotency(t *testing. _, _, err := outbox.Enqueue(ctx, colliding, channeladapter.DefaultRetryPolicy()) require.NoError(t, err) audit := platform.NewInMemoryAuditSink() - svc := NewService(registry, idempotency, store, WithAuditSink(audit)) + svc := NewService( + registry, + idempotency, + store, + WithAuditSink(audit), + WithMessageEventSink(messageEvents), + ) _, err = svc.HandleInbound(ctx, msg) @@ -404,6 +425,7 @@ func TestServiceHandleInboundOutboxFailureDoesNotCompleteIdempotency(t *testing. require.Len(t, audit.Records(), 1) assert.NotEmpty(t, audit.Records()[0].AuditID) assert.Equal(t, "outbound_error", audit.Records()[0].Decision) + assert.Empty(t, messageEvents.Events()) } func TestServiceHandleInboundDuplicateReplyFailedReusesStoredOutbound(t *testing.T) { @@ -735,7 +757,13 @@ func TestServiceHandleInboundRunnerErrorDoesNotComplete(t *testing.T) { r := &recordingRunner{runErr: runnerErr} registerRuntime(t, registry, "tenant-a", r) store := platform.NewInMemoryIdempotencyStore() - svc := NewService(registry, store, NewInMemoryOutboundStore()) + messageEvents := platform.NewInMemoryMessageEventSink() + svc := NewService( + registry, + store, + NewInMemoryOutboundStore(), + WithMessageEventSink(messageEvents), + ) msg := inbound("tenant-a", "msg-1", "user-1", "hello") _, err := svc.HandleInbound(ctx, msg) @@ -746,6 +774,7 @@ func TestServiceHandleInboundRunnerErrorDoesNotComplete(t *testing.T) { require.True(t, ok) assert.Equal(t, platform.IdempotencyStatusProcessing, record.Status) assert.Empty(t, record.ResultRef) + assert.Empty(t, messageEvents.Events()) } func TestServiceHandleInboundRunnerErrorRedactsAuditReason(t *testing.T) { @@ -781,11 +810,13 @@ func TestServiceHandleInboundUsesRequestIDAndStreamsText(t *testing.T) { r := &recordingRunner{chunks: []string{"hel", "lo"}} registerRuntime(t, registry, "tenant-a", r) audit := platform.NewInMemoryAuditSink() + messageEvents := platform.NewInMemoryMessageEventSink() svc := NewService( registry, platform.NewInMemoryIdempotencyStore(), NewInMemoryOutboundStore(), WithAuditSink(audit), + WithMessageEventSink(messageEvents), ) msg := inbound("tenant-a", "msg-1", "user-1", "hello") msg.TraceContext = map[string]string{"request_id": "req-123"} @@ -802,6 +833,17 @@ func TestServiceHandleInboundUsesRequestIDAndStreamsText(t *testing.T) { require.Len(t, audit.Records(), 1) assert.Equal(t, "req-123", audit.Records()[0].RequestID) assert.Equal(t, "req-123", audit.Records()[0].TraceID) + events := messageEvents.Events() + require.Len(t, events, 2) + assert.Equal(t, "req-123", events[0].TraceID) + assert.Equal(t, "req-123", events[1].TraceID) + assert.Equal(t, result.SessionID, events[0].SessionID) + assert.Equal(t, result.SessionID, events[1].SessionID) + assert.Equal(t, platform.MessageEventRoleUser, events[0].Role) + assert.Equal(t, platform.MessageEventRoleAssistant, events[1].Role) + assert.Equal(t, result.Outbound.TraceID, audit.Records()[0].TraceID) + assert.Equal(t, result.Outbound.TraceID, events[0].TraceID) + assert.Equal(t, result.Outbound.TraceID, events[1].TraceID) } func TestServiceHandleInboundEmitsTraceSkeleton(t *testing.T) { diff --git a/platform/message_event.go b/platform/message_event.go new file mode 100644 index 0000000000..1509e8083e --- /dev/null +++ b/platform/message_event.go @@ -0,0 +1,54 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "context" + "sync" +) + +// MessageEventSink stores immutable conversation events. +type MessageEventSink interface { + // WriteMessageEvent writes one message event. + WriteMessageEvent(ctx context.Context, event MessageEvent) error +} + +// InMemoryMessageEventSink is a concurrency-safe message event sink for tests and demos. +type InMemoryMessageEventSink struct { + mu sync.Mutex + events []MessageEvent +} + +// NewInMemoryMessageEventSink creates an in-memory message event sink. +func NewInMemoryMessageEventSink() *InMemoryMessageEventSink { + return &InMemoryMessageEventSink{} +} + +// WriteMessageEvent writes one message event. +func (s *InMemoryMessageEventSink) WriteMessageEvent(ctx context.Context, event MessageEvent) error { + if err := ctx.Err(); err != nil { + return err + } + if err := event.Validate(); err != nil { + return err + } + s.mu.Lock() + defer s.mu.Unlock() + s.events = append(s.events, event) + return nil +} + +// Events returns a snapshot of written message events. +func (s *InMemoryMessageEventSink) Events() []MessageEvent { + s.mu.Lock() + defer s.mu.Unlock() + out := make([]MessageEvent, len(s.events)) + copy(out, s.events) + return out +} diff --git a/platform/message_event_test.go b/platform/message_event_test.go new file mode 100644 index 0000000000..5e23f5f5e8 --- /dev/null +++ b/platform/message_event_test.go @@ -0,0 +1,145 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "context" + "strings" + "testing" + "time" +) + +func TestMessageEventValidateAcceptsSafeRecord(t *testing.T) { + event := validMessageEvent() + event.ContentJSON = `{"text":"hello"}` + event.MetadataJSON = `{"source":"gateway"}` + + if err := event.Validate(); err != nil { + t.Fatalf("expected valid message event, got %v", err) + } +} + +func TestMessageEventValidateRequiresIdentity(t *testing.T) { + tests := []struct { + name string + mutate func(*MessageEvent) + want string + }{ + {name: "tenant", mutate: func(e *MessageEvent) { e.TenantID = " " }, want: "tenant_id"}, + {name: "app", mutate: func(e *MessageEvent) { e.AppID = " " }, want: "app_id"}, + {name: "session", mutate: func(e *MessageEvent) { e.SessionID = " " }, want: "session_id"}, + {name: "event", mutate: func(e *MessageEvent) { e.EventID = " " }, want: "event_id"}, + {name: "idempotency", mutate: func(e *MessageEvent) { e.IdempotencyKey = " " }, want: "idempotency_key"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + event := validMessageEvent() + tt.mutate(&event) + if err := event.Validate(); err == nil || !strings.Contains(err.Error(), tt.want) { + t.Fatalf("expected %s validation, got %v", tt.want, err) + } + }) + } +} + +func TestMessageEventValidateRejectsInvalidRoleTypeAndSequence(t *testing.T) { + tests := []struct { + name string + mutate func(*MessageEvent) + want string + }{ + {name: "role", mutate: func(e *MessageEvent) { e.Role = "admin" }, want: "role"}, + {name: "type", mutate: func(e *MessageEvent) { e.EventType = "secret" }, want: "event_type"}, + {name: "sequence", mutate: func(e *MessageEvent) { e.Sequence = 0 }, want: "sequence"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + event := validMessageEvent() + tt.mutate(&event) + if err := event.Validate(); err == nil || !strings.Contains(err.Error(), tt.want) { + t.Fatalf("expected %s validation, got %v", tt.want, err) + } + }) + } +} + +func TestMessageEventValidateRejectsSensitiveTraceFields(t *testing.T) { + tests := []struct { + name string + mutate func(*MessageEvent) + want string + }{ + {name: "app", mutate: func(e *MessageEvent) { e.AppID = "api_key=sk-1234567890abcdef" }, want: "app_id"}, + {name: "session", mutate: func(e *MessageEvent) { e.SessionID = "Authorization: Bearer raw-token" }, want: "session_id"}, + {name: "event", mutate: func(e *MessageEvent) { e.EventID = "password=plain" }, want: "event_id"}, + {name: "idempotency", mutate: func(e *MessageEvent) { e.IdempotencyKey = "token=raw-token" }, want: "idempotency_key"}, + {name: "trace", mutate: func(e *MessageEvent) { e.TraceID = "Authorization: Bearer raw-token" }, want: "trace_id"}, + {name: "content", mutate: func(e *MessageEvent) { e.ContentJSON = `{"api_key":"sk-1234567890abcdef"}` }, want: "content_json"}, + {name: "tool_calls", mutate: func(e *MessageEvent) { e.ToolCallsJSON = `{"password":"plain"}` }, want: "tool_calls_json"}, + {name: "metadata", mutate: func(e *MessageEvent) { e.MetadataJSON = `{"token":"raw-token"}` }, want: "metadata_json"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + event := validMessageEvent() + tt.mutate(&event) + if err := event.Validate(); err == nil || !strings.Contains(err.Error(), tt.want) { + t.Fatalf("expected %s validation, got %v", tt.want, err) + } + }) + } +} + +func TestMessageEventSinkStoresSnapshot(t *testing.T) { + sink := NewInMemoryMessageEventSink() + event := validMessageEvent() + + if err := sink.WriteMessageEvent(context.Background(), event); err != nil { + t.Fatalf("WriteMessageEvent: %v", err) + } + events := sink.Events() + if len(events) != 1 { + t.Fatalf("expected one message event, got %d", len(events)) + } + events[0].TenantID = "changed" + if sink.Events()[0].TenantID != "tenant" { + t.Fatalf("Events should return a defensive copy") + } +} + +func TestMessageEventSinkRejectsInvalidRecord(t *testing.T) { + sink := NewInMemoryMessageEventSink() + event := validMessageEvent() + event.TraceID = "password=plain" + + err := sink.WriteMessageEvent(context.Background(), event) + if err == nil || !strings.Contains(err.Error(), "trace_id") { + t.Fatalf("expected trace validation, got %v", err) + } + if got := sink.Events(); len(got) != 0 { + t.Fatalf("expected invalid event to be rejected, got %+v", got) + } +} + +func validMessageEvent() MessageEvent { + return MessageEvent{ + TenantID: "tenant", + AppID: "app", + SessionID: "session", + EventID: "event", + Sequence: 1, + IdempotencyKey: "idempotency", + Role: MessageEventRoleUser, + EventType: MessageEventTypeMessage, + TraceID: "trace", + CreatedAt: time.Now(), + } +} diff --git a/platform/types.go b/platform/types.go index 2a57b51f5c..8daa5d52d5 100644 --- a/platform/types.go +++ b/platform/types.go @@ -132,6 +132,38 @@ const ( OutboundMessageKindStatus OutboundMessageKind = "status" ) +// MessageEventRole describes the normalized message actor. +type MessageEventRole string + +const ( + // MessageEventRoleUser records an inbound user message. + MessageEventRoleUser MessageEventRole = "user" + // MessageEventRoleAssistant records an assistant reply message. + MessageEventRoleAssistant MessageEventRole = "assistant" + // MessageEventRoleTool records a tool event. + MessageEventRoleTool MessageEventRole = "tool" + // MessageEventRoleSystem records a system event. + MessageEventRoleSystem MessageEventRole = "system" +) + +// MessageEventType describes the immutable conversation event kind. +type MessageEventType string + +const ( + // MessageEventTypeMessage records a normal conversational message. + MessageEventTypeMessage MessageEventType = "message" + // MessageEventTypeToolCall records a tool call request. + MessageEventTypeToolCall MessageEventType = "tool_call" + // MessageEventTypeToolResult records a tool call result. + MessageEventTypeToolResult MessageEventType = "tool_result" + // MessageEventTypeError records an execution error event. + MessageEventTypeError MessageEventType = "error" + // MessageEventTypeRevoke records a revoked prior event. + MessageEventTypeRevoke MessageEventType = "revoke" + // MessageEventTypeEdit records an edited prior event. + MessageEventTypeEdit MessageEventType = "edit" +) + // IdempotencyStatus is the state of one inbound platform message. type IdempotencyStatus string @@ -381,6 +413,23 @@ type OutboundMessage struct { TraceID string } +// MessageEvent stores immutable conversation event metadata for trace correlation. +type MessageEvent struct { + TenantID string + AppID string + SessionID string + EventID string + Sequence int64 + IdempotencyKey string + Role MessageEventRole + EventType MessageEventType + ContentJSON string + ToolCallsJSON string + MetadataJSON string + TraceID string + CreatedAt time.Time +} + // IdempotencyRecord stores duplicate delivery state for an inbound message. type IdempotencyRecord struct { TenantID string diff --git a/platform/validation.go b/platform/validation.go index 1848b04bc7..52d6de6f69 100644 --- a/platform/validation.go +++ b/platform/validation.go @@ -350,6 +350,61 @@ func (r AuditRecord) Validate() error { return nil } +// Validate checks that a message event has required identity and safe trace metadata. +func (e MessageEvent) Validate() error { + if strings.TrimSpace(e.TenantID) == "" { + return ErrTenantIDRequired + } + if strings.TrimSpace(e.AppID) == "" { + return ErrAppIDRequired + } + if strings.TrimSpace(e.SessionID) == "" { + return fmt.Errorf("session_id is required") + } + if strings.TrimSpace(e.EventID) == "" { + return fmt.Errorf("event_id is required") + } + if e.Sequence <= 0 { + return fmt.Errorf("sequence must be greater than 0") + } + if strings.TrimSpace(e.IdempotencyKey) == "" { + return fmt.Errorf("idempotency_key is required") + } + for field, value := range map[string]string{ + "tenant_id": e.TenantID, + "app_id": e.AppID, + "session_id": e.SessionID, + "event_id": e.EventID, + "idempotency_key": e.IdempotencyKey, + } { + if err := validateAuditRedactedText(field, value); err != nil { + return err + } + } + switch e.Role { + case MessageEventRoleUser, MessageEventRoleAssistant, MessageEventRoleTool, MessageEventRoleSystem: + default: + return fmt.Errorf("invalid role %q", e.Role) + } + switch e.EventType { + case MessageEventTypeMessage, MessageEventTypeToolCall, MessageEventTypeToolResult, + MessageEventTypeError, MessageEventTypeRevoke, MessageEventTypeEdit: + default: + return fmt.Errorf("invalid event_type %q", e.EventType) + } + for field, value := range map[string]string{ + "trace_id": e.TraceID, + "content_json": e.ContentJSON, + "tool_calls_json": e.ToolCallsJSON, + "metadata_json": e.MetadataJSON, + } { + if err := validateAuditRedactedText(field, value); err != nil { + return err + } + } + return nil +} + // Validate checks that a usage record has required identity and safe accounting values. func (r UsageRecord) Validate() error { if strings.TrimSpace(r.TenantID) == "" { From f05fa12c70ad7b524e0a05606b0d16d3fda65f1f Mon Sep 17 00:00:00 2001 From: Nene7ko_ <141395478+XnLemon@users.noreply.github.com> Date: Thu, 9 Jul 2026 15:32:59 +0800 Subject: [PATCH 46/95] feat(platform): mark tool call trace spans Adds safe platform tool-call trace contract spans. Independent re-review reported P0/P1 clear; CodeRabbit only reported a trivial test-helper nitpick. --- internal/flow/processor/functioncall.go | 2 + internal/flow/processor/functioncall_test.go | 219 +++++++++++++++++++ internal/telemetry/trace.go | 61 ++++-- internal/telemetry/trace_test.go | 155 +++++++++++++ telemetry/semconv/trace/trace.go | 6 + 5 files changed, 427 insertions(+), 16 deletions(-) diff --git a/internal/flow/processor/functioncall.go b/internal/flow/processor/functioncall.go index 95b9837856..f17df244ae 100644 --- a/internal/flow/processor/functioncall.go +++ b/internal/flow/processor/functioncall.go @@ -781,6 +781,7 @@ func (p *FunctionCallResponseProcessor) executeSingleToolCallSequentialResult( ) (toolResult, error) { ctx, span, startedSpan := itrace.StartSpan(ctx, invocation, itelemetry.NewExecuteToolSpanName(toolCall.Function.Name)) if startedSpan { + itelemetry.MarkToolCallSpan(span) defer span.End() } startTime := time.Now() @@ -1023,6 +1024,7 @@ func (p *FunctionCallResponseProcessor) runParallelToolCall( // Trace the tool execution for observability. ctx, span, startedSpan := itrace.StartSpan(ctx, invocation, itelemetry.NewExecuteToolSpanName(tc.Function.Name)) if startedSpan { + itelemetry.MarkToolCallSpan(span) defer span.End() } startTime := time.Now() diff --git a/internal/flow/processor/functioncall_test.go b/internal/flow/processor/functioncall_test.go index 81301d4132..d020f081f9 100644 --- a/internal/flow/processor/functioncall_test.go +++ b/internal/flow/processor/functioncall_test.go @@ -35,6 +35,7 @@ import ( "trpc.group/trpc-go/trpc-agent-go/plugin" "trpc.group/trpc-go/trpc-agent-go/session" skillstate "trpc.group/trpc-go/trpc-agent-go/skill" + semconvtrace "trpc.group/trpc-go/trpc-agent-go/telemetry/semconv/trace" "trpc.group/trpc-go/trpc-agent-go/telemetry/trace" "trpc.group/trpc-go/trpc-agent-go/tool" agenttool "trpc.group/trpc-go/trpc-agent-go/tool/agent" @@ -92,6 +93,22 @@ func useSpanRecorder(t *testing.T) *tracetest.SpanRecorder { return recorder } +func requireRecordedSpanAttribute(t *testing.T, recorder *tracetest.SpanRecorder, spanName, key, value string) { + t.Helper() + for _, span := range recorder.Ended() { + if span.Name() != spanName { + continue + } + for _, attr := range span.Attributes() { + if string(attr.Key) == key && attr.Value.AsString() == value { + return + } + } + t.Fatalf("span %q missing attribute %s=%q; attributes=%v", spanName, key, value, span.Attributes()) + } + t.Fatalf("span %q not recorded; ended spans=%v", spanName, recorder.Ended()) +} + // Minimal callable tool used by tests above type mockCallableTool struct { declaration *tool.Declaration @@ -226,6 +243,90 @@ func TestExecuteSingleToolCallSequential_DisableTracingSkipsSpanCreation(t *test require.Empty(t, recorder.Ended()) } +func TestExecuteSingleToolCallSequential_RecordsToolCallTraceContract(t *testing.T) { + recorder := useSpanRecorder(t) + p := NewFunctionCallResponseProcessor(false, nil) + invocation := agent.NewInvocation() + invocation.AgentName = "test-agent" + response := &model.Response{Model: "mock-model"} + toolCall := model.ToolCall{ + ID: "call-1", + Function: model.FunctionDefinitionParam{ + Name: "echo", + Arguments: []byte(`{"message":"hello"}`), + }, + } + tools := map[string]tool.Tool{ + "echo": &mockCallableTool{ + declaration: &tool.Declaration{Name: "echo"}, + callFn: func(context.Context, []byte) (any, error) { + return "ok", nil + }, + }, + } + + toolEvent, err := p.executeSingleToolCallSequential( + context.Background(), + invocation, + response, + tools, + make(chan *event.Event, 1), + 0, + toolCall, + ) + require.NoError(t, err) + require.NotNil(t, toolEvent) + requireRecordedSpanAttribute( + t, + recorder, + "execute_tool echo", + semconvtrace.KeyTRPCAgentGoTraceSpan, + "tool.call", + ) +} + +func TestExecuteSingleToolCallSequential_RecordsToolCallTraceContractOnCriticalError(t *testing.T) { + recorder := useSpanRecorder(t) + p := NewFunctionCallResponseProcessor(false, nil) + invocation := agent.NewInvocation() + invocation.AgentName = "test-agent" + response := &model.Response{Model: "mock-model"} + toolCall := model.ToolCall{ + ID: "call-stop", + Function: model.FunctionDefinitionParam{ + Name: "stopper", + Arguments: []byte(`{}`), + }, + } + tools := map[string]tool.Tool{ + "stopper": &mockCallableTool{ + declaration: &tool.Declaration{Name: "stopper"}, + callFn: func(context.Context, []byte) (any, error) { + return nil, agent.NewStopError("stop") + }, + }, + } + + toolEvent, err := p.executeSingleToolCallSequential( + context.Background(), + invocation, + response, + tools, + make(chan *event.Event, 1), + 0, + toolCall, + ) + require.Error(t, err) + require.Nil(t, toolEvent) + requireRecordedSpanAttribute( + t, + recorder, + "execute_tool stopper", + semconvtrace.KeyTRPCAgentGoTraceSpan, + "tool.call", + ) +} + func TestExecuteSingleToolCallSequential_AddsToolCallArgsExtension(t *testing.T) { const ( originalArgs = `{"action":"query"}` @@ -328,6 +429,124 @@ func TestExecuteToolCallsInParallel_DisableTracingSkipsSpanCreation(t *testing.T require.Empty(t, recorder.Ended()) } +func TestExecuteToolCallsInParallel_RecordsToolCallTraceContract(t *testing.T) { + recorder := useSpanRecorder(t) + p := NewFunctionCallResponseProcessor(true, nil) + invocation := agent.NewInvocation() + invocation.AgentName = "test-agent" + response := &model.Response{Model: "mock-model"} + toolCalls := []model.ToolCall{ + { + ID: "call-1", + Function: model.FunctionDefinitionParam{ + Name: "tool1", + Arguments: []byte(`{}`), + }, + }, + { + ID: "call-2", + Function: model.FunctionDefinitionParam{ + Name: "tool2", + Arguments: []byte(`{}`), + }, + }, + } + tools := map[string]tool.Tool{ + "tool1": &mockCallableTool{ + declaration: &tool.Declaration{Name: "tool1"}, + callFn: func(context.Context, []byte) (any, error) { + return "ok-1", nil + }, + }, + "tool2": &mockCallableTool{ + declaration: &tool.Declaration{Name: "tool2"}, + callFn: func(context.Context, []byte) (any, error) { + return "ok-2", nil + }, + }, + } + + mergedEvent, err := p.executeToolCallsInParallel( + context.Background(), + invocation, + response, + toolCalls, + tools, + make(chan *event.Event, 2), + ) + require.NoError(t, err) + require.NotNil(t, mergedEvent) + for _, spanName := range []string{ + "execute_tool tool1", + "execute_tool tool2", + "execute_tool (merged tools)", + } { + requireRecordedSpanAttribute( + t, + recorder, + spanName, + semconvtrace.KeyTRPCAgentGoTraceSpan, + "tool.call", + ) + } +} + +func TestExecuteToolCallsInParallel_RecordsToolCallTraceContractOnCriticalError(t *testing.T) { + recorder := useSpanRecorder(t) + p := NewFunctionCallResponseProcessor(true, nil) + invocation := agent.NewInvocation() + invocation.AgentName = "test-agent" + response := &model.Response{Model: "mock-model"} + toolCalls := []model.ToolCall{ + { + ID: "call-stop", + Function: model.FunctionDefinitionParam{ + Name: "stopper", + Arguments: []byte(`{}`), + }, + }, + { + ID: "call-ok", + Function: model.FunctionDefinitionParam{ + Name: "echo", + Arguments: []byte(`{}`), + }, + }, + } + tools := map[string]tool.Tool{ + "stopper": &mockCallableTool{ + declaration: &tool.Declaration{Name: "stopper"}, + callFn: func(context.Context, []byte) (any, error) { + return nil, agent.NewStopError("stop") + }, + }, + "echo": &mockCallableTool{ + declaration: &tool.Declaration{Name: "echo"}, + callFn: func(context.Context, []byte) (any, error) { + return "ok", nil + }, + }, + } + + mergedEvent, err := p.executeToolCallsInParallel( + context.Background(), + invocation, + response, + toolCalls, + tools, + make(chan *event.Event, 2), + ) + require.Error(t, err) + require.NotNil(t, mergedEvent) + requireRecordedSpanAttribute( + t, + recorder, + "execute_tool stopper", + semconvtrace.KeyTRPCAgentGoTraceSpan, + "tool.call", + ) +} + type mockInvocationStateDeltaTool struct { declaration *tool.Declaration callFn func(context.Context, []byte) (any, error) diff --git a/internal/telemetry/trace.go b/internal/telemetry/trace.go index e65abef0d1..119491be80 100644 --- a/internal/telemetry/trace.go +++ b/internal/telemetry/trace.go @@ -45,6 +45,7 @@ const ( SpanNamePrefixExecuteTool = "execute_tool" OperationExecuteTool = "execute_tool" + OperationToolCall = "tool.call" OperationChat = "chat" OperationGenerateContent = "generate_content" OperationInvokeAgent = "invoke_agent" @@ -63,6 +64,19 @@ func NewExecuteToolSpanName(toolName string) string { return OperationExecuteTool + " " + toolName } +// NewToolCallSpanName creates the stable platform tool-call span contract name. +func NewToolCallSpanName() string { + return OperationToolCall +} + +// MarkToolCallSpan marks a span with the stable platform tool-call contract. +func MarkToolCallSpan(span trace.Span) { + if !span.IsRecording() { + return + } + span.SetAttributes(attribute.String(semconvtrace.KeyTRPCAgentGoTraceSpan, NewToolCallSpanName())) +} + // WorkflowType is the normalized type vocabulary used by workflow spans. type WorkflowType string @@ -201,7 +215,9 @@ func TraceToolCall(span trace.Span, sess *session.Session, declaration *tool.Dec attribute.String(semconvtrace.KeyGenAIOperationName, OperationExecuteTool), attribute.String(semconvtrace.KeyGenAIToolName, declaration.Name), attribute.String(semconvtrace.KeyGenAIToolDescription, declaration.Description), + attribute.Bool(semconvtrace.KeyGenAIToolCallArgumentsPresent, len(args) > 0), ) + MarkToolCallSpan(span) if rspEvent != nil { span.SetAttributes(attribute.String(semconvtrace.KeyEventID, rspEvent.ID)) } @@ -212,23 +228,24 @@ func TraceToolCall(span trace.Span, sess *session.Session, declaration *tool.Dec ) } - // args is json-encoded. - setBytesAttribute(span, OperationExecuteTool, semconvtrace.KeyGenAIToolCallArguments, args) if rspEvent != nil && rspEvent.Response != nil { if e := rspEvent.Response.Error; e != nil { - span.SetStatus(codes.Error, e.Message) - span.SetAttributes(responseErrorAttributes(e, semconvtrace.ValueDefaultErrorType)...) + errorType := FormatResponseErrorLabel(e, semconvtrace.ValueDefaultErrorType) + span.SetStatus(codes.Error, errorType) + span.SetAttributes(attribute.String(semconvtrace.KeyErrorType, errorType)) } else if err != nil { - span.SetStatus(codes.Error, err.Error()) - span.SetAttributes(attribute.String(semconvtrace.KeyErrorType, ToErrorType(err, semconvtrace.ValueDefaultErrorType)), attribute.String(semconvtrace.KeyErrorMessage, err.Error())) + errorType := ToErrorType(err, semconvtrace.ValueDefaultErrorType) + span.SetStatus(codes.Error, errorType) + span.SetAttributes(attribute.String(semconvtrace.KeyErrorType, errorType)) } if callIDs := rspEvent.Response.GetToolCallIDs(); len(callIDs) > 0 { span.SetAttributes(attribute.String(semconvtrace.KeyGenAIToolCallID, callIDs[0])) } - setStringAttribute(span, OperationExecuteTool, semconvtrace.KeyGenAIToolCallResult, "", func() ([]byte, error) { - return json.Marshal(rspEvent.Response) - }) + span.SetAttributes(attribute.Bool( + semconvtrace.KeyGenAIToolCallResultPresent, + toolCallResultPresent(rspEvent.Response), + )) } // Setting empty llm request and response (as UI expect these) while not @@ -251,21 +268,23 @@ func TraceMergedToolCalls(span trace.Span, rspEvent *event.Event) { attribute.String(semconvtrace.KeyGenAIOperationName, OperationExecuteTool), attribute.String(semconvtrace.KeyGenAIToolName, ToolNameMergedTools), attribute.String(semconvtrace.KeyGenAIToolDescription, "(merged tools)"), - attribute.String(semconvtrace.KeyGenAIToolCallArguments, "N/A"), + attribute.Bool(semconvtrace.KeyGenAIToolCallArgumentsPresent, false), ) + MarkToolCallSpan(span) if rspEvent != nil && rspEvent.Response != nil { if callIDs := rspEvent.Response.GetToolCallIDs(); len(callIDs) > 0 { span.SetAttributes(attribute.String(semconvtrace.KeyGenAIToolCallID, callIDs[0])) } if e := rspEvent.Response.Error; e != nil { - span.SetStatus(codes.Error, e.Message) - span.SetAttributes(responseErrorAttributes(e, semconvtrace.ValueDefaultErrorType)...) + errorType := FormatResponseErrorLabel(e, semconvtrace.ValueDefaultErrorType) + span.SetStatus(codes.Error, errorType) + span.SetAttributes(attribute.String(semconvtrace.KeyErrorType, errorType)) } span.SetAttributes(attribute.String(semconvtrace.KeyEventID, rspEvent.ID)) - - setStringAttribute(span, OperationExecuteTool, semconvtrace.KeyGenAIToolCallResult, "", func() ([]byte, error) { - return json.Marshal(rspEvent.Response) - }) + span.SetAttributes(attribute.Bool( + semconvtrace.KeyGenAIToolCallResultPresent, + toolCallResultPresent(rspEvent.Response), + )) } // Setting empty llm request and response (as UI expect these) while not @@ -276,6 +295,16 @@ func TraceMergedToolCalls(span trace.Span, rspEvent *event.Event) { ) } +func toolCallResultPresent(rsp *model.Response) bool { + if rsp == nil { + return false + } + if rsp.Error != nil { + return true + } + return len(rsp.Choices) > 0 +} + func resolveInvocationAgentIdentity(invoke *agent.Invocation) (string, string) { if invoke == nil { return "", "" diff --git a/internal/telemetry/trace_test.go b/internal/telemetry/trace_test.go index adb5d8c53e..ad0415682d 100644 --- a/internal/telemetry/trace_test.go +++ b/internal/telemetry/trace_test.go @@ -13,6 +13,7 @@ import ( "context" "encoding/json" "errors" + "strings" "testing" "time" @@ -144,6 +145,41 @@ func attrStringValue(attrs []attribute.KeyValue, key string) (string, bool) { return "", false } +func attrBoolValue(attrs []attribute.KeyValue, key string) (bool, bool) { + for _, kv := range attrs { + if string(kv.Key) == key { + return kv.Value.AsBool(), true + } + } + return false, false +} + +func spanText(span *recordingSpan) string { + var b strings.Builder + b.WriteString(span.statusDesc) + for _, attr := range span.attrs { + b.WriteString(string(attr.Key)) + b.WriteString("=") + b.WriteString(attr.Value.Emit()) + b.WriteString("\n") + } + for _, err := range span.recordedErrors { + if err != nil { + b.WriteString(err.Error()) + b.WriteString("\n") + } + } + return b.String() +} + +func requireSpanOmitsSensitiveText(t *testing.T, span *recordingSpan, secrets ...string) { + t.Helper() + got := spanText(span) + for _, secret := range secrets { + require.NotContains(t, got, secret) + } +} + func TestNewWorkflowSpanName(t *testing.T) { require.Equal(t, "workflow myflow", NewWorkflowSpanName("myflow")) } @@ -524,6 +560,16 @@ func TestNewExecuteToolSpanName(t *testing.T) { } } +func TestNewToolCallSpanName(t *testing.T) { + require.Equal(t, OperationToolCall, NewToolCallSpanName()) +} + +func TestMarkToolCallSpan(t *testing.T) { + span := newRecordingSpan() + MarkToolCallSpan(span) + require.True(t, hasAttr(span.attrs, semconvtrace.KeyTRPCAgentGoTraceSpan, OperationToolCall)) +} + func TestNewSummarizeTaskType(t *testing.T) { tests := []struct { name string @@ -598,6 +644,7 @@ func TestTraceToolCall_NilPaths(t *testing.T) { // Verify basic attributes are always set require.True(t, hasAttr(span.attrs, semconvtrace.KeyGenAISystem, semconvtrace.SystemTRPCGoAgent)) require.True(t, hasAttr(span.attrs, semconvtrace.KeyGenAIOperationName, OperationExecuteTool)) + require.True(t, hasAttr(span.attrs, semconvtrace.KeyTRPCAgentGoTraceSpan, OperationToolCall)) require.True(t, hasAttr(span.attrs, semconvtrace.KeyGenAIToolName, "test_tool")) // Verify error status when err is provided @@ -608,6 +655,113 @@ func TestTraceToolCall_NilPaths(t *testing.T) { } } +func TestTraceToolCall_ContractOmitsRawPayloadAndError(t *testing.T) { + span := newRecordingSpan() + decl := &tool.Declaration{Name: "search", Description: "safe description"} + args := []byte(`{"query":"customer user text","api_key":"sk-live-secret"}`) + rspEvt := event.New("evt-safe", "author") + code := "tool_failed" + rspEvt.Response = &model.Response{ + Error: &model.ResponseError{ + Code: &code, + Message: "Authorization: Bearer raw-token for customer user text", + }, + } + + TraceToolCall( + span, + &session.Session{ID: "session-1", UserID: "user-1"}, + decl, + args, + rspEvt, + errors.New("password=raw-password token=raw-token"), + ) + + require.True(t, hasAttr(span.attrs, semconvtrace.KeyTRPCAgentGoTraceSpan, OperationToolCall)) + require.False(t, hasAttrKey(span.attrs, semconvtrace.KeyGenAIToolCallArguments)) + require.False(t, hasAttrKey(span.attrs, semconvtrace.KeyGenAIToolCallResult)) + require.False(t, hasAttrKey(span.attrs, semconvtrace.KeyErrorMessage)) + require.True(t, hasAttr(span.attrs, semconvtrace.KeyErrorType, "_OTHER_tool_failed")) + gotArgsPresent, ok := attrBoolValue(span.attrs, semconvtrace.KeyGenAIToolCallArgumentsPresent) + require.True(t, ok) + require.True(t, gotArgsPresent) + gotResultPresent, ok := attrBoolValue(span.attrs, semconvtrace.KeyGenAIToolCallResultPresent) + require.True(t, ok) + require.True(t, gotResultPresent) + require.Equal(t, codes.Error, span.status) + require.Equal(t, "_OTHER_tool_failed", span.statusDesc) + requireSpanOmitsSensitiveText( + t, + span, + "customer user text", + "sk-live-secret", + "Authorization", + "raw-token", + "raw-password", + ) +} + +func TestTraceToolCall_ContractOmitsGenericRawError(t *testing.T) { + span := newRecordingSpan() + decl := &tool.Declaration{Name: "lookup", Description: "safe description"} + rspEvt := event.New("evt-safe", "author") + rspEvt.Response = &model.Response{} + + TraceToolCall( + span, + nil, + decl, + []byte(`{"password":"raw-password"}`), + rspEvt, + errors.New("api_key=secret-key in user text"), + ) + + require.False(t, hasAttrKey(span.attrs, semconvtrace.KeyGenAIToolCallArguments)) + require.False(t, hasAttrKey(span.attrs, semconvtrace.KeyGenAIToolCallResult)) + require.False(t, hasAttrKey(span.attrs, semconvtrace.KeyErrorMessage)) + require.True(t, hasAttr(span.attrs, semconvtrace.KeyErrorType, semconvtrace.ValueDefaultErrorType)) + require.Equal(t, semconvtrace.ValueDefaultErrorType, span.statusDesc) + requireSpanOmitsSensitiveText(t, span, "raw-password", "secret-key", "user text") +} + +func TestTraceMergedToolCalls_ContractOmitsRawPayloadAndError(t *testing.T) { + span := newRecordingSpan() + rspEvt := event.New("evt-safe", "author") + code := "merged_failed" + rspEvt.Response = &model.Response{ + Error: &model.ResponseError{ + Code: &code, + Message: "api_key=secret-key Authorization Bearer raw-token", + }, + Choices: []model.Choice{{ + Message: model.Message{Content: "private tool output"}, + }}, + } + + TraceMergedToolCalls(span, rspEvt) + + require.True(t, hasAttr(span.attrs, semconvtrace.KeyTRPCAgentGoTraceSpan, OperationToolCall)) + require.False(t, hasAttrKey(span.attrs, semconvtrace.KeyGenAIToolCallArguments)) + require.False(t, hasAttrKey(span.attrs, semconvtrace.KeyGenAIToolCallResult)) + require.False(t, hasAttrKey(span.attrs, semconvtrace.KeyErrorMessage)) + require.True(t, hasAttr(span.attrs, semconvtrace.KeyErrorType, "_OTHER_merged_failed")) + gotArgsPresent, ok := attrBoolValue(span.attrs, semconvtrace.KeyGenAIToolCallArgumentsPresent) + require.True(t, ok) + require.False(t, gotArgsPresent) + gotResultPresent, ok := attrBoolValue(span.attrs, semconvtrace.KeyGenAIToolCallResultPresent) + require.True(t, ok) + require.True(t, gotResultPresent) + require.Equal(t, "_OTHER_merged_failed", span.statusDesc) + requireSpanOmitsSensitiveText( + t, + span, + "secret-key", + "Authorization", + "raw-token", + "private tool output", + ) +} + func TestTraceMergedToolCalls_NilPaths(t *testing.T) { tests := []struct { name string @@ -630,6 +784,7 @@ func TestTraceMergedToolCalls_NilPaths(t *testing.T) { // Verify basic attributes are always set require.True(t, hasAttr(span.attrs, semconvtrace.KeyGenAISystem, semconvtrace.SystemTRPCGoAgent)) + require.True(t, hasAttr(span.attrs, semconvtrace.KeyTRPCAgentGoTraceSpan, OperationToolCall)) require.True(t, hasAttr(span.attrs, semconvtrace.KeyGenAIToolName, ToolNameMergedTools)) }) } diff --git a/telemetry/semconv/trace/trace.go b/telemetry/semconv/trace/trace.go index bec8f780f5..554c39f04b 100644 --- a/telemetry/semconv/trace/trace.go +++ b/telemetry/semconv/trace/trace.go @@ -46,6 +46,8 @@ const ( KeyTRPCAgentGoUserID = "trpc_go_agent.user.id" // KeyTRPCAgentGoClientTimeToFirstToken is the attribute key for time to first token metric. KeyTRPCAgentGoClientTimeToFirstToken = "trpc_agent_go.client.time_to_first_token" // #nosec G101 - this is a metric key name, not a credential. + // KeyTRPCAgentGoTraceSpan is the stable platform trace span contract name. + KeyTRPCAgentGoTraceSpan = "trpc.go.agent.trace.span" // KeyGenAIAppName is the attribute key for GenAI application name. KeyGenAIAppName = "gen_ai.app.name" @@ -137,6 +139,10 @@ const ( KeyGenAIToolCallArguments = "gen_ai.tool.call.arguments" // KeyGenAIToolCallResult is the attribute key for tool call result. KeyGenAIToolCallResult = "gen_ai.tool.call.result" + // KeyGenAIToolCallArgumentsPresent is the attribute key for whether tool call arguments were provided. + KeyGenAIToolCallArgumentsPresent = "trpc.go.agent.tool.call.arguments_present" + // KeyGenAIToolCallResultPresent is the attribute key for whether a tool call result was provided. + KeyGenAIToolCallResultPresent = "trpc.go.agent.tool.call.result_present" // KeyGenAIRequestToolDefinitions is the attribute key for tool definitions. KeyGenAIRequestToolDefinitions = "gen_ai.request.tool.definitions" From 5c52f3bc24adfa13ea76388613eb30c9aa4c4556 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <141395478+XnLemon@users.noreply.github.com> Date: Thu, 9 Jul 2026 15:40:06 +0800 Subject: [PATCH 47/95] feat(platform): trace memory search spans Adds safe memory search trace spans on the current memory.Reader boundary. Independent review reported P0/P1 clear; local focused validation and build passed. --- internal/flow/processor/content.go | 16 ++- .../flow/processor/content_memory_test.go | 97 ++++++++++++++++++ internal/telemetry/trace.go | 34 +++++++ internal/telemetry/trace_test.go | 53 ++++++++++ memory/tool/tool.go | 30 +++++- memory/tool/tool_test.go | 99 +++++++++++++++++++ telemetry/semconv/trace/trace.go | 8 ++ 7 files changed, 335 insertions(+), 2 deletions(-) diff --git a/internal/flow/processor/content.go b/internal/flow/processor/content.go index 0db00738ea..b1d7a037ed 100644 --- a/internal/flow/processor/content.go +++ b/internal/flow/processor/content.go @@ -28,6 +28,8 @@ import ( "trpc.group/trpc-go/trpc-agent-go/graph" "trpc.group/trpc-go/trpc-agent-go/internal/fileref" iflow "trpc.group/trpc-go/trpc-agent-go/internal/flow" + itelemetry "trpc.group/trpc-go/trpc-agent-go/internal/telemetry" + itrace "trpc.group/trpc-go/trpc-agent-go/internal/trace" "trpc.group/trpc-go/trpc-agent-go/internal/util/message" "trpc.group/trpc-go/trpc-agent-go/log" "trpc.group/trpc-go/trpc-agent-go/memory" @@ -3107,12 +3109,24 @@ func (p *ContentRequestProcessor) getAdaptivePreloadMemoryMessage( Deduplicate: true, HybridSearch: true, } + searchCtx, span, startedSpan := itrace.StartSpan(ctx, inv, itelemetry.NewMemorySearchSpanName()) memories, err := reader.SearchMemories( - ctx, + searchCtx, userKey, query, memory.WithSearchOptions(searchOpts), ) + if startedSpan { + itelemetry.TraceMemorySearch( + span, + searchOpts.MaxResults, + len(memories), + searchOpts.HybridSearch, + searchOpts.Deduplicate, + err, + ) + span.End() + } if err != nil { log.WarnfContext(ctx, "Failed to search memories for preload: %v", err) return p.loadPreloadMemoryMessage(ctx, inv, reader, userKey, budget) diff --git a/internal/flow/processor/content_memory_test.go b/internal/flow/processor/content_memory_test.go index 42dbcb1fcd..9ddf29077a 100644 --- a/internal/flow/processor/content_memory_test.go +++ b/internal/flow/processor/content_memory_test.go @@ -17,12 +17,16 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "go.opentelemetry.io/otel/codes" + sdktrace "go.opentelemetry.io/otel/sdk/trace" "trpc.group/trpc-go/trpc-agent-go/agent" "trpc.group/trpc-go/trpc-agent-go/event" + itelemetry "trpc.group/trpc-go/trpc-agent-go/internal/telemetry" "trpc.group/trpc-go/trpc-agent-go/memory" "trpc.group/trpc-go/trpc-agent-go/model" "trpc.group/trpc-go/trpc-agent-go/session" "trpc.group/trpc-go/trpc-agent-go/session/inmemory" + semconvtrace "trpc.group/trpc-go/trpc-agent-go/telemetry/semconv/trace" "trpc.group/trpc-go/trpc-agent-go/tool" ) @@ -367,6 +371,49 @@ func (m *mockMemoryService) Close() error { return nil } +func requireMemorySearchSpan(t *testing.T, recorder interface { + Ended() []sdktrace.ReadOnlySpan +}) sdktrace.ReadOnlySpan { + t.Helper() + for _, span := range recorder.Ended() { + if span.Name() == itelemetry.NewMemorySearchSpanName() { + return span + } + } + t.Fatalf("span %q not recorded; ended spans=%v", itelemetry.NewMemorySearchSpanName(), recorder.Ended()) + return nil +} + +func requireMemorySearchSpanAttribute(t *testing.T, span sdktrace.ReadOnlySpan, key string, want any) { + t.Helper() + for _, attr := range span.Attributes() { + if string(attr.Key) != key { + continue + } + switch v := want.(type) { + case string: + require.Equal(t, v, attr.Value.AsString()) + case int64: + require.Equal(t, v, attr.Value.AsInt64()) + case bool: + require.Equal(t, v, attr.Value.AsBool()) + default: + t.Fatalf("unsupported expected attribute type %T", want) + } + return + } + t.Fatalf("missing attribute %s=%v; attributes=%v", key, want, span.Attributes()) +} + +func requireNoMemorySearchSpanAttribute(t *testing.T, span sdktrace.ReadOnlySpan, key string) { + t.Helper() + for _, attr := range span.Attributes() { + if string(attr.Key) == key { + t.Fatalf("unexpected attribute %s present; attributes=%v", key, span.Attributes()) + } + } +} + type mockSearchableSessionService struct { session.Service searchResults []session.EventSearchResult @@ -660,6 +707,34 @@ func TestGetPreloadMemoryMessage(t *testing.T) { assert.Contains(t, msg.Content, "Relevant memory") }) + t.Run("positive preload records memory search trace contract", func(t *testing.T) { + recorder := useSpanRecorder(t) + p := NewContentRequestProcessor(WithPreloadMemory(2)) + mockSvc := &mockMemoryService{ + memories: []*memory.Entry{ + newTestMemoryEntry("mem-1", "first"), + newTestMemoryEntry("mem-2", "second"), + newTestMemoryEntry("mem-3", "third"), + }, + searchResults: []*memory.Entry{ + newTestMemoryEntry("mem-search", "Relevant memory"), + }, + } + inv := newTestInvocation(model.NewUserMessage("find relevant"), mockSvc) + + msg := p.getPreloadMemoryMessage(context.Background(), inv) + + require.NotNil(t, msg) + span := requireMemorySearchSpan(t, recorder) + requireMemorySearchSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoTraceSpan, itelemetry.OperationMemorySearch) + requireMemorySearchSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoMemorySearchMaxResults, int64(2)) + requireMemorySearchSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoMemorySearchResultCount, int64(1)) + requireMemorySearchSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoMemorySearchHybrid, true) + requireMemorySearchSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoMemorySearchDeduplicate, true) + requireNoMemorySearchSpanAttribute(t, span, "trpc.go.agent.memory.search.query") + require.NotEqual(t, codes.Error, span.Status().Code) + }) + t.Run("positive preload falls back to recent load when query is empty", func(t *testing.T) { p := NewContentRequestProcessor(WithPreloadMemory(2)) mockSvc := &mockMemoryService{ @@ -699,6 +774,28 @@ func TestGetPreloadMemoryMessage(t *testing.T) { assert.NotContains(t, msg.Content, "third") }) + t.Run("positive preload records memory search trace contract on search error", func(t *testing.T) { + recorder := useSpanRecorder(t) + p := NewContentRequestProcessor(WithPreloadMemory(2)) + mockSvc := &mockMemoryService{ + memories: []*memory.Entry{ + newTestMemoryEntry("mem-1", "first"), + newTestMemoryEntry("mem-2", "second"), + newTestMemoryEntry("mem-3", "third"), + }, + searchErr: assert.AnError, + } + inv := newTestInvocation(model.NewUserMessage("hello"), mockSvc) + + msg := p.getPreloadMemoryMessage(context.Background(), inv) + + require.NotNil(t, msg) + span := requireMemorySearchSpan(t, recorder) + requireMemorySearchSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoTraceSpan, itelemetry.OperationMemorySearch) + requireMemorySearchSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoMemorySearchResultCount, int64(0)) + require.Equal(t, codes.Error, span.Status().Code) + }) + t.Run("positive preload falls back to recent load when search is empty", func(t *testing.T) { p := NewContentRequestProcessor(WithPreloadMemory(2)) mockSvc := &mockMemoryService{ diff --git a/internal/telemetry/trace.go b/internal/telemetry/trace.go index 119491be80..0c2a05aed5 100644 --- a/internal/telemetry/trace.go +++ b/internal/telemetry/trace.go @@ -13,6 +13,7 @@ package telemetry import ( "encoding/json" + "errors" "fmt" "time" @@ -46,6 +47,7 @@ const ( OperationExecuteTool = "execute_tool" OperationToolCall = "tool.call" + OperationMemorySearch = "memory.search" OperationChat = "chat" OperationGenerateContent = "generate_content" OperationInvokeAgent = "invoke_agent" @@ -69,6 +71,11 @@ func NewToolCallSpanName() string { return OperationToolCall } +// NewMemorySearchSpanName creates the stable platform memory-search span contract name. +func NewMemorySearchSpanName() string { + return OperationMemorySearch +} + // MarkToolCallSpan marks a span with the stable platform tool-call contract. func MarkToolCallSpan(span trace.Span) { if !span.IsRecording() { @@ -77,6 +84,33 @@ func MarkToolCallSpan(span trace.Span) { span.SetAttributes(attribute.String(semconvtrace.KeyTRPCAgentGoTraceSpan, NewToolCallSpanName())) } +// TraceMemorySearch marks a memory search span with stable low-cardinality attributes. +func TraceMemorySearch(span trace.Span, maxResults int, resultCount int, hybridSearch bool, deduplicate bool, err error) { + if !span.IsRecording() { + return + } + span.SetAttributes( + attribute.String(semconvtrace.KeyTRPCAgentGoTraceSpan, NewMemorySearchSpanName()), + attribute.Int(semconvtrace.KeyTRPCAgentGoMemorySearchMaxResults, maxResults), + attribute.Int(semconvtrace.KeyTRPCAgentGoMemorySearchResultCount, resultCount), + attribute.Bool(semconvtrace.KeyTRPCAgentGoMemorySearchHybrid, hybridSearch), + attribute.Bool(semconvtrace.KeyTRPCAgentGoMemorySearchDeduplicate, deduplicate), + ) + if err != nil { + recordSafeSpanError(span, err, semconvtrace.ValueDefaultErrorType) + } +} + +func recordSafeSpanError(span trace.Span, err error, fallback string) { + if err == nil { + return + } + errorType := ToErrorType(err, fallback) + span.SetAttributes(attribute.String(semconvtrace.KeyErrorType, errorType)) + span.SetStatus(codes.Error, errorType) + span.RecordError(errors.New(errorType)) +} + // WorkflowType is the normalized type vocabulary used by workflow spans. type WorkflowType string diff --git a/internal/telemetry/trace_test.go b/internal/telemetry/trace_test.go index ad0415682d..42b3abe134 100644 --- a/internal/telemetry/trace_test.go +++ b/internal/telemetry/trace_test.go @@ -564,12 +564,65 @@ func TestNewToolCallSpanName(t *testing.T) { require.Equal(t, OperationToolCall, NewToolCallSpanName()) } +func TestNewMemorySearchSpanName(t *testing.T) { + require.Equal(t, OperationMemorySearch, NewMemorySearchSpanName()) +} + func TestMarkToolCallSpan(t *testing.T) { span := newRecordingSpan() MarkToolCallSpan(span) require.True(t, hasAttr(span.attrs, semconvtrace.KeyTRPCAgentGoTraceSpan, OperationToolCall)) } +func TestTraceMemorySearch(t *testing.T) { + span := newRecordingSpan() + + TraceMemorySearch(span, 3, 2, true, true, nil) + + require.True(t, hasAttr(span.attrs, semconvtrace.KeyTRPCAgentGoTraceSpan, OperationMemorySearch)) + require.True(t, hasAttr(span.attrs, semconvtrace.KeyTRPCAgentGoMemorySearchMaxResults, int64(3))) + require.True(t, hasAttr(span.attrs, semconvtrace.KeyTRPCAgentGoMemorySearchResultCount, int64(2))) + require.True(t, hasAttr(span.attrs, semconvtrace.KeyTRPCAgentGoMemorySearchHybrid, true)) + require.True(t, hasAttr(span.attrs, semconvtrace.KeyTRPCAgentGoMemorySearchDeduplicate, true)) + require.False(t, hasAttrKey(span.attrs, "trpc.go.agent.memory.search.query")) + require.NotEqual(t, codes.Error, span.status) +} + +func TestTraceMemorySearch_Error(t *testing.T) { + span := newRecordingSpan() + err := errors.New("boom") + + TraceMemorySearch(span, 1, 0, false, false, err) + + require.True(t, hasAttr(span.attrs, semconvtrace.KeyTRPCAgentGoTraceSpan, OperationMemorySearch)) + require.True(t, hasAttr(span.attrs, semconvtrace.KeyTRPCAgentGoMemorySearchResultCount, int64(0))) + require.Equal(t, codes.Error, span.status) + require.Equal(t, semconvtrace.ValueDefaultErrorType, span.statusDesc) + require.True(t, hasAttr(span.attrs, semconvtrace.KeyErrorType, semconvtrace.ValueDefaultErrorType)) + require.Len(t, span.recordedErrors, 1) + require.Equal(t, semconvtrace.ValueDefaultErrorType, span.recordedErrors[0].Error()) +} + +func TestTraceMemorySearch_ErrorDoesNotExposeRawErrorText(t *testing.T) { + span := newRecordingSpan() + err := errors.New("query=reset password Authorization: Bearer raw-token api_key=sk-1234567890abcdef") + + TraceMemorySearch(span, 1, 0, false, false, err) + + traceText := span.statusDesc + for _, recorded := range span.recordedErrors { + traceText += "\n" + recorded.Error() + } + for _, attr := range span.attrs { + traceText += "\n" + attr.Value.AsString() + } + require.Equal(t, codes.Error, span.status) + require.NotContains(t, traceText, "reset password") + require.NotContains(t, traceText, "raw-token") + require.NotContains(t, traceText, "sk-1234567890abcdef") + require.NotContains(t, traceText, err.Error()) +} + func TestNewSummarizeTaskType(t *testing.T) { tests := []struct { name string diff --git a/memory/tool/tool.go b/memory/tool/tool.go index e6f4d760b6..62d3a53367 100644 --- a/memory/tool/tool.go +++ b/memory/tool/tool.go @@ -18,6 +18,8 @@ import ( "time" "trpc.group/trpc-go/trpc-agent-go/agent" + itelemetry "trpc.group/trpc-go/trpc-agent-go/internal/telemetry" + itrace "trpc.group/trpc-go/trpc-agent-go/internal/trace" "trpc.group/trpc-go/trpc-agent-go/memory" "trpc.group/trpc-go/trpc-agent-go/tool" "trpc.group/trpc-go/trpc-agent-go/tool/function" @@ -272,11 +274,37 @@ func NewSearchTool() tool.CallableTool { userKey := memory.UserKey{AppName: appName, UserID: userID} opts := buildSearchOptions(req) - memories, err := memoryService.SearchMemories(ctx, userKey, + searchCtx := ctx + var spanStarted bool + var spanErr error + var spanResultCount int + if invocation, ok := agent.InvocationFromContext(ctx); ok { + tracedCtx, span, startedSpan := itrace.StartSpan(ctx, invocation, itelemetry.NewMemorySearchSpanName()) + searchCtx = tracedCtx + spanStarted = startedSpan + if startedSpan { + defer func() { + itelemetry.TraceMemorySearch( + span, + opts.MaxResults, + spanResultCount, + opts.HybridSearch, + opts.Deduplicate, + spanErr, + ) + span.End() + }() + } + } + memories, err := memoryService.SearchMemories(searchCtx, userKey, opts.Query, memory.WithSearchOptions(opts)) if err != nil { + spanErr = err return nil, fmt.Errorf("failed to search memories: %v", err) } + if spanStarted { + spanResultCount = len(memories) + } // Convert MemoryEntry to MemoryResult. results := make([]Result, len(memories)) diff --git a/memory/tool/tool_test.go b/memory/tool/tool_test.go index 9f1827b37e..2aa18ee3aa 100644 --- a/memory/tool/tool_test.go +++ b/memory/tool/tool_test.go @@ -20,11 +20,17 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "go.opentelemetry.io/otel/codes" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + "go.opentelemetry.io/otel/sdk/trace/tracetest" "trpc.group/trpc-go/trpc-agent-go/agent" "trpc.group/trpc-go/trpc-agent-go/event" + itelemetry "trpc.group/trpc-go/trpc-agent-go/internal/telemetry" "trpc.group/trpc-go/trpc-agent-go/memory" "trpc.group/trpc-go/trpc-agent-go/session" + semconvtrace "trpc.group/trpc-go/trpc-agent-go/telemetry/semconv/trace" + tracetelemetry "trpc.group/trpc-go/trpc-agent-go/telemetry/trace" "trpc.group/trpc-go/trpc-agent-go/tool" "trpc.group/trpc-go/trpc-agent-go/tool/function" ) @@ -157,6 +163,54 @@ func createMockContext(appName, userID string, service memory.Service) context.C return agent.NewInvocationContext(context.Background(), mockInvocation) } +func useMemoryToolSpanRecorder(t *testing.T) *tracetest.SpanRecorder { + t.Helper() + recorder := tracetest.NewSpanRecorder() + provider := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(recorder)) + originalProvider := tracetelemetry.TracerProvider + originalTracer := tracetelemetry.Tracer + tracetelemetry.TracerProvider = provider + tracetelemetry.Tracer = provider.Tracer("memory-tool-test") + t.Cleanup(func() { + _ = provider.Shutdown(context.Background()) + tracetelemetry.TracerProvider = originalProvider + tracetelemetry.Tracer = originalTracer + }) + return recorder +} + +func requireMemoryToolSpan(t *testing.T, recorder *tracetest.SpanRecorder) sdktrace.ReadOnlySpan { + t.Helper() + for _, span := range recorder.Ended() { + if span.Name() == itelemetry.NewMemorySearchSpanName() { + return span + } + } + t.Fatalf("span %q not recorded; ended spans=%v", itelemetry.NewMemorySearchSpanName(), recorder.Ended()) + return nil +} + +func requireMemoryToolSpanAttribute(t *testing.T, span sdktrace.ReadOnlySpan, key string, want any) { + t.Helper() + for _, attr := range span.Attributes() { + if string(attr.Key) != key { + continue + } + switch v := want.(type) { + case string: + require.Equal(t, v, attr.Value.AsString()) + case int64: + require.Equal(t, v, attr.Value.AsInt64()) + case bool: + require.Equal(t, v, attr.Value.AsBool()) + default: + t.Fatalf("unsupported expected attribute type %T", want) + } + return + } + t.Fatalf("missing attribute %s=%v; attributes=%v", key, want, span.Attributes()) +} + func TestMemoryTool_AddMemory(t *testing.T) { service := newMockMemoryService() tool := NewAddTool() @@ -266,6 +320,51 @@ func TestMemoryTool_SearchMemory(t *testing.T) { assert.Equal(t, "User likes coffee", response.Results[0].Memory, "Expected memory 'User likes coffee', got '%s'", response.Results[0].Memory) } +func TestMemoryTool_SearchMemory_RecordsMemorySearchTraceContract(t *testing.T) { + recorder := useMemoryToolSpanRecorder(t) + service := newMockMemoryService() + userKey := memory.UserKey{AppName: "test-app", UserID: "test-user"} + require.NoError(t, service.AddMemory(context.Background(), userKey, "User likes coffee", []string{"preferences"})) + + tool := NewSearchTool() + ctx := createMockContext("test-app", "test-user", service) + jsonArgs, err := json.Marshal(map[string]any{ + "query": "coffee", + }) + require.NoError(t, err) + + result, err := tool.Call(ctx, jsonArgs) + + require.NoError(t, err) + require.NotNil(t, result) + span := requireMemoryToolSpan(t, recorder) + requireMemoryToolSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoTraceSpan, itelemetry.OperationMemorySearch) + requireMemoryToolSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoMemorySearchMaxResults, int64(0)) + requireMemoryToolSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoMemorySearchResultCount, int64(1)) + requireMemoryToolSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoMemorySearchHybrid, true) + requireMemoryToolSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoMemorySearchDeduplicate, true) + require.NotEqual(t, codes.Error, span.Status().Code) +} + +func TestMemoryTool_SearchMemory_RecordsMemorySearchTraceContractOnError(t *testing.T) { + recorder := useMemoryToolSpanRecorder(t) + tool := NewSearchTool() + ctx := createMockContext("test-app", "test-user", &mockMemoryServiceWithError{}) + jsonArgs, err := json.Marshal(map[string]any{ + "query": "coffee", + }) + require.NoError(t, err) + + result, err := tool.Call(ctx, jsonArgs) + + require.Error(t, err) + require.Nil(t, result) + span := requireMemoryToolSpan(t, recorder) + requireMemoryToolSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoTraceSpan, itelemetry.OperationMemorySearch) + requireMemoryToolSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoMemorySearchResultCount, int64(0)) + require.Equal(t, codes.Error, span.Status().Code) +} + func TestMemoryTool_LoadMemory(t *testing.T) { service := newMockMemoryService() diff --git a/telemetry/semconv/trace/trace.go b/telemetry/semconv/trace/trace.go index 554c39f04b..f68883494a 100644 --- a/telemetry/semconv/trace/trace.go +++ b/telemetry/semconv/trace/trace.go @@ -48,6 +48,14 @@ const ( KeyTRPCAgentGoClientTimeToFirstToken = "trpc_agent_go.client.time_to_first_token" // #nosec G101 - this is a metric key name, not a credential. // KeyTRPCAgentGoTraceSpan is the stable platform trace span contract name. KeyTRPCAgentGoTraceSpan = "trpc.go.agent.trace.span" + // KeyTRPCAgentGoMemorySearchMaxResults is the configured memory search result cap. + KeyTRPCAgentGoMemorySearchMaxResults = "trpc.go.agent.memory.search.max_results" + // KeyTRPCAgentGoMemorySearchResultCount is the number of returned memory search results. + KeyTRPCAgentGoMemorySearchResultCount = "trpc.go.agent.memory.search.result_count" + // KeyTRPCAgentGoMemorySearchHybrid is whether hybrid memory search was requested. + KeyTRPCAgentGoMemorySearchHybrid = "trpc.go.agent.memory.search.hybrid" + // KeyTRPCAgentGoMemorySearchDeduplicate is whether memory search deduplication was requested. + KeyTRPCAgentGoMemorySearchDeduplicate = "trpc.go.agent.memory.search.deduplicate" // KeyGenAIAppName is the attribute key for GenAI application name. KeyGenAIAppName = "gen_ai.app.name" From 19f09a05b19ecf965618f54ca73c33816c5e6054 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <141395478+XnLemon@users.noreply.github.com> Date: Thu, 9 Jul 2026 15:45:48 +0800 Subject: [PATCH 48/95] feat(platform): trace memory write spans Adds safe memory write trace spans for add/update/delete/clear. Independent review reported P0/P1 clear; local focused validation and build passed. --- internal/telemetry/trace.go | 28 ++++ internal/telemetry/trace_test.go | 32 +++++ memory/tool/tool.go | 60 +++++++- memory/tool/tool_test.go | 228 ++++++++++++++++++++++++++++++- telemetry/semconv/trace/trace.go | 2 + 5 files changed, 340 insertions(+), 10 deletions(-) diff --git a/internal/telemetry/trace.go b/internal/telemetry/trace.go index 0c2a05aed5..0714965064 100644 --- a/internal/telemetry/trace.go +++ b/internal/telemetry/trace.go @@ -48,6 +48,7 @@ const ( OperationExecuteTool = "execute_tool" OperationToolCall = "tool.call" OperationMemorySearch = "memory.search" + OperationMemoryWrite = "memory.write" OperationChat = "chat" OperationGenerateContent = "generate_content" OperationInvokeAgent = "invoke_agent" @@ -56,6 +57,14 @@ const ( OperationWorkflow = "workflow" ) +// Memory write operation values. +const ( + MemoryWriteOperationAdd = "add" + MemoryWriteOperationUpdate = "update" + MemoryWriteOperationDelete = "delete" + MemoryWriteOperationClear = "clear" +) + // NewChatSpanName creates a new chat span name. func NewChatSpanName(requestModel string) string { return newInferenceSpanName(OperationChat, requestModel) @@ -76,6 +85,11 @@ func NewMemorySearchSpanName() string { return OperationMemorySearch } +// NewMemoryWriteSpanName creates the stable platform memory-write span contract name. +func NewMemoryWriteSpanName() string { + return OperationMemoryWrite +} + // MarkToolCallSpan marks a span with the stable platform tool-call contract. func MarkToolCallSpan(span trace.Span) { if !span.IsRecording() { @@ -101,6 +115,20 @@ func TraceMemorySearch(span trace.Span, maxResults int, resultCount int, hybridS } } +// TraceMemoryWrite marks a memory write span with stable low-cardinality attributes. +func TraceMemoryWrite(span trace.Span, operation string, err error) { + if !span.IsRecording() { + return + } + span.SetAttributes( + attribute.String(semconvtrace.KeyTRPCAgentGoTraceSpan, NewMemoryWriteSpanName()), + attribute.String(semconvtrace.KeyTRPCAgentGoMemoryWriteOperation, operation), + ) + if err != nil { + recordSafeSpanError(span, err, semconvtrace.ValueDefaultErrorType) + } +} + func recordSafeSpanError(span trace.Span, err error, fallback string) { if err == nil { return diff --git a/internal/telemetry/trace_test.go b/internal/telemetry/trace_test.go index 42b3abe134..d192407735 100644 --- a/internal/telemetry/trace_test.go +++ b/internal/telemetry/trace_test.go @@ -568,6 +568,10 @@ func TestNewMemorySearchSpanName(t *testing.T) { require.Equal(t, OperationMemorySearch, NewMemorySearchSpanName()) } +func TestNewMemoryWriteSpanName(t *testing.T) { + require.Equal(t, OperationMemoryWrite, NewMemoryWriteSpanName()) +} + func TestMarkToolCallSpan(t *testing.T) { span := newRecordingSpan() MarkToolCallSpan(span) @@ -623,6 +627,34 @@ func TestTraceMemorySearch_ErrorDoesNotExposeRawErrorText(t *testing.T) { require.NotContains(t, traceText, err.Error()) } +func TestTraceMemoryWrite(t *testing.T) { + span := newRecordingSpan() + + TraceMemoryWrite(span, MemoryWriteOperationAdd, nil) + + require.True(t, hasAttr(span.attrs, semconvtrace.KeyTRPCAgentGoTraceSpan, OperationMemoryWrite)) + require.True(t, hasAttr(span.attrs, semconvtrace.KeyTRPCAgentGoMemoryWriteOperation, MemoryWriteOperationAdd)) + require.False(t, hasAttrKey(span.attrs, "trpc.go.agent.memory.write.memory")) + require.False(t, hasAttrKey(span.attrs, "trpc.go.agent.memory.write.memory_id")) + require.NotEqual(t, codes.Error, span.status) +} + +func TestTraceMemoryWrite_Error(t *testing.T) { + span := newRecordingSpan() + err := errors.New("boom") + + TraceMemoryWrite(span, MemoryWriteOperationDelete, err) + + require.True(t, hasAttr(span.attrs, semconvtrace.KeyTRPCAgentGoTraceSpan, OperationMemoryWrite)) + require.True(t, hasAttr(span.attrs, semconvtrace.KeyTRPCAgentGoMemoryWriteOperation, MemoryWriteOperationDelete)) + require.True(t, hasAttr(span.attrs, semconvtrace.KeyErrorType, semconvtrace.ValueDefaultErrorType)) + require.Equal(t, codes.Error, span.status) + require.Equal(t, semconvtrace.ValueDefaultErrorType, span.statusDesc) + require.Len(t, span.recordedErrors, 1) + require.Equal(t, semconvtrace.ValueDefaultErrorType, span.recordedErrors[0].Error()) + require.NotContains(t, spanText(span), err.Error()) +} + func TestNewSummarizeTaskType(t *testing.T) { tests := []struct { name string diff --git a/memory/tool/tool.go b/memory/tool/tool.go index 62d3a53367..b5edada7f1 100644 --- a/memory/tool/tool.go +++ b/memory/tool/tool.go @@ -83,8 +83,21 @@ func NewAddTool() tool.CallableTool { if ep != nil { opts = append(opts, memory.WithMetadata(ep)) } - err = memoryService.AddMemory(ctx, userKey, req.Memory, req.Topics, opts...) + writeCtx := ctx + var spanErr error + if invocation, ok := agent.InvocationFromContext(ctx); ok { + tracedCtx, span, startedSpan := itrace.StartSpan(ctx, invocation, itelemetry.NewMemoryWriteSpanName()) + writeCtx = tracedCtx + if startedSpan { + defer func() { + itelemetry.TraceMemoryWrite(span, itelemetry.MemoryWriteOperationAdd, spanErr) + span.End() + }() + } + } + err = memoryService.AddMemory(writeCtx, userKey, req.Memory, req.Topics, opts...) if err != nil { + spanErr = err return nil, fmt.Errorf("failed to add memory: %v", err) } @@ -143,8 +156,21 @@ func NewUpdateTool() tool.CallableTool { if ep != nil { opts = append(opts, memory.WithUpdateMetadata(ep)) } - err = memoryService.UpdateMemory(ctx, memoryKey, req.Memory, req.Topics, opts...) + writeCtx := ctx + var spanErr error + if invocation, ok := agent.InvocationFromContext(ctx); ok { + tracedCtx, span, startedSpan := itrace.StartSpan(ctx, invocation, itelemetry.NewMemoryWriteSpanName()) + writeCtx = tracedCtx + if startedSpan { + defer func() { + itelemetry.TraceMemoryWrite(span, itelemetry.MemoryWriteOperationUpdate, spanErr) + span.End() + }() + } + } + err = memoryService.UpdateMemory(writeCtx, memoryKey, req.Memory, req.Topics, opts...) if err != nil { + spanErr = err return nil, fmt.Errorf("failed to update memory: %v", err) } @@ -188,8 +214,21 @@ func NewDeleteTool() tool.CallableTool { } memoryKey := memory.Key{AppName: appName, UserID: userID, MemoryID: req.MemoryID} - err = memoryService.DeleteMemory(ctx, memoryKey) + writeCtx := ctx + var spanErr error + if invocation, ok := agent.InvocationFromContext(ctx); ok { + tracedCtx, span, startedSpan := itrace.StartSpan(ctx, invocation, itelemetry.NewMemoryWriteSpanName()) + writeCtx = tracedCtx + if startedSpan { + defer func() { + itelemetry.TraceMemoryWrite(span, itelemetry.MemoryWriteOperationDelete, spanErr) + span.End() + }() + } + } + err = memoryService.DeleteMemory(writeCtx, memoryKey) if err != nil { + spanErr = err return nil, fmt.Errorf("failed to delete memory: %v", err) } @@ -227,8 +266,21 @@ func NewClearTool() tool.CallableTool { } userKey := memory.UserKey{AppName: appName, UserID: userID} - err = memoryService.ClearMemories(ctx, userKey) + writeCtx := ctx + var spanErr error + if invocation, ok := agent.InvocationFromContext(ctx); ok { + tracedCtx, span, startedSpan := itrace.StartSpan(ctx, invocation, itelemetry.NewMemoryWriteSpanName()) + writeCtx = tracedCtx + if startedSpan { + defer func() { + itelemetry.TraceMemoryWrite(span, itelemetry.MemoryWriteOperationClear, spanErr) + span.End() + }() + } + } + err = memoryService.ClearMemories(writeCtx, userKey) if err != nil { + spanErr = err return nil, fmt.Errorf("memory clear tool: failed to clear memories: %v", err) } diff --git a/memory/tool/tool_test.go b/memory/tool/tool_test.go index 2aa18ee3aa..162022c517 100644 --- a/memory/tool/tool_test.go +++ b/memory/tool/tool_test.go @@ -180,13 +180,17 @@ func useMemoryToolSpanRecorder(t *testing.T) *tracetest.SpanRecorder { } func requireMemoryToolSpan(t *testing.T, recorder *tracetest.SpanRecorder) sdktrace.ReadOnlySpan { + return requireMemoryToolSpanNamed(t, recorder, itelemetry.NewMemorySearchSpanName()) +} + +func requireMemoryToolSpanNamed(t *testing.T, recorder *tracetest.SpanRecorder, spanName string) sdktrace.ReadOnlySpan { t.Helper() for _, span := range recorder.Ended() { - if span.Name() == itelemetry.NewMemorySearchSpanName() { + if span.Name() == spanName { return span } } - t.Fatalf("span %q not recorded; ended spans=%v", itelemetry.NewMemorySearchSpanName(), recorder.Ended()) + t.Fatalf("span %q not recorded; ended spans=%v", spanName, recorder.Ended()) return nil } @@ -211,6 +215,36 @@ func requireMemoryToolSpanAttribute(t *testing.T, span sdktrace.ReadOnlySpan, ke t.Fatalf("missing attribute %s=%v; attributes=%v", key, want, span.Attributes()) } +func requireNoMemoryToolSpanAttribute(t *testing.T, span sdktrace.ReadOnlySpan, key string) { + t.Helper() + for _, attr := range span.Attributes() { + if string(attr.Key) == key { + t.Fatalf("unexpected attribute %s present; attributes=%v", key, span.Attributes()) + } + } +} + +func requireMemoryToolSpanSafeError(t *testing.T, span sdktrace.ReadOnlySpan, rawValues ...string) { + t.Helper() + require.Equal(t, codes.Error, span.Status().Code) + require.Equal(t, semconvtrace.ValueDefaultErrorType, span.Status().Description) + requireMemoryToolSpanAttribute(t, span, semconvtrace.KeyErrorType, semconvtrace.ValueDefaultErrorType) + + traceText := span.Status().Description + for _, attr := range span.Attributes() { + traceText += "\n" + string(attr.Key) + "=" + attr.Value.AsString() + } + for _, event := range span.Events() { + traceText += "\n" + event.Name + for _, attr := range event.Attributes { + traceText += "\n" + string(attr.Key) + "=" + attr.Value.AsString() + } + } + for _, raw := range rawValues { + require.NotContains(t, traceText, raw) + } +} + func TestMemoryTool_AddMemory(t *testing.T) { service := newMockMemoryService() tool := NewAddTool() @@ -245,6 +279,29 @@ func TestMemoryTool_AddMemory(t *testing.T) { assert.Equal(t, "User's name is John Doe", memories[0].Memory.Memory, "Expected memory 'User's name is John Doe', got '%s'", memories[0].Memory.Memory) } +func TestMemoryTool_AddMemory_RecordsMemoryWriteTraceContract(t *testing.T) { + recorder := useMemoryToolSpanRecorder(t) + service := newMockMemoryService() + tool := NewAddTool() + ctx := createMockContext("test-app", "test-user", service) + jsonArgs, err := json.Marshal(map[string]any{ + "memory": "User's name is John Doe", + "topics": []string{"personal"}, + }) + require.NoError(t, err) + + result, err := tool.Call(ctx, jsonArgs) + + require.NoError(t, err) + require.NotNil(t, result) + span := requireMemoryToolSpanNamed(t, recorder, itelemetry.NewMemoryWriteSpanName()) + requireMemoryToolSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoTraceSpan, itelemetry.OperationMemoryWrite) + requireMemoryToolSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoMemoryWriteOperation, itelemetry.MemoryWriteOperationAdd) + requireNoMemoryToolSpanAttribute(t, span, "trpc.go.agent.memory.write.memory") + requireNoMemoryToolSpanAttribute(t, span, "trpc.go.agent.memory.write.memory_id") + require.NotEqual(t, codes.Error, span.Status().Code) +} + func TestMemoryTool_AddMemory_WithoutTopics(t *testing.T) { service := newMockMemoryService() tool := NewAddTool() @@ -442,6 +499,35 @@ func TestMemoryTool_UpdateMemory(t *testing.T) { assert.Equal(t, "User loves coffee and tea", updatedMemories[0].Memory.Memory, "Expected updated memory content") } +func TestMemoryTool_UpdateMemory_RecordsMemoryWriteTraceContract(t *testing.T) { + recorder := useMemoryToolSpanRecorder(t) + service := newMockMemoryService() + userKey := memory.UserKey{AppName: "test-app", UserID: "test-user"} + require.NoError(t, service.AddMemory(context.Background(), userKey, "User likes coffee", []string{"preferences"})) + memories, err := service.ReadMemories(context.Background(), userKey, 1) + require.NoError(t, err) + require.Len(t, memories, 1) + + tool := NewUpdateTool() + ctx := createMockContext("test-app", "test-user", service) + jsonArgs, err := json.Marshal(map[string]any{ + "memory_id": memories[0].ID, + "memory": "User loves coffee and tea", + }) + require.NoError(t, err) + + result, err := tool.Call(ctx, jsonArgs) + + require.NoError(t, err) + require.NotNil(t, result) + span := requireMemoryToolSpanNamed(t, recorder, itelemetry.NewMemoryWriteSpanName()) + requireMemoryToolSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoTraceSpan, itelemetry.OperationMemoryWrite) + requireMemoryToolSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoMemoryWriteOperation, itelemetry.MemoryWriteOperationUpdate) + requireNoMemoryToolSpanAttribute(t, span, "trpc.go.agent.memory.write.memory") + requireNoMemoryToolSpanAttribute(t, span, "trpc.go.agent.memory.write.memory_id") + require.NotEqual(t, codes.Error, span.Status().Code) +} + func TestMemoryTool_UpdateMemory_WithoutTopics(t *testing.T) { service := newMockMemoryService() @@ -581,6 +667,57 @@ func TestMemoryTool_DeleteMemory(t *testing.T) { assert.Len(t, deletedMemories, 0, "Expected 0 memories after deletion") } +func TestMemoryTool_DeleteMemory_RecordsMemoryWriteTraceContract(t *testing.T) { + recorder := useMemoryToolSpanRecorder(t) + service := newMockMemoryService() + userKey := memory.UserKey{AppName: "test-app", UserID: "test-user"} + require.NoError(t, service.AddMemory(context.Background(), userKey, "User likes coffee", []string{"preferences"})) + memories, err := service.ReadMemories(context.Background(), userKey, 1) + require.NoError(t, err) + require.Len(t, memories, 1) + + tool := NewDeleteTool() + ctx := createMockContext("test-app", "test-user", service) + jsonArgs, err := json.Marshal(map[string]any{ + "memory_id": memories[0].ID, + }) + require.NoError(t, err) + + result, err := tool.Call(ctx, jsonArgs) + + require.NoError(t, err) + require.NotNil(t, result) + span := requireMemoryToolSpanNamed(t, recorder, itelemetry.NewMemoryWriteSpanName()) + requireMemoryToolSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoTraceSpan, itelemetry.OperationMemoryWrite) + requireMemoryToolSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoMemoryWriteOperation, itelemetry.MemoryWriteOperationDelete) + requireNoMemoryToolSpanAttribute(t, span, "trpc.go.agent.memory.write.memory") + requireNoMemoryToolSpanAttribute(t, span, "trpc.go.agent.memory.write.memory_id") + require.NotEqual(t, codes.Error, span.Status().Code) +} + +func TestMemoryTool_ClearMemory_RecordsMemoryWriteTraceContract(t *testing.T) { + recorder := useMemoryToolSpanRecorder(t) + service := newMockMemoryService() + userKey := memory.UserKey{AppName: "test-app", UserID: "test-user"} + require.NoError(t, service.AddMemory(context.Background(), userKey, "User likes coffee", []string{"preferences"})) + + tool := NewClearTool() + ctx := createMockContext("test-app", "test-user", service) + jsonArgs, err := json.Marshal(map[string]any{}) + require.NoError(t, err) + + result, err := tool.Call(ctx, jsonArgs) + + require.NoError(t, err) + require.NotNil(t, result) + span := requireMemoryToolSpanNamed(t, recorder, itelemetry.NewMemoryWriteSpanName()) + requireMemoryToolSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoTraceSpan, itelemetry.OperationMemoryWrite) + requireMemoryToolSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoMemoryWriteOperation, itelemetry.MemoryWriteOperationClear) + requireNoMemoryToolSpanAttribute(t, span, "trpc.go.agent.memory.write.memory") + requireNoMemoryToolSpanAttribute(t, span, "trpc.go.agent.memory.write.memory_id") + require.NotEqual(t, codes.Error, span.Status().Code) +} + func TestMemoryTool_DeleteMemory_InvalidID(t *testing.T) { service := newMockMemoryService() tool := NewDeleteTool() @@ -1210,6 +1347,85 @@ func TestMemoryTool_AddMemory_ServiceError(t *testing.T) { assert.Contains(t, err.Error(), "failed to add memory") } +func TestMemoryTool_AddMemory_RecordsMemoryWriteTraceContractOnError(t *testing.T) { + recorder := useMemoryToolSpanRecorder(t) + service := &mockMemoryServiceWithError{} + tool := NewAddTool() + ctx := createMockContext("test-app", "test-user", service) + jsonArgs, err := json.Marshal(map[string]any{ + "memory": "reset password with Authorization: Bearer raw-token and api_key=sk-1234567890abcdef", + }) + require.NoError(t, err) + + result, err := tool.Call(ctx, jsonArgs) + + require.Error(t, err) + require.Nil(t, result) + span := requireMemoryToolSpanNamed(t, recorder, itelemetry.NewMemoryWriteSpanName()) + requireMemoryToolSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoTraceSpan, itelemetry.OperationMemoryWrite) + requireMemoryToolSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoMemoryWriteOperation, itelemetry.MemoryWriteOperationAdd) + requireMemoryToolSpanSafeError(t, span, "reset password", "raw-token", "sk-1234567890abcdef") +} + +func TestMemoryTool_UpdateMemory_RecordsMemoryWriteTraceContractOnError(t *testing.T) { + recorder := useMemoryToolSpanRecorder(t) + service := &mockMemoryServiceWithError{} + tool := NewUpdateTool() + ctx := createMockContext("test-app", "test-user", service) + jsonArgs, err := json.Marshal(map[string]any{ + "memory_id": "memory-1-raw-token", + "memory": "updated reset password with Authorization: Bearer raw-token and api_key=sk-1234567890abcdef", + }) + require.NoError(t, err) + + result, err := tool.Call(ctx, jsonArgs) + + require.Error(t, err) + require.Nil(t, result) + span := requireMemoryToolSpanNamed(t, recorder, itelemetry.NewMemoryWriteSpanName()) + requireMemoryToolSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoTraceSpan, itelemetry.OperationMemoryWrite) + requireMemoryToolSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoMemoryWriteOperation, itelemetry.MemoryWriteOperationUpdate) + requireMemoryToolSpanSafeError(t, span, "memory-1-raw-token", "reset password", "raw-token", "sk-1234567890abcdef") +} + +func TestMemoryTool_DeleteMemory_RecordsMemoryWriteTraceContractOnError(t *testing.T) { + recorder := useMemoryToolSpanRecorder(t) + service := &mockMemoryServiceWithError{} + tool := NewDeleteTool() + ctx := createMockContext("test-app", "test-user", service) + jsonArgs, err := json.Marshal(map[string]any{ + "memory_id": "memory-1-raw-token", + }) + require.NoError(t, err) + + result, err := tool.Call(ctx, jsonArgs) + + require.Error(t, err) + require.Nil(t, result) + span := requireMemoryToolSpanNamed(t, recorder, itelemetry.NewMemoryWriteSpanName()) + requireMemoryToolSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoTraceSpan, itelemetry.OperationMemoryWrite) + requireMemoryToolSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoMemoryWriteOperation, itelemetry.MemoryWriteOperationDelete) + requireMemoryToolSpanSafeError(t, span, "memory-1-raw-token", "raw-token") +} + +func TestMemoryTool_ClearMemory_RecordsMemoryWriteTraceContractOnError(t *testing.T) { + recorder := useMemoryToolSpanRecorder(t) + service := &mockMemoryServiceWithError{} + tool := NewClearTool() + ctx := createMockContext("test-app", "test-user", service) + jsonArgs, err := json.Marshal(map[string]any{}) + require.NoError(t, err) + + result, err := tool.Call(ctx, jsonArgs) + + require.Error(t, err) + require.Nil(t, result) + span := requireMemoryToolSpanNamed(t, recorder, itelemetry.NewMemoryWriteSpanName()) + requireMemoryToolSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoTraceSpan, itelemetry.OperationMemoryWrite) + requireMemoryToolSpanAttribute(t, span, semconvtrace.KeyTRPCAgentGoMemoryWriteOperation, itelemetry.MemoryWriteOperationClear) + requireMemoryToolSpanSafeError(t, span, "reset password", "raw-token", "sk-1234567890abcdef") +} + func TestMemoryTool_SearchMemory_ServiceError(t *testing.T) { service := &mockMemoryServiceWithError{} tool := NewSearchTool() @@ -1250,19 +1466,19 @@ func TestMemoryTool_LoadMemory_ServiceError(t *testing.T) { type mockMemoryServiceWithError struct{} func (m *mockMemoryServiceWithError) AddMemory(ctx context.Context, userKey memory.UserKey, memoryStr string, topics []string, opts ...memory.AddOption) error { - return fmt.Errorf("mock add error") + return fmt.Errorf("mock add error: memory=%s Authorization: Bearer raw-token api_key=sk-1234567890abcdef", memoryStr) } func (m *mockMemoryServiceWithError) UpdateMemory(ctx context.Context, memoryKey memory.Key, mem string, topics []string, opts ...memory.UpdateOption) error { - return fmt.Errorf("mock update error") + return fmt.Errorf("mock update error: memory_id=%s memory=%s Authorization: Bearer raw-token api_key=sk-1234567890abcdef", memoryKey.MemoryID, mem) } func (m *mockMemoryServiceWithError) DeleteMemory(ctx context.Context, memoryKey memory.Key) error { - return fmt.Errorf("mock delete error") + return fmt.Errorf("mock delete error: memory_id=%s Authorization: Bearer raw-token", memoryKey.MemoryID) } func (m *mockMemoryServiceWithError) ClearMemories(ctx context.Context, userKey memory.UserKey) error { - return fmt.Errorf("mock clear error") + return fmt.Errorf("mock clear error: reset password Authorization: Bearer raw-token api_key=sk-1234567890abcdef") } func (m *mockMemoryServiceWithError) ReadMemories(ctx context.Context, userKey memory.UserKey, limit int) ([]*memory.Entry, error) { diff --git a/telemetry/semconv/trace/trace.go b/telemetry/semconv/trace/trace.go index f68883494a..3d453bcfda 100644 --- a/telemetry/semconv/trace/trace.go +++ b/telemetry/semconv/trace/trace.go @@ -56,6 +56,8 @@ const ( KeyTRPCAgentGoMemorySearchHybrid = "trpc.go.agent.memory.search.hybrid" // KeyTRPCAgentGoMemorySearchDeduplicate is whether memory search deduplication was requested. KeyTRPCAgentGoMemorySearchDeduplicate = "trpc.go.agent.memory.search.deduplicate" + // KeyTRPCAgentGoMemoryWriteOperation is the memory write operation type. + KeyTRPCAgentGoMemoryWriteOperation = "trpc.go.agent.memory.write.operation" // KeyGenAIAppName is the attribute key for GenAI application name. KeyGenAIAppName = "gen_ai.app.name" From f919d4ddde571c3dc7d157f81a0152229af536ae Mon Sep 17 00:00:00 2001 From: Nene7ko_ <141395478+XnLemon@users.noreply.github.com> Date: Thu, 9 Jul 2026 15:50:13 +0800 Subject: [PATCH 49/95] feat(platform): mark summary create trace spans Adds summary-create trace contract markers for Redis/PostgreSQL summary creation spans. Independent review reported P0/P1 clear; local focused validation and build passed. --- internal/telemetry/trace.go | 14 ++++ internal/telemetry/trace_test.go | 12 +++ session/postgres/service_tracing_test.go | 102 +++++++++++++++++++++++ session/postgres/summary.go | 7 ++ session/redis/service_tracing_test.go | 3 + session/redis/summary.go | 2 + 6 files changed, 140 insertions(+) create mode 100644 session/postgres/service_tracing_test.go diff --git a/internal/telemetry/trace.go b/internal/telemetry/trace.go index 0714965064..88ffae2f23 100644 --- a/internal/telemetry/trace.go +++ b/internal/telemetry/trace.go @@ -49,6 +49,7 @@ const ( OperationToolCall = "tool.call" OperationMemorySearch = "memory.search" OperationMemoryWrite = "memory.write" + OperationSummaryCreate = "summary.create" OperationChat = "chat" OperationGenerateContent = "generate_content" OperationInvokeAgent = "invoke_agent" @@ -90,6 +91,11 @@ func NewMemoryWriteSpanName() string { return OperationMemoryWrite } +// NewSummaryCreateSpanName creates the stable platform summary-create span contract name. +func NewSummaryCreateSpanName() string { + return OperationSummaryCreate +} + // MarkToolCallSpan marks a span with the stable platform tool-call contract. func MarkToolCallSpan(span trace.Span) { if !span.IsRecording() { @@ -129,6 +135,14 @@ func TraceMemoryWrite(span trace.Span, operation string, err error) { } } +// MarkSummaryCreateSpan marks a span with the stable platform summary-create contract. +func MarkSummaryCreateSpan(span trace.Span) { + if !span.IsRecording() { + return + } + span.SetAttributes(attribute.String(semconvtrace.KeyTRPCAgentGoTraceSpan, NewSummaryCreateSpanName())) +} + func recordSafeSpanError(span trace.Span, err error, fallback string) { if err == nil { return diff --git a/internal/telemetry/trace_test.go b/internal/telemetry/trace_test.go index d192407735..61f0ed499c 100644 --- a/internal/telemetry/trace_test.go +++ b/internal/telemetry/trace_test.go @@ -572,6 +572,10 @@ func TestNewMemoryWriteSpanName(t *testing.T) { require.Equal(t, OperationMemoryWrite, NewMemoryWriteSpanName()) } +func TestNewSummaryCreateSpanName(t *testing.T) { + require.Equal(t, OperationSummaryCreate, NewSummaryCreateSpanName()) +} + func TestMarkToolCallSpan(t *testing.T) { span := newRecordingSpan() MarkToolCallSpan(span) @@ -655,6 +659,14 @@ func TestTraceMemoryWrite_Error(t *testing.T) { require.NotContains(t, spanText(span), err.Error()) } +func TestMarkSummaryCreateSpan(t *testing.T) { + span := newRecordingSpan() + + MarkSummaryCreateSpan(span) + + require.True(t, hasAttr(span.attrs, semconvtrace.KeyTRPCAgentGoTraceSpan, OperationSummaryCreate)) +} + func TestNewSummarizeTaskType(t *testing.T) { tests := []struct { name string diff --git a/session/postgres/service_tracing_test.go b/session/postgres/service_tracing_test.go new file mode 100644 index 0000000000..fef97086e5 --- /dev/null +++ b/session/postgres/service_tracing_test.go @@ -0,0 +1,102 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// +// + +package postgres + +import ( + "context" + "testing" + "time" + + "github.com/DATA-DOG/go-sqlmock" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.opentelemetry.io/otel" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + "go.opentelemetry.io/otel/sdk/trace/tracetest" + itelemetry "trpc.group/trpc-go/trpc-agent-go/internal/telemetry" + "trpc.group/trpc-go/trpc-agent-go/session" + semconvtrace "trpc.group/trpc-go/trpc-agent-go/telemetry/semconv/trace" + atrace "trpc.group/trpc-go/trpc-agent-go/telemetry/trace" +) + +func setupTracingProvider(t *testing.T) (*tracetest.InMemoryExporter, func()) { + t.Helper() + + origTracer := atrace.Tracer + origProvider := atrace.TracerProvider + + exporter := tracetest.NewInMemoryExporter() + tp := sdktrace.NewTracerProvider( + sdktrace.WithSyncer(exporter), + sdktrace.WithSampler(sdktrace.AlwaysSample()), + ) + + otel.SetTracerProvider(tp) + atrace.TracerProvider = tp + atrace.Tracer = tp.Tracer("test") + + cleanup := func() { + _ = tp.Shutdown(context.Background()) + atrace.Tracer = origTracer + atrace.TracerProvider = origProvider + otel.SetTracerProvider(origProvider) + } + + return exporter, cleanup +} + +func findSpan(spans tracetest.SpanStubs, name string) *tracetest.SpanStub { + for i := range spans { + if spans[i].Name == name { + return &spans[i] + } + } + return nil +} + +func spanAttr(s *tracetest.SpanStub, key string) string { + for _, a := range s.Attributes { + if string(a.Key) == key { + return a.Value.AsString() + } + } + return "" +} + +func TestCreateSessionSummary_WithTracing(t *testing.T) { + exporter, cleanupTP := setupTracingProvider(t) + defer cleanupTP() + + summarizer := &mockSummarizerImpl{ + summaryText: "traced summary", + shouldSummarize: true, + } + s, mock, db := setupMockService(t, &TestServiceOpts{summarizer: summarizer}) + defer db.Close() + + sess := &session.Session{ + ID: "trace-session", + AppName: "trace-app", + UserID: "trace-user", + UpdatedAt: time.Now(), + } + + mock.ExpectExec("INSERT INTO session_summaries"). + WithArgs("trace-app", "trace-user", "trace-session", "", sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg()). + WillReturnResult(sqlmock.NewResult(1, 1)) + + err := s.CreateSessionSummary(context.Background(), sess, "", true) + require.NoError(t, err) + require.NoError(t, mock.ExpectationsWereMet()) + + span := findSpan(exporter.GetSpans(), "create_session_summary") + require.NotNil(t, span, "expected create_session_summary span") + assert.Equal(t, itelemetry.OperationSummaryCreate, spanAttr(span, semconvtrace.KeyTRPCAgentGoTraceSpan)) +} diff --git a/session/postgres/summary.go b/session/postgres/summary.go index c4453b7b87..e900cab890 100644 --- a/session/postgres/summary.go +++ b/session/postgres/summary.go @@ -16,8 +16,10 @@ import ( "fmt" "time" + itelemetry "trpc.group/trpc-go/trpc-agent-go/internal/telemetry" "trpc.group/trpc-go/trpc-agent-go/session" isummary "trpc.group/trpc-go/trpc-agent-go/session/internal/summary" + atrace "trpc.group/trpc-go/trpc-agent-go/telemetry/trace" ) // CreateSessionSummary is the internal implementation that returns the summary. @@ -39,6 +41,11 @@ func (s *Service) CreateSessionSummary( if err := key.CheckSessionKey(); err != nil { return fmt.Errorf("check session key failed: %w", err) } + + ctx, span := atrace.Tracer.Start(ctx, "create_session_summary") + itelemetry.MarkSummaryCreateSpan(span) + defer span.End() + if !isummary.NewSummaryDispatchPolicy( s.opts.summaryFilterAllowlist, s.opts.shouldCascadeFullSessionSummary(), diff --git a/session/redis/service_tracing_test.go b/session/redis/service_tracing_test.go index 076e55d33d..58503eb93f 100644 --- a/session/redis/service_tracing_test.go +++ b/session/redis/service_tracing_test.go @@ -22,8 +22,10 @@ import ( sdktrace "go.opentelemetry.io/otel/sdk/trace" "go.opentelemetry.io/otel/sdk/trace/tracetest" "trpc.group/trpc-go/trpc-agent-go/event" + itelemetry "trpc.group/trpc-go/trpc-agent-go/internal/telemetry" "trpc.group/trpc-go/trpc-agent-go/model" "trpc.group/trpc-go/trpc-agent-go/session" + semconvtrace "trpc.group/trpc-go/trpc-agent-go/telemetry/semconv/trace" atrace "trpc.group/trpc-go/trpc-agent-go/telemetry/trace" ) @@ -336,6 +338,7 @@ func TestCreateSessionSummary_WithTracing(t *testing.T) { s := findSpan(spans, "create_session_summary") require.NotNil(t, s, "expected create_session_summary span") assert.Equal(t, "css1", spanAttr(s, "session_id")) + assert.Equal(t, itelemetry.OperationSummaryCreate, spanAttr(s, semconvtrace.KeyTRPCAgentGoTraceSpan)) } // ============================================================================ diff --git a/session/redis/summary.go b/session/redis/summary.go index 9f64ce94fe..26c38e0bc8 100644 --- a/session/redis/summary.go +++ b/session/redis/summary.go @@ -14,6 +14,7 @@ import ( "fmt" "time" + itelemetry "trpc.group/trpc-go/trpc-agent-go/internal/telemetry" "trpc.group/trpc-go/trpc-agent-go/log" "trpc.group/trpc-go/trpc-agent-go/session" isummary "trpc.group/trpc-go/trpc-agent-go/session/internal/summary" @@ -34,6 +35,7 @@ func (s *Service) CreateSessionSummary(ctx context.Context, sess *session.Sessio key := session.Key{AppName: sess.AppName, UserID: sess.UserID, SessionID: sess.ID} ctx, span := s.startSpan(ctx, "create_session_summary", key) + itelemetry.MarkSummaryCreateSpan(span) defer span.End() if err := key.CheckSessionKey(); err != nil { From dc28c707eb6dfc5a7d396abf80b0b6d88fbac71b Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Thu, 9 Jul 2026 17:03:59 +0800 Subject: [PATCH 50/95] fix(platform): address CI checks --- memory/mysql/go.mod | 17 ++ memory/mysql/go.sum | 12 +- memory/mysqlvec/go.mod | 20 +- memory/mysqlvec/go.sum | 9 +- memory/pgvector/go.mod | 16 + memory/pgvector/go.sum | 7 +- memory/postgres/go.mod | 16 + memory/postgres/go.sum | 7 +- memory/redis/go.mod | 17 ++ memory/redis/go.sum | 12 +- memory/sqlite/go.mod | 17 ++ memory/sqlite/go.sum | 12 +- memory/sqlitevec/go.mod | 16 + memory/sqlitevec/go.sum | 12 +- platform/backend_migration_status.go | 18 ++ platform/budget_audit.go | 47 ++- platform/config_operation_summary.go | 104 +++++-- platform/gateway/service.go | 426 ++++++++++++++++++--------- platform/gateway/service_test.go | 6 +- platform/toolpolicy/policy.go | 36 ++- session/postgres/go.mod | 4 +- 21 files changed, 620 insertions(+), 211 deletions(-) diff --git a/memory/mysql/go.mod b/memory/mysql/go.mod index 7cb4a19d1e..c3c62da7de 100644 --- a/memory/mysql/go.mod +++ b/memory/mysql/go.mod @@ -17,15 +17,32 @@ require ( require ( filippo.io/edwards25519 v1.1.1 // indirect + github.com/cenkalti/backoff/v4 v4.3.0 // indirect github.com/davecgh/go-spew v1.1.1 // indirect github.com/go-ego/gse v1.0.0 // indirect + github.com/go-logr/logr v1.4.2 // indirect + github.com/go-logr/stdr v1.2.2 // indirect github.com/google/uuid v1.6.0 // indirect + github.com/grpc-ecosystem/grpc-gateway/v2 v2.22.0 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect github.com/vcaesar/cedar v0.20.2 // indirect go.opentelemetry.io/otel v1.29.0 // indirect + go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.29.0 // indirect + go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.29.0 // indirect + go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.29.0 // indirect + go.opentelemetry.io/otel/metric v1.29.0 // indirect + go.opentelemetry.io/otel/sdk v1.29.0 // indirect go.opentelemetry.io/otel/trace v1.29.0 // indirect + go.opentelemetry.io/proto/otlp v1.3.1 // indirect go.uber.org/multierr v1.10.0 // indirect go.uber.org/zap v1.27.0 // indirect + golang.org/x/net v0.34.0 // indirect + golang.org/x/sys v0.30.0 // indirect + golang.org/x/text v0.21.0 // indirect + google.golang.org/genproto/googleapis/api v0.0.0-20240822170219-fc7c04adadcd // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20240822170219-fc7c04adadcd // indirect + google.golang.org/grpc v1.65.0 // indirect + google.golang.org/protobuf v1.34.2 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect trpc.group/trpc-go/trpc-a2a-go v0.2.5 // indirect ) diff --git a/memory/mysql/go.sum b/memory/mysql/go.sum index 9a406885e1..682b5de524 100644 --- a/memory/mysql/go.sum +++ b/memory/mysql/go.sum @@ -12,6 +12,7 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/go-ego/gse v1.0.0 h1:GNbtH1WP7Yd1VvCZ85fIK6eVEe7RctmgmnwliEPUMNA= github.com/go-ego/gse v1.0.0/go.mod h1:Gt3A9Ry1Eso2Kza4MRaiZ7f2DTAvActmETY46Lxg0gU= +github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= github.com/go-logr/logr v1.4.2 h1:6pFjapn8bFcIbiKo3XT4j/BhANplGihG6tvd+8rYgrY= github.com/go-logr/logr v1.4.2/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= @@ -25,8 +26,14 @@ github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+ github.com/grpc-ecosystem/grpc-gateway/v2 v2.22.0 h1:asbCHRVmodnJTuQ3qamDwqVOIjwqUPTYmYuemVOx+Ys= github.com/grpc-ecosystem/grpc-gateway/v2 v2.22.0/go.mod h1:ggCgvZ2r7uOoQjOyu2Y1NhHmEPPzzuhWgcza5M1Ji1I= github.com/kisielk/sqlstruct v0.0.0-20201105191214-5f3e10d3ab46/go.mod h1:yyMNCyc/Ib3bDTKd379tNMpB/7/H5TjM2Y9QJ5THLbE= +github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= +github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= +github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= +github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/rogpeppe/go-internal v1.12.0 h1:exVL4IDcn6na9z1rAb56Vxr+CgyK3nn3O+epU5NdKM8= +github.com/rogpeppe/go-internal v1.12.0/go.mod h1:E+RYuTGaKKdloAfM02xzb0FW3Paa99yedzYV+kq4uf4= github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= github.com/vcaesar/cedar v0.20.2 h1:TDx7AdZhilKcfE1WvdToTJf5VrC/FXcUOW+KY1upLZ4= @@ -45,6 +52,8 @@ go.opentelemetry.io/otel/metric v1.29.0 h1:vPf/HFWTNkPu1aYeIsc98l4ktOQaL6LeSoeV2 go.opentelemetry.io/otel/metric v1.29.0/go.mod h1:auu/QWieFVWx+DmQOUMgj0F8LHWdgalxXqvp7BII/W8= go.opentelemetry.io/otel/sdk v1.29.0 h1:vkqKjk7gwhS8VaWb0POZKmIEDimRCMsopNYnriHyryo= go.opentelemetry.io/otel/sdk v1.29.0/go.mod h1:pM8Dx5WKnvxLCb+8lG1PRNIDxu9g9b9g59Qr7hfAAok= +go.opentelemetry.io/otel/sdk/metric v1.29.0 h1:K2CfmJohnRgvZ9UAj2/FhIf/okdWcNdBwe1m8xFXiSY= +go.opentelemetry.io/otel/sdk/metric v1.29.0/go.mod h1:6zZLdCl2fkauYoZIOn/soQIDSWFmNSRcICarHfuhNJQ= go.opentelemetry.io/otel/trace v1.29.0 h1:J/8ZNK4XgR7a21DZUAsbF8pZ5Jcw1VhACmnYt39JTi4= go.opentelemetry.io/otel/trace v1.29.0/go.mod h1:eHl3w0sp3paPkYstJOmAimxhiFXPg+MMTlEh3nsQgWQ= go.opentelemetry.io/proto/otlp v1.3.1 h1:TrMUixzpM0yuc/znrFTP9MMRh8trP93mkCiDVeXrui0= @@ -69,8 +78,9 @@ google.golang.org/grpc v1.65.0 h1:bs/cUb4lp1G5iImFFd3u5ixQzweKizoZJAwBNLR42lc= google.golang.org/grpc v1.65.0/go.mod h1:WgYC2ypjlB0EiQi6wdKixMqukr6lBc0Vo+oOgjrM5ZQ= google.golang.org/protobuf v1.34.2 h1:6xV6lTsCfpGD21XK49h7MhtcApnLqkfYgPcdHftf6hg= google.golang.org/protobuf v1.34.2/go.mod h1:qYOHts0dSfpeUzUFpOMr/WGzszTmLH+DiWniOlNbLDw= -gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= trpc.group/trpc-go/trpc-a2a-go v0.2.5 h1:X3pAlWD128LaS9TtXsUDZoJWPVuPZDkZKUecKRxmWn4= diff --git a/memory/mysqlvec/go.mod b/memory/mysqlvec/go.mod index 6afbbd5c41..28a3e0efad 100644 --- a/memory/mysqlvec/go.mod +++ b/memory/mysqlvec/go.mod @@ -16,19 +16,33 @@ require ( require ( filippo.io/edwards25519 v1.1.1 // indirect + github.com/cenkalti/backoff/v4 v4.3.0 // indirect github.com/davecgh/go-spew v1.1.1 // indirect github.com/go-ego/gse v1.0.0 // indirect + github.com/go-logr/logr v1.4.2 // indirect + github.com/go-logr/stdr v1.2.2 // indirect github.com/go-sql-driver/mysql v1.9.3 // indirect github.com/google/uuid v1.6.0 // indirect - github.com/kr/pretty v0.3.1 // indirect + github.com/grpc-ecosystem/grpc-gateway/v2 v2.22.0 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect - github.com/rogpeppe/go-internal v1.12.0 // indirect github.com/vcaesar/cedar v0.20.2 // indirect go.opentelemetry.io/otel v1.29.0 // indirect + go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.29.0 // indirect + go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.29.0 // indirect + go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.29.0 // indirect + go.opentelemetry.io/otel/metric v1.29.0 // indirect + go.opentelemetry.io/otel/sdk v1.29.0 // indirect go.opentelemetry.io/otel/trace v1.29.0 // indirect + go.opentelemetry.io/proto/otlp v1.3.1 // indirect go.uber.org/multierr v1.10.0 // indirect go.uber.org/zap v1.27.0 // indirect - gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect + golang.org/x/net v0.34.0 // indirect + golang.org/x/sys v0.30.0 // indirect + golang.org/x/text v0.21.0 // indirect + google.golang.org/genproto/googleapis/api v0.0.0-20240822170219-fc7c04adadcd // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20240822170219-fc7c04adadcd // indirect + google.golang.org/grpc v1.65.0 // indirect + google.golang.org/protobuf v1.34.2 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect trpc.group/trpc-go/trpc-a2a-go v0.2.5 // indirect ) diff --git a/memory/mysqlvec/go.sum b/memory/mysqlvec/go.sum index fcfbd022e4..682b5de524 100644 --- a/memory/mysqlvec/go.sum +++ b/memory/mysqlvec/go.sum @@ -6,13 +6,13 @@ github.com/bmatcuk/doublestar/v4 v4.9.1 h1:X8jg9rRZmJd4yRy7ZeNDRnM+T3ZfHv15JiBJ/ github.com/bmatcuk/doublestar/v4 v4.9.1/go.mod h1:xBQ8jztBU6kakFMg+8WGxn0c6z1fTSPVIjEY1Wr7jzc= github.com/cenkalti/backoff/v4 v4.3.0 h1:MyRJ/UdXutAwSAT+s3wNd7MfTIcy71VQueUuFK343L8= github.com/cenkalti/backoff/v4 v4.3.0/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE= -github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= github.com/creack/pty v1.1.24 h1:bJrF4RRfyJnbTJqzRLHzcGaZK1NeM5kTC9jGgovnR1s= github.com/creack/pty v1.1.24/go.mod h1:08sCNb52WyoAwi2QDyzUCTgcvVFhUzewun7wtTfvcwE= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/go-ego/gse v1.0.0 h1:GNbtH1WP7Yd1VvCZ85fIK6eVEe7RctmgmnwliEPUMNA= github.com/go-ego/gse v1.0.0/go.mod h1:Gt3A9Ry1Eso2Kza4MRaiZ7f2DTAvActmETY46Lxg0gU= +github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= github.com/go-logr/logr v1.4.2 h1:6pFjapn8bFcIbiKo3XT4j/BhANplGihG6tvd+8rYgrY= github.com/go-logr/logr v1.4.2/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= @@ -26,17 +26,12 @@ github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+ github.com/grpc-ecosystem/grpc-gateway/v2 v2.22.0 h1:asbCHRVmodnJTuQ3qamDwqVOIjwqUPTYmYuemVOx+Ys= github.com/grpc-ecosystem/grpc-gateway/v2 v2.22.0/go.mod h1:ggCgvZ2r7uOoQjOyu2Y1NhHmEPPzzuhWgcza5M1Ji1I= github.com/kisielk/sqlstruct v0.0.0-20201105191214-5f3e10d3ab46/go.mod h1:yyMNCyc/Ib3bDTKd379tNMpB/7/H5TjM2Y9QJ5THLbE= -github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= -github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= -github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= -github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= -github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs= github.com/rogpeppe/go-internal v1.12.0 h1:exVL4IDcn6na9z1rAb56Vxr+CgyK3nn3O+epU5NdKM8= github.com/rogpeppe/go-internal v1.12.0/go.mod h1:E+RYuTGaKKdloAfM02xzb0FW3Paa99yedzYV+kq4uf4= github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= @@ -57,6 +52,8 @@ go.opentelemetry.io/otel/metric v1.29.0 h1:vPf/HFWTNkPu1aYeIsc98l4ktOQaL6LeSoeV2 go.opentelemetry.io/otel/metric v1.29.0/go.mod h1:auu/QWieFVWx+DmQOUMgj0F8LHWdgalxXqvp7BII/W8= go.opentelemetry.io/otel/sdk v1.29.0 h1:vkqKjk7gwhS8VaWb0POZKmIEDimRCMsopNYnriHyryo= go.opentelemetry.io/otel/sdk v1.29.0/go.mod h1:pM8Dx5WKnvxLCb+8lG1PRNIDxu9g9b9g59Qr7hfAAok= +go.opentelemetry.io/otel/sdk/metric v1.29.0 h1:K2CfmJohnRgvZ9UAj2/FhIf/okdWcNdBwe1m8xFXiSY= +go.opentelemetry.io/otel/sdk/metric v1.29.0/go.mod h1:6zZLdCl2fkauYoZIOn/soQIDSWFmNSRcICarHfuhNJQ= go.opentelemetry.io/otel/trace v1.29.0 h1:J/8ZNK4XgR7a21DZUAsbF8pZ5Jcw1VhACmnYt39JTi4= go.opentelemetry.io/otel/trace v1.29.0/go.mod h1:eHl3w0sp3paPkYstJOmAimxhiFXPg+MMTlEh3nsQgWQ= go.opentelemetry.io/proto/otlp v1.3.1 h1:TrMUixzpM0yuc/znrFTP9MMRh8trP93mkCiDVeXrui0= diff --git a/memory/pgvector/go.mod b/memory/pgvector/go.mod index bb5d14935a..16603aa5ab 100644 --- a/memory/pgvector/go.mod +++ b/memory/pgvector/go.mod @@ -17,9 +17,13 @@ require ( ) require ( + github.com/cenkalti/backoff/v4 v4.3.0 // indirect github.com/davecgh/go-spew v1.1.1 // indirect github.com/go-ego/gse v1.0.0 // indirect + github.com/go-logr/logr v1.4.2 // indirect + github.com/go-logr/stdr v1.2.2 // indirect github.com/google/uuid v1.6.0 // indirect + github.com/grpc-ecosystem/grpc-gateway/v2 v2.22.0 // indirect github.com/jackc/pgpassfile v1.0.0 // indirect github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect github.com/jackc/pgx/v5 v5.7.1 // indirect @@ -27,12 +31,24 @@ require ( github.com/pmezard/go-difflib v1.0.0 // indirect github.com/vcaesar/cedar v0.20.2 // indirect go.opentelemetry.io/otel v1.29.0 // indirect + go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.29.0 // indirect + go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.29.0 // indirect + go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.29.0 // indirect + go.opentelemetry.io/otel/metric v1.29.0 // indirect + go.opentelemetry.io/otel/sdk v1.29.0 // indirect go.opentelemetry.io/otel/trace v1.29.0 // indirect + go.opentelemetry.io/proto/otlp v1.3.1 // indirect go.uber.org/multierr v1.10.0 // indirect go.uber.org/zap v1.27.0 // indirect golang.org/x/crypto v0.32.0 // indirect + golang.org/x/net v0.34.0 // indirect golang.org/x/sync v0.10.0 // indirect + golang.org/x/sys v0.30.0 // indirect golang.org/x/text v0.21.0 // indirect + google.golang.org/genproto/googleapis/api v0.0.0-20240822170219-fc7c04adadcd // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20240822170219-fc7c04adadcd // indirect + google.golang.org/grpc v1.65.0 // indirect + google.golang.org/protobuf v1.34.2 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect trpc.group/trpc-go/trpc-a2a-go v0.2.5 // indirect ) diff --git a/memory/pgvector/go.sum b/memory/pgvector/go.sum index b649ad606c..de1e5ecd79 100644 --- a/memory/pgvector/go.sum +++ b/memory/pgvector/go.sum @@ -13,6 +13,7 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/go-ego/gse v1.0.0 h1:GNbtH1WP7Yd1VvCZ85fIK6eVEe7RctmgmnwliEPUMNA= github.com/go-ego/gse v1.0.0/go.mod h1:Gt3A9Ry1Eso2Kza4MRaiZ7f2DTAvActmETY46Lxg0gU= +github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= github.com/go-logr/logr v1.4.2 h1:6pFjapn8bFcIbiKo3XT4j/BhANplGihG6tvd+8rYgrY= github.com/go-logr/logr v1.4.2/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= @@ -42,8 +43,8 @@ github.com/jinzhu/now v1.1.5/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/ github.com/jmoiron/sqlx v1.3.5 h1:vFFPA71p1o5gAeqtEAwLU4dnX2napprKtHr7PYIcN3g= github.com/jmoiron/sqlx v1.3.5/go.mod h1:nRVWtLre0KfCLJvgxzCsLVMogSvQ1zNJtpYr2Ccp0mQ= github.com/kisielk/sqlstruct v0.0.0-20201105191214-5f3e10d3ab46/go.mod h1:yyMNCyc/Ib3bDTKd379tNMpB/7/H5TjM2Y9QJ5THLbE= -github.com/kr/pretty v0.3.0 h1:WgNl7dwNpEZ6jJ9k1snq4pZsg7DOEN8hP9Xw0Tsjwk0= -github.com/kr/pretty v0.3.0/go.mod h1:640gp4NfQd8pI5XOwp5fnNeVWj67G7CFk/SaSQn7NBk= +github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= +github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/lib/pq v1.10.9 h1:YXG7RB+JIjhP29X+OtkiDnYaXQwpS4JEWq7dtCCRUEw= @@ -91,6 +92,8 @@ go.opentelemetry.io/otel/metric v1.29.0 h1:vPf/HFWTNkPu1aYeIsc98l4ktOQaL6LeSoeV2 go.opentelemetry.io/otel/metric v1.29.0/go.mod h1:auu/QWieFVWx+DmQOUMgj0F8LHWdgalxXqvp7BII/W8= go.opentelemetry.io/otel/sdk v1.29.0 h1:vkqKjk7gwhS8VaWb0POZKmIEDimRCMsopNYnriHyryo= go.opentelemetry.io/otel/sdk v1.29.0/go.mod h1:pM8Dx5WKnvxLCb+8lG1PRNIDxu9g9b9g59Qr7hfAAok= +go.opentelemetry.io/otel/sdk/metric v1.29.0 h1:K2CfmJohnRgvZ9UAj2/FhIf/okdWcNdBwe1m8xFXiSY= +go.opentelemetry.io/otel/sdk/metric v1.29.0/go.mod h1:6zZLdCl2fkauYoZIOn/soQIDSWFmNSRcICarHfuhNJQ= go.opentelemetry.io/otel/trace v1.29.0 h1:J/8ZNK4XgR7a21DZUAsbF8pZ5Jcw1VhACmnYt39JTi4= go.opentelemetry.io/otel/trace v1.29.0/go.mod h1:eHl3w0sp3paPkYstJOmAimxhiFXPg+MMTlEh3nsQgWQ= go.opentelemetry.io/proto/otlp v1.3.1 h1:TrMUixzpM0yuc/znrFTP9MMRh8trP93mkCiDVeXrui0= diff --git a/memory/postgres/go.mod b/memory/postgres/go.mod index 11d8739bf6..c35f36184c 100644 --- a/memory/postgres/go.mod +++ b/memory/postgres/go.mod @@ -15,9 +15,13 @@ require ( ) require ( + github.com/cenkalti/backoff/v4 v4.3.0 // indirect github.com/davecgh/go-spew v1.1.1 // indirect github.com/go-ego/gse v1.0.0 // indirect + github.com/go-logr/logr v1.4.2 // indirect + github.com/go-logr/stdr v1.2.2 // indirect github.com/google/uuid v1.6.0 // indirect + github.com/grpc-ecosystem/grpc-gateway/v2 v2.22.0 // indirect github.com/jackc/pgpassfile v1.0.0 // indirect github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect github.com/jackc/pgx/v5 v5.7.2 // indirect @@ -25,12 +29,24 @@ require ( github.com/pmezard/go-difflib v1.0.0 // indirect github.com/vcaesar/cedar v0.20.2 // indirect go.opentelemetry.io/otel v1.29.0 // indirect + go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.29.0 // indirect + go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.29.0 // indirect + go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.29.0 // indirect + go.opentelemetry.io/otel/metric v1.29.0 // indirect + go.opentelemetry.io/otel/sdk v1.29.0 // indirect go.opentelemetry.io/otel/trace v1.29.0 // indirect + go.opentelemetry.io/proto/otlp v1.3.1 // indirect go.uber.org/multierr v1.10.0 // indirect go.uber.org/zap v1.27.0 // indirect golang.org/x/crypto v0.32.0 // indirect + golang.org/x/net v0.34.0 // indirect golang.org/x/sync v0.10.0 // indirect + golang.org/x/sys v0.30.0 // indirect golang.org/x/text v0.21.0 // indirect + google.golang.org/genproto/googleapis/api v0.0.0-20240822170219-fc7c04adadcd // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20240822170219-fc7c04adadcd // indirect + google.golang.org/grpc v1.65.0 // indirect + google.golang.org/protobuf v1.34.2 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect trpc.group/trpc-go/trpc-a2a-go v0.2.5 // indirect ) diff --git a/memory/postgres/go.sum b/memory/postgres/go.sum index 4d19483438..f2eaee8d33 100644 --- a/memory/postgres/go.sum +++ b/memory/postgres/go.sum @@ -11,6 +11,7 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/go-ego/gse v1.0.0 h1:GNbtH1WP7Yd1VvCZ85fIK6eVEe7RctmgmnwliEPUMNA= github.com/go-ego/gse v1.0.0/go.mod h1:Gt3A9Ry1Eso2Kza4MRaiZ7f2DTAvActmETY46Lxg0gU= +github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= github.com/go-logr/logr v1.4.2 h1:6pFjapn8bFcIbiKo3XT4j/BhANplGihG6tvd+8rYgrY= github.com/go-logr/logr v1.4.2/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= @@ -30,8 +31,8 @@ github.com/jackc/pgx/v5 v5.7.2/go.mod h1:ncY89UGWxg82EykZUwSpUKEfccBGGYq1xjrOpsb github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo= github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= github.com/kisielk/sqlstruct v0.0.0-20201105191214-5f3e10d3ab46/go.mod h1:yyMNCyc/Ib3bDTKd379tNMpB/7/H5TjM2Y9QJ5THLbE= -github.com/kr/pretty v0.3.0 h1:WgNl7dwNpEZ6jJ9k1snq4pZsg7DOEN8hP9Xw0Tsjwk0= -github.com/kr/pretty v0.3.0/go.mod h1:640gp4NfQd8pI5XOwp5fnNeVWj67G7CFk/SaSQn7NBk= +github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= +github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= @@ -59,6 +60,8 @@ go.opentelemetry.io/otel/metric v1.29.0 h1:vPf/HFWTNkPu1aYeIsc98l4ktOQaL6LeSoeV2 go.opentelemetry.io/otel/metric v1.29.0/go.mod h1:auu/QWieFVWx+DmQOUMgj0F8LHWdgalxXqvp7BII/W8= go.opentelemetry.io/otel/sdk v1.29.0 h1:vkqKjk7gwhS8VaWb0POZKmIEDimRCMsopNYnriHyryo= go.opentelemetry.io/otel/sdk v1.29.0/go.mod h1:pM8Dx5WKnvxLCb+8lG1PRNIDxu9g9b9g59Qr7hfAAok= +go.opentelemetry.io/otel/sdk/metric v1.29.0 h1:K2CfmJohnRgvZ9UAj2/FhIf/okdWcNdBwe1m8xFXiSY= +go.opentelemetry.io/otel/sdk/metric v1.29.0/go.mod h1:6zZLdCl2fkauYoZIOn/soQIDSWFmNSRcICarHfuhNJQ= go.opentelemetry.io/otel/trace v1.29.0 h1:J/8ZNK4XgR7a21DZUAsbF8pZ5Jcw1VhACmnYt39JTi4= go.opentelemetry.io/otel/trace v1.29.0/go.mod h1:eHl3w0sp3paPkYstJOmAimxhiFXPg+MMTlEh3nsQgWQ= go.opentelemetry.io/proto/otlp v1.3.1 h1:TrMUixzpM0yuc/znrFTP9MMRh8trP93mkCiDVeXrui0= diff --git a/memory/redis/go.mod b/memory/redis/go.mod index b90caded99..b0d4e4663a 100644 --- a/memory/redis/go.mod +++ b/memory/redis/go.mod @@ -16,18 +16,35 @@ require ( ) require ( + github.com/cenkalti/backoff/v4 v4.3.0 // indirect github.com/cespare/xxhash/v2 v2.3.0 // indirect github.com/davecgh/go-spew v1.1.1 // indirect github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect github.com/go-ego/gse v1.0.0 // indirect + github.com/go-logr/logr v1.4.2 // indirect + github.com/go-logr/stdr v1.2.2 // indirect github.com/google/uuid v1.6.0 // indirect + github.com/grpc-ecosystem/grpc-gateway/v2 v2.22.0 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect github.com/vcaesar/cedar v0.20.2 // indirect github.com/yuin/gopher-lua v1.1.1 // indirect go.opentelemetry.io/otel v1.29.0 // indirect + go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.29.0 // indirect + go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.29.0 // indirect + go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.29.0 // indirect + go.opentelemetry.io/otel/metric v1.29.0 // indirect + go.opentelemetry.io/otel/sdk v1.29.0 // indirect go.opentelemetry.io/otel/trace v1.29.0 // indirect + go.opentelemetry.io/proto/otlp v1.3.1 // indirect go.uber.org/multierr v1.10.0 // indirect go.uber.org/zap v1.27.0 // indirect + golang.org/x/net v0.34.0 // indirect + golang.org/x/sys v0.30.0 // indirect + golang.org/x/text v0.21.0 // indirect + google.golang.org/genproto/googleapis/api v0.0.0-20240822170219-fc7c04adadcd // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20240822170219-fc7c04adadcd // indirect + google.golang.org/grpc v1.65.0 // indirect + google.golang.org/protobuf v1.34.2 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect trpc.group/trpc-go/trpc-a2a-go v0.2.5 // indirect ) diff --git a/memory/redis/go.sum b/memory/redis/go.sum index c767ca0caa..f93c4b6794 100644 --- a/memory/redis/go.sum +++ b/memory/redis/go.sum @@ -18,6 +18,7 @@ github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/r github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= github.com/go-ego/gse v1.0.0 h1:GNbtH1WP7Yd1VvCZ85fIK6eVEe7RctmgmnwliEPUMNA= github.com/go-ego/gse v1.0.0/go.mod h1:Gt3A9Ry1Eso2Kza4MRaiZ7f2DTAvActmETY46Lxg0gU= +github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= github.com/go-logr/logr v1.4.2 h1:6pFjapn8bFcIbiKo3XT4j/BhANplGihG6tvd+8rYgrY= github.com/go-logr/logr v1.4.2/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= @@ -28,10 +29,16 @@ github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/grpc-ecosystem/grpc-gateway/v2 v2.22.0 h1:asbCHRVmodnJTuQ3qamDwqVOIjwqUPTYmYuemVOx+Ys= github.com/grpc-ecosystem/grpc-gateway/v2 v2.22.0/go.mod h1:ggCgvZ2r7uOoQjOyu2Y1NhHmEPPzzuhWgcza5M1Ji1I= +github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= +github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= +github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= +github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/redis/go-redis/v9 v9.11.0 h1:E3S08Gl/nJNn5vkxd2i78wZxWAPNZgUNTp8WIJUAiIs= github.com/redis/go-redis/v9 v9.11.0/go.mod h1:huWgSWd8mW6+m0VPhJjSSQ+d6Nh1VICQ6Q5lHuCH/Iw= +github.com/rogpeppe/go-internal v1.12.0 h1:exVL4IDcn6na9z1rAb56Vxr+CgyK3nn3O+epU5NdKM8= +github.com/rogpeppe/go-internal v1.12.0/go.mod h1:E+RYuTGaKKdloAfM02xzb0FW3Paa99yedzYV+kq4uf4= github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= github.com/vcaesar/cedar v0.20.2 h1:TDx7AdZhilKcfE1WvdToTJf5VrC/FXcUOW+KY1upLZ4= @@ -52,6 +59,8 @@ go.opentelemetry.io/otel/metric v1.29.0 h1:vPf/HFWTNkPu1aYeIsc98l4ktOQaL6LeSoeV2 go.opentelemetry.io/otel/metric v1.29.0/go.mod h1:auu/QWieFVWx+DmQOUMgj0F8LHWdgalxXqvp7BII/W8= go.opentelemetry.io/otel/sdk v1.29.0 h1:vkqKjk7gwhS8VaWb0POZKmIEDimRCMsopNYnriHyryo= go.opentelemetry.io/otel/sdk v1.29.0/go.mod h1:pM8Dx5WKnvxLCb+8lG1PRNIDxu9g9b9g59Qr7hfAAok= +go.opentelemetry.io/otel/sdk/metric v1.29.0 h1:K2CfmJohnRgvZ9UAj2/FhIf/okdWcNdBwe1m8xFXiSY= +go.opentelemetry.io/otel/sdk/metric v1.29.0/go.mod h1:6zZLdCl2fkauYoZIOn/soQIDSWFmNSRcICarHfuhNJQ= go.opentelemetry.io/otel/trace v1.29.0 h1:J/8ZNK4XgR7a21DZUAsbF8pZ5Jcw1VhACmnYt39JTi4= go.opentelemetry.io/otel/trace v1.29.0/go.mod h1:eHl3w0sp3paPkYstJOmAimxhiFXPg+MMTlEh3nsQgWQ= go.opentelemetry.io/proto/otlp v1.3.1 h1:TrMUixzpM0yuc/znrFTP9MMRh8trP93mkCiDVeXrui0= @@ -76,8 +85,9 @@ google.golang.org/grpc v1.65.0 h1:bs/cUb4lp1G5iImFFd3u5ixQzweKizoZJAwBNLR42lc= google.golang.org/grpc v1.65.0/go.mod h1:WgYC2ypjlB0EiQi6wdKixMqukr6lBc0Vo+oOgjrM5ZQ= google.golang.org/protobuf v1.34.2 h1:6xV6lTsCfpGD21XK49h7MhtcApnLqkfYgPcdHftf6hg= google.golang.org/protobuf v1.34.2/go.mod h1:qYOHts0dSfpeUzUFpOMr/WGzszTmLH+DiWniOlNbLDw= -gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= trpc.group/trpc-go/trpc-a2a-go v0.2.5 h1:X3pAlWD128LaS9TtXsUDZoJWPVuPZDkZKUecKRxmWn4= diff --git a/memory/sqlite/go.mod b/memory/sqlite/go.mod index c0333151b5..9b7b50380e 100644 --- a/memory/sqlite/go.mod +++ b/memory/sqlite/go.mod @@ -11,15 +11,32 @@ require ( ) require ( + github.com/cenkalti/backoff/v4 v4.3.0 // indirect github.com/davecgh/go-spew v1.1.1 // indirect github.com/go-ego/gse v1.0.0 // indirect + github.com/go-logr/logr v1.4.2 // indirect + github.com/go-logr/stdr v1.2.2 // indirect github.com/google/uuid v1.6.0 // indirect + github.com/grpc-ecosystem/grpc-gateway/v2 v2.22.0 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect github.com/vcaesar/cedar v0.20.2 // indirect go.opentelemetry.io/otel v1.29.0 // indirect + go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.29.0 // indirect + go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.29.0 // indirect + go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.29.0 // indirect + go.opentelemetry.io/otel/metric v1.29.0 // indirect + go.opentelemetry.io/otel/sdk v1.29.0 // indirect go.opentelemetry.io/otel/trace v1.29.0 // indirect + go.opentelemetry.io/proto/otlp v1.3.1 // indirect go.uber.org/multierr v1.10.0 // indirect go.uber.org/zap v1.27.0 // indirect + golang.org/x/net v0.34.0 // indirect + golang.org/x/sys v0.30.0 // indirect + golang.org/x/text v0.21.0 // indirect + google.golang.org/genproto/googleapis/api v0.0.0-20240822170219-fc7c04adadcd // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20240822170219-fc7c04adadcd // indirect + google.golang.org/grpc v1.65.0 // indirect + google.golang.org/protobuf v1.34.2 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect trpc.group/trpc-go/trpc-a2a-go v0.2.5 // indirect ) diff --git a/memory/sqlite/go.sum b/memory/sqlite/go.sum index 563a80b39a..4029d8cac3 100644 --- a/memory/sqlite/go.sum +++ b/memory/sqlite/go.sum @@ -8,6 +8,7 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/go-ego/gse v1.0.0 h1:GNbtH1WP7Yd1VvCZ85fIK6eVEe7RctmgmnwliEPUMNA= github.com/go-ego/gse v1.0.0/go.mod h1:Gt3A9Ry1Eso2Kza4MRaiZ7f2DTAvActmETY46Lxg0gU= +github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= github.com/go-logr/logr v1.4.2 h1:6pFjapn8bFcIbiKo3XT4j/BhANplGihG6tvd+8rYgrY= github.com/go-logr/logr v1.4.2/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= @@ -18,10 +19,16 @@ github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/grpc-ecosystem/grpc-gateway/v2 v2.22.0 h1:asbCHRVmodnJTuQ3qamDwqVOIjwqUPTYmYuemVOx+Ys= github.com/grpc-ecosystem/grpc-gateway/v2 v2.22.0/go.mod h1:ggCgvZ2r7uOoQjOyu2Y1NhHmEPPzzuhWgcza5M1Ji1I= +github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= +github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= +github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= +github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/mattn/go-sqlite3 v1.14.32 h1:JD12Ag3oLy1zQA+BNn74xRgaBbdhbNIDYvQUEuuErjs= github.com/mattn/go-sqlite3 v1.14.32/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/rogpeppe/go-internal v1.12.0 h1:exVL4IDcn6na9z1rAb56Vxr+CgyK3nn3O+epU5NdKM8= +github.com/rogpeppe/go-internal v1.12.0/go.mod h1:E+RYuTGaKKdloAfM02xzb0FW3Paa99yedzYV+kq4uf4= github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= github.com/vcaesar/cedar v0.20.2 h1:TDx7AdZhilKcfE1WvdToTJf5VrC/FXcUOW+KY1upLZ4= @@ -40,6 +47,8 @@ go.opentelemetry.io/otel/metric v1.29.0 h1:vPf/HFWTNkPu1aYeIsc98l4ktOQaL6LeSoeV2 go.opentelemetry.io/otel/metric v1.29.0/go.mod h1:auu/QWieFVWx+DmQOUMgj0F8LHWdgalxXqvp7BII/W8= go.opentelemetry.io/otel/sdk v1.29.0 h1:vkqKjk7gwhS8VaWb0POZKmIEDimRCMsopNYnriHyryo= go.opentelemetry.io/otel/sdk v1.29.0/go.mod h1:pM8Dx5WKnvxLCb+8lG1PRNIDxu9g9b9g59Qr7hfAAok= +go.opentelemetry.io/otel/sdk/metric v1.29.0 h1:K2CfmJohnRgvZ9UAj2/FhIf/okdWcNdBwe1m8xFXiSY= +go.opentelemetry.io/otel/sdk/metric v1.29.0/go.mod h1:6zZLdCl2fkauYoZIOn/soQIDSWFmNSRcICarHfuhNJQ= go.opentelemetry.io/otel/trace v1.29.0 h1:J/8ZNK4XgR7a21DZUAsbF8pZ5Jcw1VhACmnYt39JTi4= go.opentelemetry.io/otel/trace v1.29.0/go.mod h1:eHl3w0sp3paPkYstJOmAimxhiFXPg+MMTlEh3nsQgWQ= go.opentelemetry.io/proto/otlp v1.3.1 h1:TrMUixzpM0yuc/znrFTP9MMRh8trP93mkCiDVeXrui0= @@ -64,8 +73,9 @@ google.golang.org/grpc v1.65.0 h1:bs/cUb4lp1G5iImFFd3u5ixQzweKizoZJAwBNLR42lc= google.golang.org/grpc v1.65.0/go.mod h1:WgYC2ypjlB0EiQi6wdKixMqukr6lBc0Vo+oOgjrM5ZQ= google.golang.org/protobuf v1.34.2 h1:6xV6lTsCfpGD21XK49h7MhtcApnLqkfYgPcdHftf6hg= google.golang.org/protobuf v1.34.2/go.mod h1:qYOHts0dSfpeUzUFpOMr/WGzszTmLH+DiWniOlNbLDw= -gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= trpc.group/trpc-go/trpc-a2a-go v0.2.5 h1:X3pAlWD128LaS9TtXsUDZoJWPVuPZDkZKUecKRxmWn4= diff --git a/memory/sqlitevec/go.mod b/memory/sqlitevec/go.mod index a8e3769777..389e5a60e8 100644 --- a/memory/sqlitevec/go.mod +++ b/memory/sqlitevec/go.mod @@ -14,19 +14,35 @@ require ( ) require ( + github.com/cenkalti/backoff/v4 v4.3.0 // indirect github.com/davecgh/go-spew v1.1.1 // indirect github.com/go-ego/gse v1.0.0 // indirect + github.com/go-logr/logr v1.4.2 // indirect + github.com/go-logr/stdr v1.2.2 // indirect github.com/google/uuid v1.6.0 // indirect + github.com/grpc-ecosystem/grpc-gateway/v2 v2.22.0 // indirect github.com/ncruces/go-sqlite3 v0.32.0 // indirect github.com/ncruces/julianday v1.0.0 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect github.com/tetratelabs/wazero v1.11.0 // indirect github.com/vcaesar/cedar v0.20.2 // indirect go.opentelemetry.io/otel v1.29.0 // indirect + go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.29.0 // indirect + go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.29.0 // indirect + go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.29.0 // indirect + go.opentelemetry.io/otel/metric v1.29.0 // indirect + go.opentelemetry.io/otel/sdk v1.29.0 // indirect go.opentelemetry.io/otel/trace v1.29.0 // indirect + go.opentelemetry.io/proto/otlp v1.3.1 // indirect go.uber.org/multierr v1.10.0 // indirect go.uber.org/zap v1.27.0 // indirect + golang.org/x/net v0.34.0 // indirect golang.org/x/sys v0.41.0 // indirect + golang.org/x/text v0.34.0 // indirect + google.golang.org/genproto/googleapis/api v0.0.0-20240822170219-fc7c04adadcd // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20240822170219-fc7c04adadcd // indirect + google.golang.org/grpc v1.65.0 // indirect + google.golang.org/protobuf v1.34.2 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect trpc.group/trpc-go/trpc-a2a-go v0.2.5 // indirect ) diff --git a/memory/sqlitevec/go.sum b/memory/sqlitevec/go.sum index 0a407e46b2..b08dcf26df 100644 --- a/memory/sqlitevec/go.sum +++ b/memory/sqlitevec/go.sum @@ -10,6 +10,7 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/go-ego/gse v1.0.0 h1:GNbtH1WP7Yd1VvCZ85fIK6eVEe7RctmgmnwliEPUMNA= github.com/go-ego/gse v1.0.0/go.mod h1:Gt3A9Ry1Eso2Kza4MRaiZ7f2DTAvActmETY46Lxg0gU= +github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= github.com/go-logr/logr v1.4.2 h1:6pFjapn8bFcIbiKo3XT4j/BhANplGihG6tvd+8rYgrY= github.com/go-logr/logr v1.4.2/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= @@ -20,6 +21,10 @@ github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/grpc-ecosystem/grpc-gateway/v2 v2.22.0 h1:asbCHRVmodnJTuQ3qamDwqVOIjwqUPTYmYuemVOx+Ys= github.com/grpc-ecosystem/grpc-gateway/v2 v2.22.0/go.mod h1:ggCgvZ2r7uOoQjOyu2Y1NhHmEPPzzuhWgcza5M1Ji1I= +github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= +github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= +github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= +github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/mattn/go-sqlite3 v1.14.32 h1:JD12Ag3oLy1zQA+BNn74xRgaBbdhbNIDYvQUEuuErjs= github.com/mattn/go-sqlite3 v1.14.32/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y= github.com/ncruces/go-sqlite3 v0.32.0 h1:hNBUXp88LrfQCsuyXLqWTbTUG35sUuktDsqhhgHvU20= @@ -28,6 +33,8 @@ github.com/ncruces/julianday v1.0.0 h1:fH0OKwa7NWvniGQtxdJRxAgkBMolni2BjDHaWTxqt github.com/ncruces/julianday v1.0.0/go.mod h1:Dusn2KvZrrovOMJuOt0TNXL6tB7U2E8kvza5fFc9G7g= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/rogpeppe/go-internal v1.12.0 h1:exVL4IDcn6na9z1rAb56Vxr+CgyK3nn3O+epU5NdKM8= +github.com/rogpeppe/go-internal v1.12.0/go.mod h1:E+RYuTGaKKdloAfM02xzb0FW3Paa99yedzYV+kq4uf4= github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= github.com/tetratelabs/wazero v1.11.0 h1:+gKemEuKCTevU4d7ZTzlsvgd1uaToIDtlQlmNbwqYhA= @@ -48,6 +55,8 @@ go.opentelemetry.io/otel/metric v1.29.0 h1:vPf/HFWTNkPu1aYeIsc98l4ktOQaL6LeSoeV2 go.opentelemetry.io/otel/metric v1.29.0/go.mod h1:auu/QWieFVWx+DmQOUMgj0F8LHWdgalxXqvp7BII/W8= go.opentelemetry.io/otel/sdk v1.29.0 h1:vkqKjk7gwhS8VaWb0POZKmIEDimRCMsopNYnriHyryo= go.opentelemetry.io/otel/sdk v1.29.0/go.mod h1:pM8Dx5WKnvxLCb+8lG1PRNIDxu9g9b9g59Qr7hfAAok= +go.opentelemetry.io/otel/sdk/metric v1.29.0 h1:K2CfmJohnRgvZ9UAj2/FhIf/okdWcNdBwe1m8xFXiSY= +go.opentelemetry.io/otel/sdk/metric v1.29.0/go.mod h1:6zZLdCl2fkauYoZIOn/soQIDSWFmNSRcICarHfuhNJQ= go.opentelemetry.io/otel/trace v1.29.0 h1:J/8ZNK4XgR7a21DZUAsbF8pZ5Jcw1VhACmnYt39JTi4= go.opentelemetry.io/otel/trace v1.29.0/go.mod h1:eHl3w0sp3paPkYstJOmAimxhiFXPg+MMTlEh3nsQgWQ= go.opentelemetry.io/proto/otlp v1.3.1 h1:TrMUixzpM0yuc/znrFTP9MMRh8trP93mkCiDVeXrui0= @@ -72,8 +81,9 @@ google.golang.org/grpc v1.65.0 h1:bs/cUb4lp1G5iImFFd3u5ixQzweKizoZJAwBNLR42lc= google.golang.org/grpc v1.65.0/go.mod h1:WgYC2ypjlB0EiQi6wdKixMqukr6lBc0Vo+oOgjrM5ZQ= google.golang.org/protobuf v1.34.2 h1:6xV6lTsCfpGD21XK49h7MhtcApnLqkfYgPcdHftf6hg= google.golang.org/protobuf v1.34.2/go.mod h1:qYOHts0dSfpeUzUFpOMr/WGzszTmLH+DiWniOlNbLDw= -gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= trpc.group/trpc-go/trpc-a2a-go v0.2.5 h1:X3pAlWD128LaS9TtXsUDZoJWPVuPZDkZKUecKRxmWn4= diff --git a/platform/backend_migration_status.go b/platform/backend_migration_status.go index 67ab59b735..b70aec8ddb 100644 --- a/platform/backend_migration_status.go +++ b/platform/backend_migration_status.go @@ -140,6 +140,16 @@ func NewBackendMigrationStatusReport(input BackendMigrationStatusInput) (Backend // Validate checks that a backend migration status report is safe to expose or store. func (r BackendMigrationStatusReport) Validate() error { + if err := r.validateBackendMigrationIdentity(); err != nil { + return err + } + if err := r.validateBackendMigrationState(); err != nil { + return err + } + return r.validateBackendMigrationSafeText() +} + +func (r BackendMigrationStatusReport) validateBackendMigrationIdentity() error { if strings.TrimSpace(r.TenantID) == "" { return ErrTenantIDRequired } @@ -175,6 +185,10 @@ func (r BackendMigrationStatusReport) Validate() error { if strings.TrimSpace(r.SourceBackendID) == strings.TrimSpace(r.TargetBackendID) { return fmt.Errorf("source_backend_id and target_backend_id must differ") } + return nil +} + +func (r BackendMigrationStatusReport) validateBackendMigrationState() error { mode, err := NormalizeStorageMigrationMode(string(r.MigrationMode)) if err != nil { return err @@ -200,6 +214,10 @@ func (r BackendMigrationStatusReport) Validate() error { if r.UpdatedAt.IsZero() { return fmt.Errorf("updated_at is required") } + return nil +} + +func (r BackendMigrationStatusReport) validateBackendMigrationSafeText() error { for field, value := range map[string]string{ "app_id": r.AppID, "profile_id": r.ProfileID, diff --git a/platform/budget_audit.go b/platform/budget_audit.go index 77e882da06..efee53372f 100644 --- a/platform/budget_audit.go +++ b/platform/budget_audit.go @@ -126,6 +126,26 @@ func NewBudgetDecisionAuditRecord(input BudgetDecisionAuditInput) (AuditRecord, // Validate checks that the summary is complete and safe to expose. func (s BudgetDecisionSummary) Validate() error { + if err := s.validateIdentityAndText(); err != nil { + return err + } + expected, err := s.validateEstimateAndQuota() + if err != nil { + return err + } + if err := s.validateOutcome(expected); err != nil { + return err + } + if strings.TrimSpace(s.RedactionVersion) == "" { + return fmt.Errorf("redaction_version is required") + } + if s.CreatedAt.IsZero() { + return fmt.Errorf("created_at is required") + } + return validateAuditRedactedText("redacted_detail_ref", s.DetailRef()) +} + +func (s BudgetDecisionSummary) validateIdentityAndText() error { if strings.TrimSpace(s.TenantID) == "" { return ErrTenantIDRequired } @@ -147,6 +167,10 @@ func (s BudgetDecisionSummary) Validate() error { if strings.TrimSpace(s.Reason) == "" && s.Outcome != BudgetDecisionOutcomeAllow { return fmt.Errorf("reason is required for non-allow budget outcomes") } + return nil +} + +func (s BudgetDecisionSummary) validateEstimateAndQuota() (BudgetDecision, error) { estimate := UsageEstimate{ PromptTokens: s.EstimatedPrompt, CompletionTokens: s.EstimatedCompletion, @@ -154,23 +178,27 @@ func (s BudgetDecisionSummary) Validate() error { Cost: s.EstimatedCost, } if err := validateUsageEstimate(estimate); err != nil { - return err + return BudgetDecision{}, err } canonicalTotal, err := estimate.effectiveTotalTokens() if err != nil { - return err + return BudgetDecision{}, err } if canonicalTotal != s.EstimatedTotalTokens { - return fmt.Errorf("estimated_total_tokens must match effective total tokens") + return BudgetDecision{}, fmt.Errorf("estimated_total_tokens must match effective total tokens") } quota := s.quota() if err := quota.Validate(); err != nil { - return err + return BudgetDecision{}, err } expected, err := quota.Check(estimate) if err != nil { - return err + return BudgetDecision{}, err } + return expected, nil +} + +func (s BudgetDecisionSummary) validateOutcome(expected BudgetDecision) error { switch s.Outcome { case BudgetDecisionOutcomeAllow: if !expected.Allowed { @@ -207,15 +235,6 @@ func (s BudgetDecisionSummary) Validate() error { default: return fmt.Errorf("invalid budget decision outcome %q", s.Outcome) } - if strings.TrimSpace(s.RedactionVersion) == "" { - return fmt.Errorf("redaction_version is required") - } - if s.CreatedAt.IsZero() { - return fmt.Errorf("created_at is required") - } - if err := validateAuditRedactedText("redacted_detail_ref", s.DetailRef()); err != nil { - return err - } return nil } diff --git a/platform/config_operation_summary.go b/platform/config_operation_summary.go index 489b99f82f..d45ab15662 100644 --- a/platform/config_operation_summary.go +++ b/platform/config_operation_summary.go @@ -106,6 +106,19 @@ func NewAppConfigOperationSummary(input AppConfigOperationSummaryInput) (AppConf // Validate checks that a config operation summary is safe to expose or store. func (s AppConfigOperationSummary) Validate() error { + if err := s.validateConfigOperationIdentity(); err != nil { + return err + } + if err := s.validateConfigOperationState(); err != nil { + return err + } + if err := s.validateConfigOperationLinks(); err != nil { + return err + } + return s.validateConfigOperationSafeText() +} + +func (s AppConfigOperationSummary) validateConfigOperationIdentity() error { if strings.TrimSpace(s.TenantID) == "" { return ErrTenantIDRequired } @@ -134,6 +147,10 @@ func (s AppConfigOperationSummary) Validate() error { if strings.TrimSpace(s.OperationID) == "" { return fmt.Errorf("operation_id is required") } + return nil +} + +func (s AppConfigOperationSummary) validateConfigOperationState() error { if strings.TrimSpace(s.PreviousVersion) == "" { return fmt.Errorf("previous_version is required") } @@ -156,6 +173,10 @@ func (s AppConfigOperationSummary) Validate() error { if s.CreatedAt.IsZero() { return fmt.Errorf("created_at is required") } + return nil +} + +func (s AppConfigOperationSummary) validateConfigOperationLinks() error { if err := s.CacheInvalidation.Validate(); err != nil { return fmt.Errorf("cache_invalidation: %w", err) } @@ -175,6 +196,10 @@ func (s AppConfigOperationSummary) Validate() error { s.GrayStatus.ActiveChecksum != s.NextChecksum { return fmt.Errorf("gray_status active version must match next active version") } + return nil +} + +func (s AppConfigOperationSummary) validateConfigOperationSafeText() error { for field, value := range map[string]string{ "summary_id": s.SummaryID, "operation_id": s.OperationID, @@ -218,41 +243,16 @@ func validateConfigOperationInvalidation(s AppConfigOperationSummary) error { } func validateConfigOperationGrayStatus(s AppConfigOperationSummary) error { - for field, value := range map[string]string{ - "gray_active_version": s.GrayStatus.ActiveVersion, - "gray_active_checksum": s.GrayStatus.ActiveChecksum, - "gray_candidate_version": s.GrayStatus.CandidateVersion, - "gray_candidate_checksum": s.GrayStatus.CandidateChecksum, - "gray_rollback_version": s.GrayStatus.RollbackVersion, - "gray_rollback_checksum": s.GrayStatus.RollbackChecksum, - } { - if err := validateAuditRedactedText(field, value); err != nil { - return err - } + if err := validateConfigOperationGrayStatusText(s.GrayStatus); err != nil { + return err } if s.GrayStatus.ActiveTrafficPercent < 0 || s.GrayStatus.ActiveTrafficPercent > 100 || s.GrayStatus.CandidateGrayPercent < 0 || s.GrayStatus.CandidateGrayPercent > 100 || s.GrayStatus.CandidateTrafficPercent < 0 || s.GrayStatus.CandidateTrafficPercent > 100 { return fmt.Errorf("gray_status traffic percentages must be between 0 and 100") } - if s.GrayStatus.HasCandidate { - if strings.TrimSpace(s.GrayStatus.CandidateVersion) == "" || - strings.TrimSpace(s.GrayStatus.CandidateChecksum) == "" { - return fmt.Errorf("gray_status candidate version and checksum are required") - } - if s.GrayStatus.CandidateGrayPercent != s.GrayStatus.CandidateTrafficPercent { - return fmt.Errorf("gray_status candidate traffic must match candidate gray percent") - } - if s.GrayStatus.ActiveTrafficPercent != 100-s.GrayStatus.CandidateTrafficPercent { - return fmt.Errorf("gray_status active traffic must complement candidate traffic") - } - } else if s.GrayStatus.CandidateVersion != "" || - s.GrayStatus.CandidateChecksum != "" || - s.GrayStatus.CandidateGrayPercent != 0 || - s.GrayStatus.CandidateTrafficPercent != 0 { - return fmt.Errorf("gray_status candidate fields require has_candidate") - } else if s.GrayStatus.ActiveTrafficPercent != 100 { - return fmt.Errorf("gray_status active traffic must be 100 when there is no candidate") + if err := validateConfigOperationGrayCandidate(s.GrayStatus); err != nil { + return err } if !s.GrayStatus.HasRollback { return fmt.Errorf("gray_status rollback version is required") @@ -264,6 +264,52 @@ func validateConfigOperationGrayStatus(s AppConfigOperationSummary) error { return nil } +func validateConfigOperationGrayStatusText(status ConfigGrayStatusSummary) error { + for field, value := range map[string]string{ + "gray_active_version": status.ActiveVersion, + "gray_active_checksum": status.ActiveChecksum, + "gray_candidate_version": status.CandidateVersion, + "gray_candidate_checksum": status.CandidateChecksum, + "gray_rollback_version": status.RollbackVersion, + "gray_rollback_checksum": status.RollbackChecksum, + } { + if err := validateAuditRedactedText(field, value); err != nil { + return err + } + } + return nil +} + +func validateConfigOperationGrayCandidate(status ConfigGrayStatusSummary) error { + if !status.HasCandidate { + return validateConfigOperationNoGrayCandidate(status) + } + if strings.TrimSpace(status.CandidateVersion) == "" || + strings.TrimSpace(status.CandidateChecksum) == "" { + return fmt.Errorf("gray_status candidate version and checksum are required") + } + if status.CandidateGrayPercent != status.CandidateTrafficPercent { + return fmt.Errorf("gray_status candidate traffic must match candidate gray percent") + } + if status.ActiveTrafficPercent != 100-status.CandidateTrafficPercent { + return fmt.Errorf("gray_status active traffic must complement candidate traffic") + } + return nil +} + +func validateConfigOperationNoGrayCandidate(status ConfigGrayStatusSummary) error { + if status.CandidateVersion != "" || + status.CandidateChecksum != "" || + status.CandidateGrayPercent != 0 || + status.CandidateTrafficPercent != 0 { + return fmt.Errorf("gray_status candidate fields require has_candidate") + } + if status.ActiveTrafficPercent != 100 { + return fmt.Errorf("gray_status active traffic must be 100 when there is no candidate") + } + return nil +} + func (i AppConfigOperationSummaryInput) normalize() (AppConfigOperationSummaryInput, error) { i.Operation = AppConfigOperation(strings.TrimSpace(string(i.Operation))) if !i.Operation.valid() { diff --git a/platform/gateway/service.go b/platform/gateway/service.go index b0592a14ff..b59d51fed0 100644 --- a/platform/gateway/service.go +++ b/platform/gateway/service.go @@ -130,36 +130,12 @@ func (s *Service) HandleInbound( routeCtx, routeSpan := telemetrytrace.Tracer.Start(ctx, "gateway.route") defer routeSpan.End() setInboundTraceAttributes(routeSpan, msg, "", requestID, "") - runtime, ok, err := s.registry.Lookup(routeCtx, msg) + runtime, err := s.lookupRuntime(routeCtx, ctx, routeSpan, msg, start) if err != nil { - recordSpanError(routeSpan, err) - return Result{}, err - } - if !ok { - err := ErrRuntimeNotFound - s.writeAudit(ctx, auditFromMessage(msg, "", "", "reject", err.Error(), start, err)) - recordSpanError(routeSpan, err) - return Result{}, err - } - if err := runtime.Validate(); err != nil { - s.writeAudit(ctx, auditFromMessage(msg, "", "", "reject", err.Error(), start, err)) - recordSpanError(routeSpan, err) - return Result{}, err - } - if !runtime.matchesInbound(msg) { - err := ErrRuntimeMismatch - s.writeAudit(ctx, auditFromMessage(msg, "", "", "reject", err.Error(), start, err)) - recordSpanError(routeSpan, err) - return Result{}, err - } - if err := authorizeBinding(runtime.Binding, msg); err != nil { - s.writeAudit(ctx, auditFromMessage(msg, "", "", "reject", err.Error(), start, err)) return Result{}, err } - text, err := inboundText(msg) + text, err := s.validateInboundContent(ctx, routeSpan, msg, start) if err != nil { - s.writeAudit(ctx, auditFromMessage(msg, "", "", "reject", err.Error(), start, err)) - recordSpanError(routeSpan, err) return Result{}, err } sessionID, err := platform.SessionIDForInbound(msg) @@ -176,45 +152,159 @@ func (s *Service) HandleInbound( msg.ChannelAccountID, msg.PlatformMessageID, ) + record, handled, result, err := s.startInboundRun( + routeCtx, + ctx, + msg, + sessionID, + requestID, + internalUserID, + key, + ) + if err != nil { + return Result{}, err + } + if handled { + return result, nil + } + defer s.releaseSessionLease(ctx, record.SessionLease) + + return s.runAndReply( + routeCtx, + ctx, + runtime, + msg, + inboundRunInput{ + Text: text, + SessionID: sessionID, + InternalUserID: internalUserID, + RequestID: requestID, + Key: key, + Start: start, + }, + ) +} + +type inboundRunRecord struct { + Record platform.IdempotencyRecord + SessionLease SessionLease +} + +type inboundRunInput struct { + Text string + SessionID string + InternalUserID string + RequestID string + Key string + Start time.Time +} + +func (s *Service) lookupRuntime( + routeCtx context.Context, + auditCtx context.Context, + routeSpan oteltrace.Span, + msg platform.InboundMessage, + start time.Time, +) (Runtime, error) { + runtime, ok, err := s.registry.Lookup(routeCtx, msg) + if err != nil { + recordSpanError(routeSpan, err) + return Runtime{}, err + } + if !ok { + err := ErrRuntimeNotFound + s.writeRejectAudit(auditCtx, msg, start, err) + recordSpanError(routeSpan, err) + return Runtime{}, err + } + if err := validateRuntimeForMessage(runtime, msg); err != nil { + s.writeRejectAudit(auditCtx, msg, start, err) + recordSpanError(routeSpan, err) + return Runtime{}, err + } + return runtime, nil +} + +func validateRuntimeForMessage(runtime Runtime, msg platform.InboundMessage) error { + if err := runtime.Validate(); err != nil { + return err + } + if !runtime.matchesInbound(msg) { + return ErrRuntimeMismatch + } + return authorizeBinding(runtime.Binding, msg) +} + +func (s *Service) validateInboundContent( + ctx context.Context, + routeSpan oteltrace.Span, + msg platform.InboundMessage, + start time.Time, +) (string, error) { + text, err := inboundText(msg) + if err != nil { + s.writeRejectAudit(ctx, msg, start, err) + recordSpanError(routeSpan, err) + return "", err + } + return text, nil +} + +func (s *Service) startInboundRun( + routeCtx context.Context, + resultCtx context.Context, + msg platform.InboundMessage, + sessionID string, + requestID string, + internalUserID string, + key string, +) (inboundRunRecord, bool, Result, error) { idempotencyCtx, idempotencySpan := telemetrytrace.Tracer.Start(routeCtx, "gateway.idempotency") + defer idempotencySpan.End() setInboundTraceAttributes(idempotencySpan, msg, sessionID, requestID, internalUserID) existing, ok, err := s.idempotencyStore.Get(idempotencyCtx, key) if err != nil { recordSpanError(idempotencySpan, err) - idempotencySpan.End() - return Result{}, err + return inboundRunRecord{}, false, Result{}, err } if ok { - idempotencySpan.End() - return s.duplicateResult(ctx, existing) - } - leaseCtx, leaseSpan := telemetrytrace.Tracer.Start(routeCtx, "gateway.session_lock") - setInboundTraceAttributes(leaseSpan, msg, sessionID, requestID, internalUserID) - lease, acquired, err := s.leaseStore.Acquire(leaseCtx, SessionLeaseKey{ - TenantID: msg.TenantID, - AppID: msg.AppID, - SessionID: sessionID, - }) - if err != nil { - recordSpanError(leaseSpan, err) - leaseSpan.End() - return Result{}, err + result, err := s.duplicateResult(resultCtx, existing) + return inboundRunRecord{}, true, result, err + } + return s.acquireSessionLeaseAndStart( + routeCtx, + resultCtx, + idempotencyCtx, + idempotencySpan, + msg, + sessionID, + requestID, + internalUserID, + key, + ) +} + +func (s *Service) acquireSessionLeaseAndStart( + routeCtx context.Context, + resultCtx context.Context, + idempotencyCtx context.Context, + idempotencySpan oteltrace.Span, + msg platform.InboundMessage, + sessionID string, + requestID string, + internalUserID string, + key string, +) (inboundRunRecord, bool, Result, error) { + lease, handled, result, err := s.acquireSessionLease( + routeCtx, + msg, + sessionID, + requestID, + internalUserID, + ) + if err != nil || handled { + return inboundRunRecord{}, handled, result, err } - if !acquired { - leaseSpan.End() - return Result{ - RequestID: requestID, - SessionID: sessionID, - Status: platform.IdempotencyStatusProcessing, - Processing: true, - }, nil - } - defer func() { - cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) - defer cancel() - _ = lease.Release(cleanupCtx) - }() - leaseSpan.End() record, started, err := s.idempotencyStore.Start(idempotencyCtx, platform.IdempotencyRecord{ TenantID: msg.TenantID, Channel: msg.Channel, @@ -225,60 +315,128 @@ func (s *Service) HandleInbound( SessionID: sessionID, }) if err != nil { + s.releaseSessionLease(resultCtx, lease) recordSpanError(idempotencySpan, err) - idempotencySpan.End() - return Result{}, err + return inboundRunRecord{}, false, Result{}, err } if !started { - idempotencySpan.End() - return s.duplicateResult(ctx, record) + s.releaseSessionLease(resultCtx, lease) + result, err := s.duplicateResult(resultCtx, record) + return inboundRunRecord{}, true, result, err + } + return inboundRunRecord{Record: record, SessionLease: lease}, false, Result{}, nil +} + +func (s *Service) acquireSessionLease( + routeCtx context.Context, + msg platform.InboundMessage, + sessionID string, + requestID string, + internalUserID string, +) (SessionLease, bool, Result, error) { + leaseCtx, leaseSpan := telemetrytrace.Tracer.Start(routeCtx, "gateway.session_lock") + defer leaseSpan.End() + setInboundTraceAttributes(leaseSpan, msg, sessionID, requestID, internalUserID) + lease, acquired, err := s.leaseStore.Acquire(leaseCtx, SessionLeaseKey{ + TenantID: msg.TenantID, + AppID: msg.AppID, + SessionID: sessionID, + }) + if err != nil { + recordSpanError(leaseSpan, err) + return nil, false, Result{}, err } - idempotencySpan.End() + if acquired { + return lease, false, Result{}, nil + } + return nil, true, Result{ + RequestID: requestID, + SessionID: sessionID, + Status: platform.IdempotencyStatusProcessing, + Processing: true, + }, nil +} + +func (s *Service) releaseSessionLease(ctx context.Context, lease SessionLease) { + cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + defer cancel() + _ = lease.Release(cleanupCtx) +} +func (s *Service) runAndReply( + routeCtx context.Context, + auditCtx context.Context, + runtime Runtime, + msg platform.InboundMessage, + input inboundRunInput, +) (Result, error) { + content, err := s.runGatewayRunner(routeCtx, auditCtx, runtime, msg, input) + if err != nil { + return Result{}, err + } + return s.writeReply(routeCtx, auditCtx, runtime, msg, input, content) +} + +func (s *Service) runGatewayRunner( + routeCtx context.Context, + auditCtx context.Context, + runtime Runtime, + msg platform.InboundMessage, + input inboundRunInput, +) (string, error) { runnerCtx, runnerSpan := telemetrytrace.Tracer.Start(routeCtx, "runner.run") - setInboundTraceAttributes(runnerSpan, msg, sessionID, requestID, internalUserID) + defer runnerSpan.End() + setInboundTraceAttributes(runnerSpan, msg, input.SessionID, input.RequestID, input.InternalUserID) ch, err := runtime.Runner.Run( runnerCtx, - internalUserID, - sessionID, - model.NewUserMessage(text), - agent.WithRequestID(requestID), + input.InternalUserID, + input.SessionID, + model.NewUserMessage(input.Text), + agent.WithRequestID(input.RequestID), agent.WithLatencyDiagnostics(true), agent.WithLatencyDiagnosticsEvents(false), ) if err != nil { - s.writeAudit(ctx, auditFromMessage(msg, sessionID, internalUserID, "runner_error", err.Error(), start, err)) + s.writeAudit(auditCtx, auditFromMessage(msg, input.SessionID, input.InternalUserID, "runner_error", err.Error(), input.Start, err)) recordSpanError(runnerSpan, err) - runnerSpan.End() - return Result{}, err + return "", err } - content, err := collectAssistantText(ctx, ch) + content, err := collectAssistantText(auditCtx, ch) if err != nil { - s.writeAudit(ctx, auditFromMessage(msg, sessionID, internalUserID, "runner_error", err.Error(), start, err)) + s.writeAudit(auditCtx, auditFromMessage(msg, input.SessionID, input.InternalUserID, "runner_error", err.Error(), input.Start, err)) recordSpanError(runnerSpan, err) - runnerSpan.End() - return Result{}, err + return "", err } - runnerSpan.End() - resultRef := key + ":outbound:1" + return content, nil +} + +func (s *Service) writeReply( + routeCtx context.Context, + auditCtx context.Context, + runtime Runtime, + msg platform.InboundMessage, + input inboundRunInput, + content string, +) (Result, error) { + resultRef := input.Key + ":outbound:1" outbound := platform.OutboundMessage{ TenantID: msg.TenantID, BindingID: msg.BindingID, Channel: msg.Channel, - SessionID: sessionID, + SessionID: input.SessionID, ReplyToPlatformMessageID: msg.PlatformMessageID, Kind: platform.OutboundMessageKindText, Content: content, Sequence: 1, DedupKey: resultRef, - TraceID: requestID, + TraceID: input.RequestID, } replyCtx, replySpan := telemetrytrace.Tracer.Start(routeCtx, "im.reply") - setInboundTraceAttributes(replySpan, msg, sessionID, requestID, internalUserID) + defer replySpan.End() + setInboundTraceAttributes(replySpan, msg, input.SessionID, input.RequestID, input.InternalUserID) if err := s.outboundStore.Save(replyCtx, resultRef, outbound); err != nil { - s.writeAudit(ctx, auditFromMessage(msg, sessionID, internalUserID, "outbound_error", err.Error(), start, err)) + s.writeAudit(auditCtx, auditFromMessage(msg, input.SessionID, input.InternalUserID, "outbound_error", err.Error(), input.Start, err)) recordSpanError(replySpan, err) - replySpan.End() return Result{}, err } if err := s.outboundStore.Enqueue( @@ -286,29 +444,25 @@ func (s *Service) HandleInbound( outbound, channeladapter.RetryPolicyForBinding(runtime.Binding), ); err != nil { - if _, markErr := s.idempotencyStore.MarkReplyFailed(replyCtx, key, resultRef); markErr != nil { + if _, markErr := s.idempotencyStore.MarkReplyFailed(replyCtx, input.Key, resultRef); markErr != nil { recordSpanError(replySpan, markErr) - replySpan.End() return Result{}, markErr } - s.writeAudit(ctx, auditFromMessage(msg, sessionID, internalUserID, "outbound_error", err.Error(), start, err)) + s.writeAudit(auditCtx, auditFromMessage(msg, input.SessionID, input.InternalUserID, "outbound_error", err.Error(), input.Start, err)) recordSpanError(replySpan, err) - replySpan.End() return Result{}, err } - record, err = s.idempotencyStore.Complete(replyCtx, key, resultRef) + record, err := s.idempotencyStore.Complete(replyCtx, input.Key, resultRef) if err != nil { recordSpanError(replySpan, err) - replySpan.End() return Result{}, err } - replySpan.End() - s.writeMessageEvent(ctx, messageEventFromInbound(msg, sessionID, key, requestID, 1, start)) - s.writeMessageEvent(ctx, messageEventFromAssistant(msg, sessionID, resultRef, requestID, 2, s.now())) - s.writeAudit(ctx, auditFromMessage(msg, sessionID, internalUserID, "completed", "", start, nil)) + s.writeMessageEvent(auditCtx, messageEventFromInbound(msg, input.SessionID, input.Key, input.RequestID, 1, input.Start)) + s.writeMessageEvent(auditCtx, messageEventFromAssistant(msg, input.SessionID, resultRef, input.RequestID, 2, s.now())) + s.writeAudit(auditCtx, auditFromMessage(msg, input.SessionID, input.InternalUserID, "completed", "", input.Start, nil)) return Result{ - RequestID: requestID, - SessionID: sessionID, + RequestID: input.RequestID, + SessionID: input.SessionID, ResultRef: resultRef, Status: record.Status, Outbound: outbound, @@ -316,6 +470,15 @@ func (s *Service) HandleInbound( }, nil } +func (s *Service) writeRejectAudit( + ctx context.Context, + msg platform.InboundMessage, + start time.Time, + err error, +) { + s.writeAudit(ctx, auditFromMessage(msg, "", "", "reject", err.Error(), start, err)) +} + func (s *Service) validateService() error { if s.registry == nil { return fmt.Errorf("gateway registry is required") @@ -619,50 +782,39 @@ func recordSpanError(span oteltrace.Span, err error) { } func traceErrorType(err error) string { - switch { - case err == nil: + if err == nil { return "" - case errors.Is(err, context.Canceled): - return "context_canceled" - case errors.Is(err, context.DeadlineExceeded): - return "context_deadline_exceeded" - case errors.Is(err, platform.ErrTenantIDRequired): - return "tenant_id_required" - case errors.Is(err, platform.ErrAppIDRequired): - return "app_id_required" - case errors.Is(err, platform.ErrBindingIDRequired): - return "binding_id_required" - case errors.Is(err, platform.ErrChannelRequired): - return "channel_required" - case errors.Is(err, platform.ErrAccountIDRequired): - return "account_id_required" - case errors.Is(err, platform.ErrPlatformMessageIDRequired): - return "platform_message_id_required" - case errors.Is(err, platform.ErrExternalUserIDRequired): - return "external_user_id_required" - case errors.Is(err, platform.ErrExternalGroupIDRequired): - return "external_group_id_required" - case errors.Is(err, platform.ErrConversationTypeRequired): - return "conversation_type_required" - case errors.Is(err, platform.ErrInvalidConversationType): - return "invalid_conversation_type" - case errors.Is(err, ErrRuntimeNotFound): - return "runtime_not_found" - case errors.Is(err, ErrRuntimeInactive): - return "runtime_inactive" - case errors.Is(err, ErrRuntimeMismatch): - return "runtime_mismatch" - case errors.Is(err, ErrBindingAccessDenied): - return "binding_access_denied" - case errors.Is(err, ErrBindingMentionRequired): - return "binding_mention_required" - case errors.Is(err, ErrUnsupportedMessageType): - return "unsupported_message_type" - case errors.Is(err, ErrEmptyText): - return "empty_text" - case errors.Is(err, ErrRunnerResponseEmpty): - return "runner_response_empty" - default: - return "gateway_error" } + for _, candidate := range traceErrorTypes { + if errors.Is(err, candidate.err) { + return candidate.name + } + } + return "gateway_error" +} + +var traceErrorTypes = []struct { + err error + name string +}{ + {context.Canceled, "context_canceled"}, + {context.DeadlineExceeded, "context_deadline_exceeded"}, + {platform.ErrTenantIDRequired, "tenant_id_required"}, + {platform.ErrAppIDRequired, "app_id_required"}, + {platform.ErrBindingIDRequired, "binding_id_required"}, + {platform.ErrChannelRequired, "channel_required"}, + {platform.ErrAccountIDRequired, "account_id_required"}, + {platform.ErrPlatformMessageIDRequired, "platform_message_id_required"}, + {platform.ErrExternalUserIDRequired, "external_user_id_required"}, + {platform.ErrExternalGroupIDRequired, "external_group_id_required"}, + {platform.ErrConversationTypeRequired, "conversation_type_required"}, + {platform.ErrInvalidConversationType, "invalid_conversation_type"}, + {ErrRuntimeNotFound, "runtime_not_found"}, + {ErrRuntimeInactive, "runtime_inactive"}, + {ErrRuntimeMismatch, "runtime_mismatch"}, + {ErrBindingAccessDenied, "binding_access_denied"}, + {ErrBindingMentionRequired, "binding_mention_required"}, + {ErrUnsupportedMessageType, "unsupported_message_type"}, + {ErrEmptyText, "empty_text"}, + {ErrRunnerResponseEmpty, "runner_response_empty"}, } diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go index 30382f0da4..ead64ed97d 100644 --- a/platform/gateway/service_test.go +++ b/platform/gateway/service_test.go @@ -807,7 +807,7 @@ func TestServiceHandleInboundRunnerErrorRedactsAuditReason(t *testing.T) { func TestServiceHandleInboundUsesRequestIDAndStreamsText(t *testing.T) { ctx := context.Background() registry := NewInMemoryRegistry() - r := &recordingRunner{chunks: []string{"hel", "lo"}} + r := &recordingRunner{chunks: []string{"he", "llo"}} registerRuntime(t, registry, "tenant-a", r) audit := platform.NewInMemoryAuditSink() messageEvents := platform.NewInMemoryMessageEventSink() @@ -1025,8 +1025,8 @@ func TestCollectAssistantTextStopsAtRunnerCompletion(t *testing.T) { func TestCollectAssistantTextPrefersFinalFullMessage(t *testing.T) { ch := make(chan *event.Event, 3) - ch <- chunkEvent("hel", true) - ch <- chunkEvent("lo", true) + ch <- chunkEvent("he", true) + ch <- chunkEvent("llo", true) ch <- responseEvent("hello", true) close(ch) diff --git a/platform/toolpolicy/policy.go b/platform/toolpolicy/policy.go index 11d4ab58a6..7be766973f 100644 --- a/platform/toolpolicy/policy.go +++ b/platform/toolpolicy/policy.go @@ -500,6 +500,25 @@ func (p *Policy) ApprovalSummary( // Validate checks that the summary is safe to expose outside the tool runtime. func (s ApprovalSummary) Validate() error { + if err := s.validateIdentity(); err != nil { + return err + } + if err := s.validateArguments(); err != nil { + return err + } + if err := s.validateDecision(); err != nil { + return err + } + if strings.TrimSpace(s.RedactionVersion) == "" { + return fmt.Errorf("redaction_version is required") + } + if s.CreatedAt.IsZero() { + return fmt.Errorf("created_at is required") + } + return platformSafeText("detail_ref", s.DetailRef()) +} + +func (s ApprovalSummary) validateIdentity() error { if strings.TrimSpace(s.TenantID) == "" { return fmt.Errorf("tenant_id is required") } @@ -530,6 +549,10 @@ func (s ApprovalSummary) Validate() error { if err := platformSafeText("reason", s.Reason); err != nil { return err } + return nil +} + +func (s ApprovalSummary) validateArguments() error { if s.ArgumentsBytes < 0 { return fmt.Errorf("arguments_bytes must be greater than or equal to 0") } @@ -543,6 +566,10 @@ func (s ApprovalSummary) Validate() error { if s.MaxResultSize < 0 { return fmt.Errorf("max_result_size must be greater than or equal to 0") } + return nil +} + +func (s ApprovalSummary) validateDecision() error { switch s.Decision { case tool.PermissionActionAllow: if s.RequiresApproval { @@ -561,15 +588,6 @@ func (s ApprovalSummary) Validate() error { default: return fmt.Errorf("invalid decision %q", s.Decision) } - if strings.TrimSpace(s.RedactionVersion) == "" { - return fmt.Errorf("redaction_version is required") - } - if s.CreatedAt.IsZero() { - return fmt.Errorf("created_at is required") - } - if err := platformSafeText("detail_ref", s.DetailRef()); err != nil { - return err - } return nil } diff --git a/session/postgres/go.mod b/session/postgres/go.mod index fdd3f68f81..fd328aeb3c 100644 --- a/session/postgres/go.mod +++ b/session/postgres/go.mod @@ -11,6 +11,8 @@ require ( github.com/DATA-DOG/go-sqlmock v1.5.2 github.com/google/uuid v1.6.0 github.com/stretchr/testify v1.10.0 + go.opentelemetry.io/otel v1.29.0 + go.opentelemetry.io/otel/sdk v1.29.0 trpc.group/trpc-go/trpc-agent-go v0.2.0 trpc.group/trpc-go/trpc-agent-go/storage/postgres v0.8.0 ) @@ -26,12 +28,10 @@ require ( github.com/jackc/pgx/v5 v5.7.2 // indirect github.com/jackc/puddle/v2 v2.2.2 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect - go.opentelemetry.io/otel v1.29.0 // indirect go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.29.0 // indirect go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.29.0 // indirect go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.29.0 // indirect go.opentelemetry.io/otel/metric v1.29.0 // indirect - go.opentelemetry.io/otel/sdk v1.29.0 // indirect go.opentelemetry.io/otel/trace v1.29.0 // indirect go.opentelemetry.io/proto/otlp v1.3.1 // indirect go.uber.org/multierr v1.10.0 // indirect From 3daa91354ec7df9f3eb98f846d37ce7458e68a8b Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Thu, 9 Jul 2026 17:45:17 +0800 Subject: [PATCH 51/95] fix(platform): address CodeRabbit review feedback --- internal/flow/processor/content.go | 27 +++--- internal/telemetry/trace.go | 30 ++++++ platform/audit.go | 21 +---- platform/audit_query.go | 30 +++--- platform/config_cache_invalidation.go | 3 + platform/config_cache_invalidation_test.go | 11 ++- platform/config_diff.go | 9 ++ platform/config_diff_test.go | 8 ++ platform/config_operation_summary.go | 60 +++++------- platform/gateway/registry.go | 4 +- platform/gateway/registry_test.go | 55 +++++++++++ platform/gateway/service.go | 46 ++++++---- platform/gateway/service_test.go | 13 +++ platform/identity.go | 6 +- platform/inmemory_sink.go | 42 +++++++++ platform/message_event.go | 21 +---- platform/message_event_test.go | 21 +++++ platform/migration_test.go | 6 ++ platform/operational_action_audit.go | 26 +++--- platform/operational_action_audit_test.go | 20 ++++ platform/secret_rotation_status.go | 60 +++++++----- platform/secret_rotation_status_test.go | 26 ++++++ platform/storagerouter/errors.go | 2 + platform/storagerouter/router.go | 6 +- platform/storagerouter/router_test.go | 10 ++ platform/storagerouter/status.go | 17 +++- platform/toolpolicy/policy.go | 102 +++++++++------------ platform/types_test.go | 28 ++++++ platform/usage.go | 21 +---- platform/usage_record_test.go | 21 +++++ platform/validation.go | 48 ++++++---- runner/diagnostics.go | 11 +-- session/redis/summary.go | 6 +- 33 files changed, 545 insertions(+), 272 deletions(-) create mode 100644 platform/gateway/registry_test.go create mode 100644 platform/inmemory_sink.go diff --git a/internal/flow/processor/content.go b/internal/flow/processor/content.go index b1d7a037ed..86a31372ce 100644 --- a/internal/flow/processor/content.go +++ b/internal/flow/processor/content.go @@ -3110,23 +3110,26 @@ func (p *ContentRequestProcessor) getAdaptivePreloadMemoryMessage( HybridSearch: true, } searchCtx, span, startedSpan := itrace.StartSpan(ctx, inv, itelemetry.NewMemorySearchSpanName()) - memories, err := reader.SearchMemories( + var memories []*memory.Entry + if startedSpan { + defer func() { + itelemetry.TraceMemorySearch( + span, + searchOpts.MaxResults, + len(memories), + searchOpts.HybridSearch, + searchOpts.Deduplicate, + err, + ) + span.End() + }() + } + memories, err = reader.SearchMemories( searchCtx, userKey, query, memory.WithSearchOptions(searchOpts), ) - if startedSpan { - itelemetry.TraceMemorySearch( - span, - searchOpts.MaxResults, - len(memories), - searchOpts.HybridSearch, - searchOpts.Deduplicate, - err, - ) - span.End() - } if err != nil { log.WarnfContext(ctx, "Failed to search memories for preload: %v", err) return p.loadPreloadMemoryMessage(ctx, inv, reader, userKey, budget) diff --git a/internal/telemetry/trace.go b/internal/telemetry/trace.go index 88ffae2f23..c91c8e0be7 100644 --- a/internal/telemetry/trace.go +++ b/internal/telemetry/trace.go @@ -12,9 +12,14 @@ package telemetry import ( + "crypto/hmac" + "crypto/sha256" + "encoding/hex" "encoding/json" "errors" "fmt" + "os" + "strings" "time" "go.opentelemetry.io/otel/attribute" @@ -32,6 +37,10 @@ import ( "trpc.group/trpc-go/trpc-agent-go/tool" ) +const traceSafeHashEnvKey = "TRPC_AGENT_TRACE_HASH_KEY" + +var traceSafeHashDefaultKey = []byte("trpc-agent-go-trace-safe-hash-v1") + // grpcDial is a package-level variable to allow test injection of a custom dialer. // In production, this points to grpc.Dial. var grpcDial = grpc.Dial @@ -153,6 +162,27 @@ func recordSafeSpanError(span trace.Span, err error, fallback string) { span.RecordError(errors.New(errorType)) } +// TraceSafeHash returns a stable low-cardinality HMAC digest for trace attributes. +func TraceSafeHash(scope string, value string) string { + value = strings.TrimSpace(value) + if value == "" { + return "" + } + mac := hmac.New(sha256.New, traceSafeHashKey()) + _, _ = mac.Write([]byte(scope)) + _, _ = mac.Write([]byte{0}) + _, _ = mac.Write([]byte(value)) + return scope + "_hash_" + hex.EncodeToString(mac.Sum(nil))[:24] +} + +func traceSafeHashKey() []byte { + key := strings.TrimSpace(os.Getenv(traceSafeHashEnvKey)) + if key == "" { + return traceSafeHashDefaultKey + } + return []byte(key) +} + // WorkflowType is the normalized type vocabulary used by workflow spans. type WorkflowType string diff --git a/platform/audit.go b/platform/audit.go index 921cd6f5b9..b4a14c5c13 100644 --- a/platform/audit.go +++ b/platform/audit.go @@ -10,7 +10,6 @@ package platform import ( "context" - "sync" ) // AuditSink stores audit records. @@ -21,8 +20,7 @@ type AuditSink interface { // InMemoryAuditSink is a concurrency-safe audit sink for tests and demos. type InMemoryAuditSink struct { - mu sync.Mutex - records []AuditRecord + records inMemoryRecords[AuditRecord] } // NewInMemoryAuditSink creates an in-memory audit sink. @@ -32,23 +30,10 @@ func NewInMemoryAuditSink() *InMemoryAuditSink { // WriteAudit writes one audit record. func (s *InMemoryAuditSink) WriteAudit(ctx context.Context, record AuditRecord) error { - if err := ctx.Err(); err != nil { - return err - } - if err := record.Validate(); err != nil { - return err - } - s.mu.Lock() - defer s.mu.Unlock() - s.records = append(s.records, record) - return nil + return s.records.append(ctx, record, AuditRecord.Validate) } // Records returns a snapshot of written audit records. func (s *InMemoryAuditSink) Records() []AuditRecord { - s.mu.Lock() - defer s.mu.Unlock() - out := make([]AuditRecord, len(s.records)) - copy(out, s.records) - return out + return s.records.snapshot() } diff --git a/platform/audit_query.go b/platform/audit_query.go index 47e1f8dfb7..0d91cb450b 100644 --- a/platform/audit_query.go +++ b/platform/audit_query.go @@ -67,22 +67,20 @@ func (f AuditQueryFilter) normalize() (AuditQueryFilter, error) { if f.TenantID == "" { return AuditQueryFilter{}, ErrTenantIDRequired } - for field, value := range map[string]string{ - "app_id": f.AppID, - "audit_id": f.AuditID, - "channel": f.Channel, - "binding_id": f.BindingID, - "user_id_hash": f.UserIDHash, - "session_id": f.SessionID, - "request_id": f.RequestID, - "message_id": f.MessageID, - "tool_name": f.ToolName, - "decision": f.Decision, - "trace_id": f.TraceID, - } { - if err := validateAuditRedactedText(field, value); err != nil { - return AuditQueryFilter{}, err - } + if err := validateAuditRedactedFields( + safeTextField{"app_id", f.AppID}, + safeTextField{"audit_id", f.AuditID}, + safeTextField{"channel", f.Channel}, + safeTextField{"binding_id", f.BindingID}, + safeTextField{"user_id_hash", f.UserIDHash}, + safeTextField{"session_id", f.SessionID}, + safeTextField{"request_id", f.RequestID}, + safeTextField{"message_id", f.MessageID}, + safeTextField{"tool_name", f.ToolName}, + safeTextField{"decision", f.Decision}, + safeTextField{"trace_id", f.TraceID}, + ); err != nil { + return AuditQueryFilter{}, err } if f.Limit < 0 { return AuditQueryFilter{}, fmt.Errorf("limit must be non-negative") diff --git a/platform/config_cache_invalidation.go b/platform/config_cache_invalidation.go index 2128f7fe95..06c4e81aeb 100644 --- a/platform/config_cache_invalidation.go +++ b/platform/config_cache_invalidation.go @@ -135,6 +135,9 @@ func (i AppConfigCacheInvalidationInput) normalize() (AppConfigCacheInvalidation if err := requireSameConfigOwner(i.PreviousVersion, i.NextVersion); err != nil { return AppConfigCacheInvalidationInput{}, err } + if i.PreviousVersion.Status != AppConfigVersionStatusRollback { + return AppConfigCacheInvalidationInput{}, fmt.Errorf("previous config version status must be rollback") + } if i.NextVersion.Status != AppConfigVersionStatusActive { return AppConfigCacheInvalidationInput{}, fmt.Errorf("next config version status must be active") } diff --git a/platform/config_cache_invalidation_test.go b/platform/config_cache_invalidation_test.go index a85319bd1e..c63475573c 100644 --- a/platform/config_cache_invalidation_test.go +++ b/platform/config_cache_invalidation_test.go @@ -89,13 +89,9 @@ func TestNewAppConfigCacheInvalidationBuildsRollbackMarker(t *testing.T) { } whitespace := previous - whitespace.TenantID = " tenant " - whitespace.AppID = " app " whitespace.Version = " v2 " whitespace.Checksum = " sha256:previous " whitespaceNext := next - whitespaceNext.TenantID = " tenant " - whitespaceNext.AppID = " app " whitespaceNext.Version = " v1 " whitespaceNext.Checksum = " sha256:next " trimmed, err := NewAppConfigCacheInvalidation(AppConfigCacheInvalidationInput{ @@ -160,6 +156,13 @@ func TestNewAppConfigCacheInvalidationRejectsInvalidInputs(t *testing.T) { t.Fatalf("expected next active status error, got %v", err) } + previousNotRollback := base + previousNotRollback.PreviousVersion.Status = AppConfigVersionStatusReleased + if _, err := NewAppConfigCacheInvalidation(previousNotRollback); err == nil || + !strings.Contains(err.Error(), "rollback") { + t.Fatalf("expected previous rollback status error, got %v", err) + } + sameVersion := base sameVersion.NextVersion.Version = sameVersion.PreviousVersion.Version if _, err := NewAppConfigCacheInvalidation(sameVersion); err == nil || diff --git a/platform/config_diff.go b/platform/config_diff.go index 2e8d7268d0..2b8b37744c 100644 --- a/platform/config_diff.go +++ b/platform/config_diff.go @@ -11,7 +11,9 @@ package platform import ( "bytes" "encoding/json" + "errors" "fmt" + "io" "reflect" "sort" "strings" @@ -90,6 +92,13 @@ func decodeConfigBundle(bundle string) (any, error) { if err := decoder.Decode(&value); err != nil { return nil, err } + var trailing any + if err := decoder.Decode(&trailing); !errors.Is(err, io.EOF) { + if err == nil { + return nil, fmt.Errorf("config bundle contains trailing json data") + } + return nil, err + } return value, nil } diff --git a/platform/config_diff_test.go b/platform/config_diff_test.go index 15fbc7c8df..3a7de30a67 100644 --- a/platform/config_diff_test.go +++ b/platform/config_diff_test.go @@ -96,6 +96,14 @@ func TestDiffAppConfigVersionsRejectsInvalidInputs(t *testing.T) { } } +func TestDecodeConfigBundleRejectsTrailingJSON(t *testing.T) { + _, err := decodeConfigBundle(`{"model_profile_id":"model"} {"tool_policy_id":"tools"}`) + + if err == nil || !strings.Contains(err.Error(), "trailing json") { + t.Fatalf("expected trailing json rejection, got %v", err) + } +} + func TestDiffAppConfigVersionsReportsArrayAddRemove(t *testing.T) { from := validAppConfigVersion() from.ConfigBundleJSON = `{"tools":["search","ticket"]}` diff --git a/platform/config_operation_summary.go b/platform/config_operation_summary.go index d45ab15662..89f9e8a4a6 100644 --- a/platform/config_operation_summary.go +++ b/platform/config_operation_summary.go @@ -28,8 +28,10 @@ const ( // AppConfigOperationSummaryInput describes one planned or completed config operation. type AppConfigOperationSummaryInput struct { - Operation AppConfigOperation + Operation AppConfigOperation + // PreviousActive must already carry AppConfigVersionStatusRollback. PreviousActive AppConfigVersion + // NextActive must carry AppConfigVersionStatusActive. NextActive AppConfigVersion ResultVersions []AppConfigVersion OperationID string @@ -200,20 +202,15 @@ func (s AppConfigOperationSummary) validateConfigOperationLinks() error { } func (s AppConfigOperationSummary) validateConfigOperationSafeText() error { - for field, value := range map[string]string{ - "summary_id": s.SummaryID, - "operation_id": s.OperationID, - "previous_version": s.PreviousVersion, - "previous_checksum": s.PreviousChecksum, - "next_version": s.NextVersion, - "next_checksum": s.NextChecksum, - "trace_id": s.TraceID, - } { - if err := validateAuditRedactedText(field, value); err != nil { - return err - } - } - return nil + return validateAuditRedactedFields( + safeTextField{"summary_id", s.SummaryID}, + safeTextField{"operation_id", s.OperationID}, + safeTextField{"previous_version", s.PreviousVersion}, + safeTextField{"previous_checksum", s.PreviousChecksum}, + safeTextField{"next_version", s.NextVersion}, + safeTextField{"next_checksum", s.NextChecksum}, + safeTextField{"trace_id", s.TraceID}, + ) } func validateConfigOperationInvalidation(s AppConfigOperationSummary) error { @@ -265,19 +262,14 @@ func validateConfigOperationGrayStatus(s AppConfigOperationSummary) error { } func validateConfigOperationGrayStatusText(status ConfigGrayStatusSummary) error { - for field, value := range map[string]string{ - "gray_active_version": status.ActiveVersion, - "gray_active_checksum": status.ActiveChecksum, - "gray_candidate_version": status.CandidateVersion, - "gray_candidate_checksum": status.CandidateChecksum, - "gray_rollback_version": status.RollbackVersion, - "gray_rollback_checksum": status.RollbackChecksum, - } { - if err := validateAuditRedactedText(field, value); err != nil { - return err - } - } - return nil + return validateAuditRedactedFields( + safeTextField{"gray_active_version", status.ActiveVersion}, + safeTextField{"gray_active_checksum", status.ActiveChecksum}, + safeTextField{"gray_candidate_version", status.CandidateVersion}, + safeTextField{"gray_candidate_checksum", status.CandidateChecksum}, + safeTextField{"gray_rollback_version", status.RollbackVersion}, + safeTextField{"gray_rollback_checksum", status.RollbackChecksum}, + ) } func validateConfigOperationGrayCandidate(status ConfigGrayStatusSummary) error { @@ -349,13 +341,11 @@ func (i AppConfigOperationSummaryInput) normalize() (AppConfigOperationSummaryIn return AppConfigOperationSummaryInput{}, err } } - for field, value := range map[string]string{ - "operation_id": i.OperationID, - "trace_id": i.TraceID, - } { - if err := validateAuditRedactedText(field, value); err != nil { - return AppConfigOperationSummaryInput{}, err - } + if err := validateAuditRedactedFields( + safeTextField{"operation_id", i.OperationID}, + safeTextField{"trace_id", i.TraceID}, + ); err != nil { + return AppConfigOperationSummaryInput{}, err } return i, nil } diff --git a/platform/gateway/registry.go b/platform/gateway/registry.go index 0c15b09060..c3eab62739 100644 --- a/platform/gateway/registry.go +++ b/platform/gateway/registry.go @@ -125,9 +125,9 @@ func (r *InMemoryRegistry) Lookup( func runtimeKey(tenantID, appID, bindingID, channel, accountID string) string { return platform.IdempotencyKey( - tenantID+"|"+appID, + tenantID, channel, accountID, - bindingID, + platform.IdempotencyKey(appID, channel, accountID, bindingID), ) } diff --git a/platform/gateway/registry_test.go b/platform/gateway/registry_test.go new file mode 100644 index 0000000000..078fe00f58 --- /dev/null +++ b/platform/gateway/registry_test.go @@ -0,0 +1,55 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package gateway + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "trpc.group/trpc-go/trpc-agent-go/platform" +) + +func TestInMemoryRegistryAvoidsTenantAppDelimiterCollision(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + firstRunner := &recordingRunner{response: "first"} + secondRunner := &recordingRunner{response: "second"} + first := validRuntimeForBinding("tenant|app", "alpha", "binding", "wecom", "acct", firstRunner) + second := validRuntimeForBinding("tenant", "app|alpha", "binding", "wecom", "acct", secondRunner) + require.NoError(t, registry.Register(first)) + require.NoError(t, registry.Register(second)) + + gotFirst, ok, err := registry.Lookup(ctx, inboundForRegistryRuntime(first)) + require.NoError(t, err) + require.True(t, ok) + assert.Same(t, firstRunner, gotFirst.Runner) + + gotSecond, ok, err := registry.Lookup(ctx, inboundForRegistryRuntime(second)) + require.NoError(t, err) + require.True(t, ok) + assert.Same(t, secondRunner, gotSecond.Runner) +} + +func inboundForRegistryRuntime(runtime Runtime) platform.InboundMessage { + return platform.InboundMessage{ + TenantID: runtime.Tenant.TenantID, + AppID: runtime.App.AppID, + BindingID: runtime.Binding.BindingID, + Channel: runtime.Binding.Channel, + ChannelAccountID: runtime.Binding.AccountID, + PlatformMessageID: "msg-1", + ExternalUserID: "user-1", + ConversationType: platform.ConversationTypeDM, + MessageType: platform.MessageTypeText, + ContentParts: nil, + } +} diff --git a/platform/gateway/service.go b/platform/gateway/service.go index b59d51fed0..9fa202fd52 100644 --- a/platform/gateway/service.go +++ b/platform/gateway/service.go @@ -10,8 +10,6 @@ package gateway import ( "context" - "crypto/sha256" - "encoding/hex" "errors" "fmt" "strings" @@ -23,6 +21,7 @@ import ( "trpc.group/trpc-go/trpc-agent-go/agent" "trpc.group/trpc-go/trpc-agent-go/event" + itelemetry "trpc.group/trpc-go/trpc-agent-go/internal/telemetry" "trpc.group/trpc-go/trpc-agent-go/model" "trpc.group/trpc-go/trpc-agent-go/platform" "trpc.group/trpc-go/trpc-agent-go/platform/channeladapter" @@ -418,7 +417,7 @@ func (s *Service) writeReply( input inboundRunInput, content string, ) (Result, error) { - resultRef := input.Key + ":outbound:1" + reply := newReplyPlan(input.Key, 0) outbound := platform.OutboundMessage{ TenantID: msg.TenantID, BindingID: msg.BindingID, @@ -427,14 +426,14 @@ func (s *Service) writeReply( ReplyToPlatformMessageID: msg.PlatformMessageID, Kind: platform.OutboundMessageKindText, Content: content, - Sequence: 1, - DedupKey: resultRef, + Sequence: reply.OutboundSequence, + DedupKey: reply.ResultRef, TraceID: input.RequestID, } replyCtx, replySpan := telemetrytrace.Tracer.Start(routeCtx, "im.reply") defer replySpan.End() setInboundTraceAttributes(replySpan, msg, input.SessionID, input.RequestID, input.InternalUserID) - if err := s.outboundStore.Save(replyCtx, resultRef, outbound); err != nil { + if err := s.outboundStore.Save(replyCtx, reply.ResultRef, outbound); err != nil { s.writeAudit(auditCtx, auditFromMessage(msg, input.SessionID, input.InternalUserID, "outbound_error", err.Error(), input.Start, err)) recordSpanError(replySpan, err) return Result{}, err @@ -444,7 +443,7 @@ func (s *Service) writeReply( outbound, channeladapter.RetryPolicyForBinding(runtime.Binding), ); err != nil { - if _, markErr := s.idempotencyStore.MarkReplyFailed(replyCtx, input.Key, resultRef); markErr != nil { + if _, markErr := s.idempotencyStore.MarkReplyFailed(replyCtx, input.Key, reply.ResultRef); markErr != nil { recordSpanError(replySpan, markErr) return Result{}, markErr } @@ -452,24 +451,42 @@ func (s *Service) writeReply( recordSpanError(replySpan, err) return Result{}, err } - record, err := s.idempotencyStore.Complete(replyCtx, input.Key, resultRef) + record, err := s.idempotencyStore.Complete(replyCtx, input.Key, reply.ResultRef) if err != nil { recordSpanError(replySpan, err) return Result{}, err } - s.writeMessageEvent(auditCtx, messageEventFromInbound(msg, input.SessionID, input.Key, input.RequestID, 1, input.Start)) - s.writeMessageEvent(auditCtx, messageEventFromAssistant(msg, input.SessionID, resultRef, input.RequestID, 2, s.now())) + s.writeMessageEvent(auditCtx, messageEventFromInbound(msg, input.SessionID, input.Key, input.RequestID, reply.InboundSequence, input.Start)) + s.writeMessageEvent(auditCtx, messageEventFromAssistant(msg, input.SessionID, reply.ResultRef, input.RequestID, reply.AssistantSequence, s.now())) s.writeAudit(auditCtx, auditFromMessage(msg, input.SessionID, input.InternalUserID, "completed", "", input.Start, nil)) return Result{ RequestID: input.RequestID, SessionID: input.SessionID, - ResultRef: resultRef, + ResultRef: reply.ResultRef, Status: record.Status, Outbound: outbound, CompletedAt: s.now(), }, nil } +type replyPlan struct { + ResultRef string + InboundSequence int64 + OutboundSequence int + AssistantSequence int64 +} + +func newReplyPlan(idempotencyKey string, outboundIndex int) replyPlan { + outboundSequence := outboundIndex + 1 + inboundSequence := int64(1) + return replyPlan{ + ResultRef: fmt.Sprintf("%s:outbound:%d", idempotencyKey, outboundSequence), + InboundSequence: inboundSequence, + OutboundSequence: outboundSequence, + AssistantSequence: inboundSequence + int64(outboundSequence), + } +} + func (s *Service) writeRejectAudit( ctx context.Context, msg platform.InboundMessage, @@ -763,12 +780,7 @@ func setInboundTraceAttributes( } func traceSafeHash(scope string, value string) string { - value = strings.TrimSpace(value) - if value == "" { - return "" - } - sum := sha256.Sum256([]byte(scope + "\x00" + value)) - return scope + "_hash_" + hex.EncodeToString(sum[:])[:24] + return itelemetry.TraceSafeHash(scope, value) } func recordSpanError(span oteltrace.Span, err error) { diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go index ead64ed97d..cd653b7095 100644 --- a/platform/gateway/service_test.go +++ b/platform/gateway/service_test.go @@ -1036,6 +1036,19 @@ func TestCollectAssistantTextPrefersFinalFullMessage(t *testing.T) { assert.Equal(t, "hello", content) } +func TestNewReplyPlanDerivesResultRefsAndSequences(t *testing.T) { + first := newReplyPlan("tenant:tenant-a:message:msg-1", 0) + assert.Equal(t, "tenant:tenant-a:message:msg-1:outbound:1", first.ResultRef) + assert.Equal(t, int64(1), first.InboundSequence) + assert.Equal(t, 1, first.OutboundSequence) + assert.Equal(t, int64(2), first.AssistantSequence) + + second := newReplyPlan("tenant:tenant-a:message:msg-1", 1) + assert.Equal(t, "tenant:tenant-a:message:msg-1:outbound:2", second.ResultRef) + assert.Equal(t, 2, second.OutboundSequence) + assert.Equal(t, int64(3), second.AssistantSequence) +} + func registerRuntime(t *testing.T, registry *InMemoryRegistry, tenantID string, r runnerStub) { t.Helper() err := registry.Register(validRuntime(tenantID, r)) diff --git a/platform/identity.go b/platform/identity.go index 77939f0b27..89edb68a61 100644 --- a/platform/identity.go +++ b/platform/identity.go @@ -11,11 +11,15 @@ package platform import ( "crypto/sha256" "encoding/hex" + "errors" "fmt" "net/url" "strings" ) +// ErrThreadIDRequired indicates a missing thread identifier for threaded conversations. +var ErrThreadIDRequired = errors.New("thread_id is required") + type stableIDPart struct { name string value string @@ -164,7 +168,7 @@ func sessionConversationParts( if err := validateRoutingIdentifier("external_group_id", externalGroupID, ErrExternalGroupIDRequired); err != nil { return nil, err } - if err := validateRoutingIdentifier("thread_id", threadID, fmt.Errorf("thread_id is required")); err != nil { + if err := validateRoutingIdentifier("thread_id", threadID, ErrThreadIDRequired); err != nil { return nil, err } return []stableIDPart{ diff --git a/platform/inmemory_sink.go b/platform/inmemory_sink.go new file mode 100644 index 0000000000..a4e3f401f9 --- /dev/null +++ b/platform/inmemory_sink.go @@ -0,0 +1,42 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "context" + "sync" +) + +type inMemoryRecords[T any] struct { + mu sync.Mutex + records []T +} + +func (s *inMemoryRecords[T]) append(ctx context.Context, record T, validate func(T) error) error { + if err := ctx.Err(); err != nil { + return err + } + if validate != nil { + if err := validate(record); err != nil { + return err + } + } + s.mu.Lock() + defer s.mu.Unlock() + s.records = append(s.records, record) + return nil +} + +func (s *inMemoryRecords[T]) snapshot() []T { + s.mu.Lock() + defer s.mu.Unlock() + out := make([]T, len(s.records)) + copy(out, s.records) + return out +} diff --git a/platform/message_event.go b/platform/message_event.go index 1509e8083e..1357c1e62c 100644 --- a/platform/message_event.go +++ b/platform/message_event.go @@ -10,7 +10,6 @@ package platform import ( "context" - "sync" ) // MessageEventSink stores immutable conversation events. @@ -21,8 +20,7 @@ type MessageEventSink interface { // InMemoryMessageEventSink is a concurrency-safe message event sink for tests and demos. type InMemoryMessageEventSink struct { - mu sync.Mutex - events []MessageEvent + events inMemoryRecords[MessageEvent] } // NewInMemoryMessageEventSink creates an in-memory message event sink. @@ -32,23 +30,10 @@ func NewInMemoryMessageEventSink() *InMemoryMessageEventSink { // WriteMessageEvent writes one message event. func (s *InMemoryMessageEventSink) WriteMessageEvent(ctx context.Context, event MessageEvent) error { - if err := ctx.Err(); err != nil { - return err - } - if err := event.Validate(); err != nil { - return err - } - s.mu.Lock() - defer s.mu.Unlock() - s.events = append(s.events, event) - return nil + return s.events.append(ctx, event, MessageEvent.Validate) } // Events returns a snapshot of written message events. func (s *InMemoryMessageEventSink) Events() []MessageEvent { - s.mu.Lock() - defer s.mu.Unlock() - out := make([]MessageEvent, len(s.events)) - copy(out, s.events) - return out + return s.events.snapshot() } diff --git a/platform/message_event_test.go b/platform/message_event_test.go index 5e23f5f5e8..dd1fa88565 100644 --- a/platform/message_event_test.go +++ b/platform/message_event_test.go @@ -49,6 +49,27 @@ func TestMessageEventValidateRequiresIdentity(t *testing.T) { } } +func TestMessageEventValidateRejectsNonNormalizedRoutingIdentifiers(t *testing.T) { + tests := []struct { + name string + mutate func(*MessageEvent) + want string + }{ + {name: "tenant", mutate: func(e *MessageEvent) { e.TenantID = "tenant\n" }, want: "tenant_id"}, + {name: "app", mutate: func(e *MessageEvent) { e.AppID = " app" }, want: "app_id"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + event := validMessageEvent() + tt.mutate(&event) + if err := event.Validate(); err == nil || !strings.Contains(err.Error(), tt.want) { + t.Fatalf("expected %s routing validation, got %v", tt.want, err) + } + }) + } +} + func TestMessageEventValidateRejectsInvalidRoleTypeAndSequence(t *testing.T) { tests := []struct { name string diff --git a/platform/migration_test.go b/platform/migration_test.go index e5558b2af4..0b4a390e2b 100644 --- a/platform/migration_test.go +++ b/platform/migration_test.go @@ -41,6 +41,12 @@ func TestNormalizeStorageMigrationModeAcceptsDocumentedModes(t *testing.T) { } } +func TestNormalizeStorageMigrationModeRejectsUnknown(t *testing.T) { + if _, err := NormalizeStorageMigrationMode("dual-read"); err == nil { + t.Fatalf("expected unknown migration mode error") + } +} + func TestStorageProfileValidateRejectsInvalidMigrationMode(t *testing.T) { profile := StorageProfile{ TenantID: "tenant", diff --git a/platform/operational_action_audit.go b/platform/operational_action_audit.go index 5221504ec1..14d72ed9a5 100644 --- a/platform/operational_action_audit.go +++ b/platform/operational_action_audit.go @@ -135,20 +135,18 @@ func (i OperationalActionAuditInput) normalize() (OperationalActionAuditInput, e if len(i.DetailJSON) > 0 && !json.Valid(i.DetailJSON) { return OperationalActionAuditInput{}, fmt.Errorf("detail_json must be valid json") } - for field, value := range map[string]string{ - "app_id": i.AppID, - "action": string(i.Action), - "operation_id": i.OperationID, - "resource_type": i.ResourceType, - "actor_internal_user_id": i.ActorInternalUserID, - "decision": string(i.Decision), - "decision_reason": i.DecisionReason, - "request_id": i.RequestID, - "trace_id": i.TraceID, - } { - if err := validateAuditRedactedText(field, value); err != nil { - return OperationalActionAuditInput{}, err - } + if err := validateAuditRedactedFields( + safeTextField{"app_id", i.AppID}, + safeTextField{"action", string(i.Action)}, + safeTextField{"operation_id", i.OperationID}, + safeTextField{"resource_type", i.ResourceType}, + safeTextField{"actor_internal_user_id", i.ActorInternalUserID}, + safeTextField{"decision", string(i.Decision)}, + safeTextField{"decision_reason", i.DecisionReason}, + safeTextField{"request_id", i.RequestID}, + safeTextField{"trace_id", i.TraceID}, + ); err != nil { + return OperationalActionAuditInput{}, err } return i, nil } diff --git a/platform/operational_action_audit_test.go b/platform/operational_action_audit_test.go index 1d2c9d4564..bba119dec9 100644 --- a/platform/operational_action_audit_test.go +++ b/platform/operational_action_audit_test.go @@ -91,6 +91,26 @@ func TestNewOperationalActionAuditRecordBuildsSafeRecord(t *testing.T) { } } +func TestNewOperationalActionAuditRecordAcceptsInternalActorOnly(t *testing.T) { + input := validOperationalActionAuditInput() + input.ActorUserID = "" + input.ActorInternalUserID = "usr_internal" + + record, err := NewOperationalActionAuditRecord(input) + if err != nil { + t.Fatalf("new internal actor operational action audit: %v", err) + } + if record.InternalUserID != "usr_internal" { + t.Fatalf("expected internal actor identity, got %+v", record) + } + if record.UserIDHash == "" || !strings.HasPrefix(record.UserIDHash, "user_hash_") { + t.Fatalf("expected internal actor hash, got %+v", record) + } + if strings.Contains(record.UserIDHash, "usr_internal") { + t.Fatalf("user hash leaked internal actor id: %q", record.UserIDHash) + } +} + func TestNewOperationalActionAuditRecordRejectsInvalidInputs(t *testing.T) { base := validOperationalActionAuditInput() diff --git a/platform/secret_rotation_status.go b/platform/secret_rotation_status.go index 6c193bda91..3f6c274f1f 100644 --- a/platform/secret_rotation_status.go +++ b/platform/secret_rotation_status.go @@ -135,21 +135,19 @@ func (r SecretRotationStatusReport) Validate() error { if r.UpdatedAt.IsZero() { return fmt.Errorf("updated_at is required") } - for field, value := range map[string]string{ - "app_id": r.AppID, - "rotation_id": r.RotationID, - "resource_type": r.ResourceType, - "resource_hash": r.ResourceHash, - "secret_field": r.SecretField, - "operation_id": r.OperationID, - "failure_reason": r.FailureReason, - "trace_id": r.TraceID, - } { - if err := validateAuditRedactedText(field, value); err != nil { - return err - } + if err := validateAuditRedactedFields( + safeTextField{"app_id", r.AppID}, + safeTextField{"rotation_id", r.RotationID}, + safeTextField{"resource_type", r.ResourceType}, + safeTextField{"resource_hash", r.ResourceHash}, + safeTextField{"secret_field", r.SecretField}, + safeTextField{"operation_id", r.OperationID}, + safeTextField{"failure_reason", r.FailureReason}, + safeTextField{"trace_id", r.TraceID}, + ); err != nil { + return err } - return nil + return validateSecretRotationStatusGate(r) } func (i SecretRotationStatusInput) normalize() (SecretRotationStatusInput, error) { @@ -194,21 +192,33 @@ func (i SecretRotationStatusInput) normalize() (SecretRotationStatusInput, error if i.UpdatedAt.IsZero() { return SecretRotationStatusInput{}, fmt.Errorf("updated_at is required") } - for field, value := range map[string]string{ - "app_id": i.AppID, - "resource_type": i.ResourceType, - "secret_field": i.SecretField, - "operation_id": i.OperationID, - "failure_reason": i.FailureReason, - "trace_id": i.TraceID, - } { - if err := validateAuditRedactedText(field, value); err != nil { - return SecretRotationStatusInput{}, err - } + if err := validateAuditRedactedFields( + safeTextField{"app_id", i.AppID}, + safeTextField{"resource_type", i.ResourceType}, + safeTextField{"secret_field", i.SecretField}, + safeTextField{"operation_id", i.OperationID}, + safeTextField{"failure_reason", i.FailureReason}, + safeTextField{"trace_id", i.TraceID}, + ); err != nil { + return SecretRotationStatusInput{}, err } return i, nil } +func validateSecretRotationStatusGate(r SecretRotationStatusReport) error { + switch r.Status { + case SecretRotationStatusFailed: + if strings.TrimSpace(r.FailureReason) == "" { + return fmt.Errorf("failure_reason is required when secret rotation status is failed") + } + case SecretRotationStatusActive, SecretRotationStatusRolledBack: + if strings.TrimSpace(r.PreviousRef) == "" { + return fmt.Errorf("previous_ref is required when secret rotation status is %s", r.Status) + } + } + return nil +} + func (i SecretRotationStatusInput) rotationID() string { return secretRotationIDPrefix + shortHash( i.TenantID, diff --git a/platform/secret_rotation_status_test.go b/platform/secret_rotation_status_test.go index cebb0e0ed8..87e5c5115b 100644 --- a/platform/secret_rotation_status_test.go +++ b/platform/secret_rotation_status_test.go @@ -147,6 +147,32 @@ func TestNewSecretRotationStatusReportRejectsInvalidInputs(t *testing.T) { } } +func TestSecretRotationStatusReportValidateEnforcesStatusGates(t *testing.T) { + base := validSecretRotationStatusInput() + base.Status = SecretRotationStatusFailed + base.FailureReason = "" + if _, err := NewSecretRotationStatusReport(base); err == nil || + !strings.Contains(err.Error(), "failure_reason") { + t.Fatalf("expected failed status to require failure reason, got %v", err) + } + + base = validSecretRotationStatusInput() + base.Status = SecretRotationStatusActive + base.PreviousRef = "" + if _, err := NewSecretRotationStatusReport(base); err == nil || + !strings.Contains(err.Error(), "previous_ref") { + t.Fatalf("expected active status to require previous ref, got %v", err) + } + + base = validSecretRotationStatusInput() + base.Status = SecretRotationStatusRolledBack + base.PreviousRef = "" + if _, err := NewSecretRotationStatusReport(base); err == nil || + !strings.Contains(err.Error(), "previous_ref") { + t.Fatalf("expected rolled back status to require previous ref, got %v", err) + } +} + func TestSecretRotationStatusReportValidateRejectsUnsafeReport(t *testing.T) { generated, err := NewSecretRotationStatusReport(validSecretRotationStatusInput()) if err != nil { diff --git a/platform/storagerouter/errors.go b/platform/storagerouter/errors.go index c5e03604d8..66e9371d07 100644 --- a/platform/storagerouter/errors.go +++ b/platform/storagerouter/errors.go @@ -17,6 +17,8 @@ var ( ErrTenantMismatch = errors.New("storage router tenant mismatch") // ErrBackendNotFound indicates that a requested backend is not registered. ErrBackendNotFound = errors.New("storage router backend not found") + // ErrBackendIDRequired indicates that a backend registration is missing its backend_id. + ErrBackendIDRequired = errors.New("storage router backend id required") // ErrBackendTenantMismatch indicates that a registered backend belongs to another tenant. ErrBackendTenantMismatch = errors.New("storage router backend tenant mismatch") ) diff --git a/platform/storagerouter/router.go b/platform/storagerouter/router.go index b6475bc5c4..b5643bd285 100644 --- a/platform/storagerouter/router.go +++ b/platform/storagerouter/router.go @@ -10,7 +10,6 @@ package storagerouter import ( "context" - "fmt" "strings" "sync" @@ -40,6 +39,7 @@ type Router interface { Artifact(ctx context.Context, tenantID string, profileID string) (artifact.Service, error) Knowledge(ctx context.Context, tenantID string, profileID string) (knowledge.Knowledge, error) Audit(ctx context.Context, tenantID string, profileID string) (platform.AuditSink, error) + Status(ctx context.Context, tenantID string, profileID string) (StatusSummary, error) } // InMemoryRouter is a concurrency-safe storage router for tests and demos. @@ -87,7 +87,7 @@ func (r *InMemoryRouter) RegisterBackend(backend BackendSet) error { return platform.ErrTenantIDRequired } if strings.TrimSpace(backend.BackendID) == "" { - return fmt.Errorf("backend_id is required") + return ErrBackendIDRequired } r.mu.Lock() defer r.mu.Unlock() @@ -113,6 +113,7 @@ func (r *InMemoryRouter) Profile( if !ok { return platform.StorageProfile{}, ErrProfileNotFound } + // Defensive for future persistent backends that may not key by tenant_id. if profile.TenantID != tenantID { return platform.StorageProfile{}, ErrTenantMismatch } @@ -229,6 +230,7 @@ func (r *InMemoryRouter) backend( if !ok { return BackendSet{}, ErrBackendNotFound } + // Defensive for future persistent backends that may not key by tenant_id. if backend.TenantID != tenantID { return BackendSet{}, ErrBackendTenantMismatch } diff --git a/platform/storagerouter/router_test.go b/platform/storagerouter/router_test.go index 830e76b210..1cba3b6c7a 100644 --- a/platform/storagerouter/router_test.go +++ b/platform/storagerouter/router_test.go @@ -23,6 +23,8 @@ import ( sessioninmemory "trpc.group/trpc-go/trpc-agent-go/session/inmemory" ) +var _ Router = (*InMemoryRouter)(nil) + func TestRouterResolvesTenantStorageServices(t *testing.T) { ctx := context.Background() router := NewInMemoryRouter() @@ -60,6 +62,14 @@ func TestRouterResolvesTenantStorageServices(t *testing.T) { assert.Same(t, auditSink, gotAudit) } +func TestRegisterBackendRequiresBackendID(t *testing.T) { + router := NewInMemoryRouter() + + err := router.RegisterBackend(BackendSet{TenantID: "tenant-a", BackendID: " "}) + + require.ErrorIs(t, err, ErrBackendIDRequired) +} + func TestRouterRejectsCrossTenantLookup(t *testing.T) { ctx := context.Background() router := NewInMemoryRouter() diff --git a/platform/storagerouter/status.go b/platform/storagerouter/status.go index 0964e1ee2d..12c9d0c415 100644 --- a/platform/storagerouter/status.go +++ b/platform/storagerouter/status.go @@ -12,11 +12,18 @@ import ( "context" "fmt" "strings" + "sync" "unicode" "trpc.group/trpc-go/trpc-agent-go/platform" ) +var ( + statusRedactorOnce sync.Once + statusRedactor *platform.Redactor + statusRedactorErr error +) + // ResourceStatus describes the routing readiness of one storage resource. type ResourceStatus string @@ -125,6 +132,7 @@ func (r *InMemoryRouter) resourceStatus( entry.Reason = fmt.Sprintf("%s backend is not registered", resource) return entry } + // Defensive for future persistent backends that may not key by tenant_id. if backend.TenantID != profile.TenantID { entry.Status = ResourceStatusBackendTenantMismatch entry.Reason = fmt.Sprintf("%s backend belongs to another tenant", resource) @@ -162,7 +170,7 @@ func safeBackendIDForStatus(backendID string) string { strings.ContainsAny(backendID, "=@/\\") { return "" } - redactor, err := platform.NewRedactor() + redactor, err := statusBackendIDRedactor() if err != nil || redactor.Redact(backendID) != backendID { return "" } @@ -179,3 +187,10 @@ func safeBackendIDForStatus(backendID string) string { } return backendID } + +func statusBackendIDRedactor() (*platform.Redactor, error) { + statusRedactorOnce.Do(func() { + statusRedactor, statusRedactorErr = platform.NewRedactor() + }) + return statusRedactor, statusRedactorErr +} diff --git a/platform/toolpolicy/policy.go b/platform/toolpolicy/policy.go index 7be766973f..89771b21fe 100644 --- a/platform/toolpolicy/policy.go +++ b/platform/toolpolicy/policy.go @@ -285,32 +285,35 @@ func (r *Reviewer) Review(ctx context.Context, req *review.Request) (*review.Dec } func (p *Policy) decide(req *tool.PermissionRequest, name string) (tool.PermissionDecision, string, bool) { - if contains(policyDenylist(p.policy), name) { - reason := fmt.Sprintf("tool %q is denied by platform tool policy", name) - return tool.DenyPermission(reason), reason, true - } - if len(normalizedList(p.policy.ToolWhitelist)) > 0 && - !contains(normalizedList(p.policy.ToolWhitelist), name) { - reason := fmt.Sprintf("tool %q is not in platform tool whitelist", name) - return tool.DenyPermission(reason), reason, true - } - if isHighRisk(p.policy, req, name) { - switch p.policy.DangerousToolAction { - case platform.DangerousToolActionDeny: - reason := fmt.Sprintf("high-risk tool %q is denied by platform tool policy", name) - return tool.DenyPermission(reason), reason, true - case platform.DangerousToolActionAsk: - reason := fmt.Sprintf("high-risk tool %q requires approval by platform tool policy", name) - return tool.AskPermission(reason), reason, true - case platform.DangerousToolActionAllowWithAudit, "": - reason := fmt.Sprintf("high-risk tool %q allowed with audit by platform tool policy", name) - return tool.AllowPermission(), reason, true - } - } - return tool.AllowPermission(), "", false + return p.decideWithOptions(req, name, decisionOptions{includeMetadataRisk: true}) } func (p *Policy) decideNameOnly(req *tool.PermissionRequest, name string) (tool.PermissionDecision, string, bool) { + return p.decideWithOptions(req, name, decisionOptions{ + includeNameRisk: true, + includeMetadataRisk: true, + }) +} + +func (p *Policy) decideReviewer(req *tool.PermissionRequest, name string) (tool.PermissionDecision, string, bool) { + return p.decideWithOptions(req, name, decisionOptions{ + includeNameRisk: true, + includeMetadataRisk: true, + reviewerAsk: true, + }) +} + +type decisionOptions struct { + includeNameRisk bool + includeMetadataRisk bool + reviewerAsk bool +} + +func (p *Policy) decideWithOptions( + req *tool.PermissionRequest, + name string, + opts decisionOptions, +) (tool.PermissionDecision, string, bool) { if contains(policyDenylist(p.policy), name) { reason := fmt.Sprintf("tool %q is denied by platform tool policy", name) return tool.DenyPermission(reason), reason, true @@ -320,12 +323,16 @@ func (p *Policy) decideNameOnly(req *tool.PermissionRequest, name string) (tool. reason := fmt.Sprintf("tool %q is not in platform tool whitelist", name) return tool.DenyPermission(reason), reason, true } - if contains(normalizedList(p.policy.HighRiskTools), name) { + if p.highRiskForDecision(req, name, opts) { switch p.policy.DangerousToolAction { case platform.DangerousToolActionDeny: reason := fmt.Sprintf("high-risk tool %q is denied by platform tool policy", name) return tool.DenyPermission(reason), reason, true case platform.DangerousToolActionAsk: + if opts.reviewerAsk { + reason := fmt.Sprintf("high-risk tool %q approved by platform approval reviewer", name) + return tool.AllowPermission(), reason, true + } reason := fmt.Sprintf("high-risk tool %q requires approval by platform tool policy", name) return tool.AskPermission(reason), reason, true case platform.DangerousToolActionAllowWithAudit, "": @@ -333,39 +340,22 @@ func (p *Policy) decideNameOnly(req *tool.PermissionRequest, name string) (tool. return tool.AllowPermission(), reason, true } } - if req != nil && req.Metadata != (tool.ToolMetadata{}) { - return p.decide(req, name) - } return tool.AllowPermission(), "", false } -func (p *Policy) decideReviewer(req *tool.PermissionRequest, name string) (tool.PermissionDecision, string, bool) { - if contains(policyDenylist(p.policy), name) { - reason := fmt.Sprintf("tool %q is denied by platform tool policy", name) - return tool.DenyPermission(reason), reason, true - } - if len(normalizedList(p.policy.ToolWhitelist)) > 0 && - !contains(normalizedList(p.policy.ToolWhitelist), name) { - reason := fmt.Sprintf("tool %q is not in platform tool whitelist", name) - return tool.DenyPermission(reason), reason, true - } - if contains(normalizedList(p.policy.HighRiskTools), name) { - switch p.policy.DangerousToolAction { - case platform.DangerousToolActionDeny: - reason := fmt.Sprintf("high-risk tool %q is denied by platform tool policy", name) - return tool.DenyPermission(reason), reason, true - case platform.DangerousToolActionAsk: - reason := fmt.Sprintf("high-risk tool %q approved by platform approval reviewer", name) - return tool.AllowPermission(), reason, true - case platform.DangerousToolActionAllowWithAudit, "": - reason := fmt.Sprintf("high-risk tool %q allowed with audit by platform tool policy", name) - return tool.AllowPermission(), reason, true - } +func (p *Policy) highRiskForDecision( + req *tool.PermissionRequest, + name string, + opts decisionOptions, +) bool { + if (opts.includeNameRisk || opts.includeMetadataRisk) && + contains(normalizedList(p.policy.HighRiskTools), name) { + return true } - if req != nil && req.Metadata != (tool.ToolMetadata{}) { - return p.decide(req, name) + if !opts.includeMetadataRisk || req == nil || req.Metadata == (tool.ToolMetadata{}) { + return false } - return tool.AllowPermission(), "", false + return req.Metadata.Destructive || !req.Metadata.ReadOnly || req.Metadata.OpenWorld } func validate(policy platform.ToolPolicy) error { @@ -648,14 +638,6 @@ func argumentDigest(args []byte) (string, int) { return "sha256:" + hex.EncodeToString(sum[:]), len(args) } -func argumentSummary(args []byte) string { - digest, size := argumentDigest(args) - if digest == "" { - return "" - } - return fmt.Sprintf("%s bytes:%d", digest, size) -} - var sha256DigestPattern = regexp.MustCompile(`^sha256:[a-f0-9]{64}$`) func validSHA256Digest(value string) bool { diff --git a/platform/types_test.go b/platform/types_test.go index 0e49a84977..4125bb8c77 100644 --- a/platform/types_test.go +++ b/platform/types_test.go @@ -196,6 +196,34 @@ func TestValidateRejectsNonNormalizedRoutingIdentifiers(t *testing.T) { return AgentApp{TenantID: "tenant", AppID: value}.Validate() }, }, + { + name: "AppConfigVersionTenantID", + validate: func(value string) error { + version := validAppConfigVersion() + version.TenantID = value + return version.Validate() + }, + }, + { + name: "AppConfigVersionAppID", + validate: func(value string) error { + version := validAppConfigVersion() + version.AppID = value + return version.Validate() + }, + }, + { + name: "AuditPolicyTenantID", + validate: func(value string) error { + return AuditPolicy{TenantID: value, PolicyID: "policy"}.Validate() + }, + }, + { + name: "AuditPolicyPolicyID", + validate: func(value string) error { + return AuditPolicy{TenantID: "tenant", PolicyID: value}.Validate() + }, + }, { name: "BindingID", validate: func(value string) error { diff --git a/platform/usage.go b/platform/usage.go index 037fc7d59b..640daf4d59 100644 --- a/platform/usage.go +++ b/platform/usage.go @@ -10,7 +10,6 @@ package platform import ( "context" - "sync" ) // UsageSink stores post-run usage records. @@ -21,8 +20,7 @@ type UsageSink interface { // InMemoryUsageSink is a concurrency-safe usage sink for tests and demos. type InMemoryUsageSink struct { - mu sync.Mutex - records []UsageRecord + records inMemoryRecords[UsageRecord] } // NewInMemoryUsageSink creates an in-memory usage sink. @@ -32,23 +30,10 @@ func NewInMemoryUsageSink() *InMemoryUsageSink { // WriteUsage writes one usage record. func (s *InMemoryUsageSink) WriteUsage(ctx context.Context, record UsageRecord) error { - if err := ctx.Err(); err != nil { - return err - } - if err := record.Validate(); err != nil { - return err - } - s.mu.Lock() - defer s.mu.Unlock() - s.records = append(s.records, record) - return nil + return s.records.append(ctx, record, UsageRecord.Validate) } // Records returns a snapshot of written usage records. func (s *InMemoryUsageSink) Records() []UsageRecord { - s.mu.Lock() - defer s.mu.Unlock() - out := make([]UsageRecord, len(s.records)) - copy(out, s.records) - return out + return s.records.snapshot() } diff --git a/platform/usage_record_test.go b/platform/usage_record_test.go index d3837674df..64a81df302 100644 --- a/platform/usage_record_test.go +++ b/platform/usage_record_test.go @@ -45,6 +45,27 @@ func TestUsageRecordValidateRequiresTenantAndApp(t *testing.T) { } } +func TestUsageRecordValidateRejectsNonNormalizedRoutingIdentifiers(t *testing.T) { + tests := []struct { + name string + mutate func(*UsageRecord) + want string + }{ + {name: "tenant", mutate: func(r *UsageRecord) { r.TenantID = "tenant\x00" }, want: "tenant_id"}, + {name: "app", mutate: func(r *UsageRecord) { r.AppID = "app " }, want: "app_id"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + record := validUsageRecord() + tt.mutate(&record) + if err := record.Validate(); err == nil || !strings.Contains(err.Error(), tt.want) { + t.Fatalf("expected %s routing validation, got %v", tt.want, err) + } + }) + } +} + func TestUsageRecordValidateRejectsNegativeTokens(t *testing.T) { tests := []struct { name string diff --git a/platform/validation.go b/platform/validation.go index 52d6de6f69..3804b6e5cb 100644 --- a/platform/validation.go +++ b/platform/validation.go @@ -28,6 +28,20 @@ var rawSecretPrefixes = []string{ "glpat-", } +type safeTextField struct { + name string + value string +} + +func validateAuditRedactedFields(fields ...safeTextField) error { + for _, field := range fields { + if err := validateAuditRedactedText(field.name, field.value); err != nil { + return err + } + } + return nil +} + func validateRoutingIdentifier(field, value string, requiredErr error) error { trimmed := strings.TrimSpace(value) if trimmed == "" { @@ -78,11 +92,11 @@ func (a AgentApp) Validate() error { // Validate checks that an app config version is safe to store and route. func (v AppConfigVersion) Validate() error { - if strings.TrimSpace(v.TenantID) == "" { - return ErrTenantIDRequired + if err := validateRoutingIdentifier("tenant_id", v.TenantID, ErrTenantIDRequired); err != nil { + return err } - if strings.TrimSpace(v.AppID) == "" { - return ErrAppIDRequired + if err := validateRoutingIdentifier("app_id", v.AppID, ErrAppIDRequired); err != nil { + return err } if strings.TrimSpace(v.Version) == "" { return fmt.Errorf("version is required") @@ -252,7 +266,7 @@ func (m InboundMessage) Validate() error { if err := validateRoutingIdentifier("external_group_id", m.ExternalGroupID, ErrExternalGroupIDRequired); err != nil { return err } - if err := validateRoutingIdentifier("thread_id", m.ThreadID, fmt.Errorf("thread_id is required")); err != nil { + if err := validateRoutingIdentifier("thread_id", m.ThreadID, ErrThreadIDRequired); err != nil { return err } return nil @@ -295,11 +309,11 @@ func (p StorageProfile) Validate() error { // Validate checks that audit retention and sampling policy is safe to use. func (p AuditPolicy) Validate() error { - if strings.TrimSpace(p.TenantID) == "" { - return ErrTenantIDRequired + if err := validateRoutingIdentifier("tenant_id", p.TenantID, ErrTenantIDRequired); err != nil { + return err } - if strings.TrimSpace(p.PolicyID) == "" { - return fmt.Errorf("policy_id is required") + if err := validateRoutingIdentifier("policy_id", p.PolicyID, fmt.Errorf("policy_id is required")); err != nil { + return err } if p.RetentionDays < 0 { return fmt.Errorf("retention_days must be greater than or equal to 0") @@ -352,11 +366,11 @@ func (r AuditRecord) Validate() error { // Validate checks that a message event has required identity and safe trace metadata. func (e MessageEvent) Validate() error { - if strings.TrimSpace(e.TenantID) == "" { - return ErrTenantIDRequired + if err := validateRoutingIdentifier("tenant_id", e.TenantID, ErrTenantIDRequired); err != nil { + return err } - if strings.TrimSpace(e.AppID) == "" { - return ErrAppIDRequired + if err := validateRoutingIdentifier("app_id", e.AppID, ErrAppIDRequired); err != nil { + return err } if strings.TrimSpace(e.SessionID) == "" { return fmt.Errorf("session_id is required") @@ -407,11 +421,11 @@ func (e MessageEvent) Validate() error { // Validate checks that a usage record has required identity and safe accounting values. func (r UsageRecord) Validate() error { - if strings.TrimSpace(r.TenantID) == "" { - return ErrTenantIDRequired + if err := validateRoutingIdentifier("tenant_id", r.TenantID, ErrTenantIDRequired); err != nil { + return err } - if strings.TrimSpace(r.AppID) == "" { - return ErrAppIDRequired + if err := validateRoutingIdentifier("app_id", r.AppID, ErrAppIDRequired); err != nil { + return err } for field, value := range map[string]string{ "user_id_hash": r.UserIDHash, diff --git a/runner/diagnostics.go b/runner/diagnostics.go index a697baad1e..1c1c53abd5 100644 --- a/runner/diagnostics.go +++ b/runner/diagnostics.go @@ -10,11 +10,8 @@ package runner import ( "context" - "crypto/sha256" - "encoding/hex" "errors" "fmt" - "strings" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/codes" @@ -23,6 +20,7 @@ import ( "trpc.group/trpc-go/trpc-agent-go/agent" "trpc.group/trpc-go/trpc-agent-go/event" + itelemetry "trpc.group/trpc-go/trpc-agent-go/internal/telemetry" itrace "trpc.group/trpc-go/trpc-agent-go/internal/trace" "trpc.group/trpc-go/trpc-agent-go/model" "trpc.group/trpc-go/trpc-agent-go/session" @@ -160,12 +158,7 @@ func runnerSessionAttrs(key session.Key, sess *session.Session) []attribute.KeyV } func runnerTraceSafeHash(scope string, value string) string { - value = strings.TrimSpace(value) - if value == "" { - return "" - } - sum := sha256.Sum256([]byte(scope + "\x00" + value)) - return scope + "_hash_" + hex.EncodeToString(sum[:])[:24] + return itelemetry.TraceSafeHash(scope, value) } func runnerTraceErrorDescription(err error) string { diff --git a/session/redis/summary.go b/session/redis/summary.go index 26c38e0bc8..91ff44af8b 100644 --- a/session/redis/summary.go +++ b/session/redis/summary.go @@ -34,13 +34,13 @@ func (s *Service) CreateSessionSummary(ctx context.Context, sess *session.Sessio } key := session.Key{AppName: sess.AppName, UserID: sess.UserID, SessionID: sess.ID} + if err := key.CheckSessionKey(); err != nil { + return fmt.Errorf("check session key failed: %w", err) + } ctx, span := s.startSpan(ctx, "create_session_summary", key) itelemetry.MarkSummaryCreateSpan(span) defer span.End() - if err := key.CheckSessionKey(); err != nil { - return fmt.Errorf("check session key failed: %w", err) - } if !isummary.NewSummaryDispatchPolicy( s.opts.summaryFilterAllowlist, s.opts.shouldCascadeFullSessionSummary(), From 6fb7ee89a42a3b6b20d42fd910e5dbb941268d9d Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Thu, 9 Jul 2026 18:40:06 +0800 Subject: [PATCH 52/95] fix(platform): bound in-memory sink records --- platform/audit.go | 8 ++++--- platform/inmemory_sink.go | 44 +++++++++++++++++++++++++++++++++-- platform/message_event.go | 8 ++++--- platform/usage.go | 8 ++++--- platform/usage_record_test.go | 20 ++++++++++++++++ 5 files changed, 77 insertions(+), 11 deletions(-) diff --git a/platform/audit.go b/platform/audit.go index b4a14c5c13..9bca743481 100644 --- a/platform/audit.go +++ b/platform/audit.go @@ -18,14 +18,16 @@ type AuditSink interface { WriteAudit(ctx context.Context, record AuditRecord) error } -// InMemoryAuditSink is a concurrency-safe audit sink for tests and demos. +// InMemoryAuditSink is a concurrency-safe bounded audit sink for tests and demos. type InMemoryAuditSink struct { records inMemoryRecords[AuditRecord] } // NewInMemoryAuditSink creates an in-memory audit sink. -func NewInMemoryAuditSink() *InMemoryAuditSink { - return &InMemoryAuditSink{} +func NewInMemoryAuditSink(options ...InMemorySinkOption) *InMemoryAuditSink { + return &InMemoryAuditSink{ + records: newInMemoryRecords[AuditRecord](options...), + } } // WriteAudit writes one audit record. diff --git a/platform/inmemory_sink.go b/platform/inmemory_sink.go index a4e3f401f9..07991942bf 100644 --- a/platform/inmemory_sink.go +++ b/platform/inmemory_sink.go @@ -10,12 +10,41 @@ package platform import ( "context" + "fmt" "sync" ) +const defaultInMemoryRecordLimit = 1024 + +type inMemoryRecordOptions struct { + maxRecords int +} + +// InMemorySinkOption configures in-memory platform sinks. +type InMemorySinkOption func(*inMemoryRecordOptions) + +// WithInMemorySinkMaxRecords sets how many recent records an in-memory sink +// retains. The default is bounded to prevent unbounded demo/dev growth. +func WithInMemorySinkMaxRecords(maxRecords int) InMemorySinkOption { + return func(opts *inMemoryRecordOptions) { + opts.maxRecords = maxRecords + } +} + type inMemoryRecords[T any] struct { - mu sync.Mutex - records []T + mu sync.Mutex + records []T + maxRecords int +} + +func newInMemoryRecords[T any](options ...InMemorySinkOption) inMemoryRecords[T] { + opts := inMemoryRecordOptions{maxRecords: defaultInMemoryRecordLimit} + for _, option := range options { + if option != nil { + option(&opts) + } + } + return inMemoryRecords[T]{maxRecords: opts.maxRecords} } func (s *inMemoryRecords[T]) append(ctx context.Context, record T, validate func(T) error) error { @@ -29,6 +58,17 @@ func (s *inMemoryRecords[T]) append(ctx context.Context, record T, validate func } s.mu.Lock() defer s.mu.Unlock() + maxRecords := s.maxRecords + if maxRecords == 0 { + maxRecords = defaultInMemoryRecordLimit + } + if maxRecords < 0 { + return fmt.Errorf("in-memory sink max records must be positive") + } + if len(s.records) >= maxRecords { + copy(s.records, s.records[1:]) + s.records = s.records[:maxRecords-1] + } s.records = append(s.records, record) return nil } diff --git a/platform/message_event.go b/platform/message_event.go index 1357c1e62c..38e8d0016f 100644 --- a/platform/message_event.go +++ b/platform/message_event.go @@ -18,14 +18,16 @@ type MessageEventSink interface { WriteMessageEvent(ctx context.Context, event MessageEvent) error } -// InMemoryMessageEventSink is a concurrency-safe message event sink for tests and demos. +// InMemoryMessageEventSink is a concurrency-safe bounded message event sink for tests and demos. type InMemoryMessageEventSink struct { events inMemoryRecords[MessageEvent] } // NewInMemoryMessageEventSink creates an in-memory message event sink. -func NewInMemoryMessageEventSink() *InMemoryMessageEventSink { - return &InMemoryMessageEventSink{} +func NewInMemoryMessageEventSink(options ...InMemorySinkOption) *InMemoryMessageEventSink { + return &InMemoryMessageEventSink{ + events: newInMemoryRecords[MessageEvent](options...), + } } // WriteMessageEvent writes one message event. diff --git a/platform/usage.go b/platform/usage.go index 640daf4d59..83c0f229ec 100644 --- a/platform/usage.go +++ b/platform/usage.go @@ -18,14 +18,16 @@ type UsageSink interface { WriteUsage(ctx context.Context, record UsageRecord) error } -// InMemoryUsageSink is a concurrency-safe usage sink for tests and demos. +// InMemoryUsageSink is a concurrency-safe bounded usage sink for tests and demos. type InMemoryUsageSink struct { records inMemoryRecords[UsageRecord] } // NewInMemoryUsageSink creates an in-memory usage sink. -func NewInMemoryUsageSink() *InMemoryUsageSink { - return &InMemoryUsageSink{} +func NewInMemoryUsageSink(options ...InMemorySinkOption) *InMemoryUsageSink { + return &InMemoryUsageSink{ + records: newInMemoryRecords[UsageRecord](options...), + } } // WriteUsage writes one usage record. diff --git a/platform/usage_record_test.go b/platform/usage_record_test.go index 64a81df302..619db96dad 100644 --- a/platform/usage_record_test.go +++ b/platform/usage_record_test.go @@ -137,6 +137,26 @@ func TestUsageSinkStoresSnapshot(t *testing.T) { } } +func TestUsageSinkEvictsOldRecordsAtConfiguredLimit(t *testing.T) { + sink := NewInMemoryUsageSink(WithInMemorySinkMaxRecords(2)) + + for i := 0; i < 3; i++ { + record := validUsageRecord() + record.RequestID = string(rune('a' + i)) + if err := sink.WriteUsage(context.Background(), record); err != nil { + t.Fatalf("WriteUsage(%d): %v", i, err) + } + } + + records := sink.Records() + if len(records) != 2 { + t.Fatalf("expected capped records, got %d", len(records)) + } + if records[0].RequestID != "b" || records[1].RequestID != "c" { + t.Fatalf("expected newest records to be retained, got %+v", records) + } +} + func TestUsageSinkRejectsInvalidRecord(t *testing.T) { sink := NewInMemoryUsageSink() record := validUsageRecord() From 36896aac1669db7ff491f3c532925d4b7d2fbd89 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Fri, 10 Jul 2026 14:12:25 +0800 Subject: [PATCH 53/95] storage: add tenant-aware backend router --- platform/backend_migration_status.go | 3 + platform/backend_migration_status_test.go | 1 + platform/gateway/lease.go | 32 +- platform/gateway/lease_test.go | 53 ++ platform/gateway/service.go | 6 + platform/gateway/service_test.go | 51 +- platform/migration_test.go | 31 ++ platform/storage_context.go | 27 + platform/storagerouter/adapter.go | 575 ++++++++++++++++++++++ platform/storagerouter/errors.go | 2 + platform/storagerouter/router.go | 127 +++++ platform/storagerouter/router_test.go | 169 ++++++- platform/storagerouter/status.go | 3 + platform/storagerouter/status_test.go | 8 +- platform/validation.go | 32 ++ 15 files changed, 1097 insertions(+), 23 deletions(-) create mode 100644 platform/gateway/lease_test.go create mode 100644 platform/storage_context.go create mode 100644 platform/storagerouter/adapter.go diff --git a/platform/backend_migration_status.go b/platform/backend_migration_status.go index b70aec8ddb..7cc6353bbd 100644 --- a/platform/backend_migration_status.go +++ b/platform/backend_migration_status.go @@ -22,6 +22,8 @@ type BackendMigrationResource string const ( // BackendMigrationResourceSession covers session event storage migrations. BackendMigrationResourceSession BackendMigrationResource = "session" + // BackendMigrationResourceSummary covers session summary storage migrations. + BackendMigrationResourceSummary BackendMigrationResource = "summary" // BackendMigrationResourceMemory covers memory store migrations. BackendMigrationResourceMemory BackendMigrationResource = "memory" // BackendMigrationResourceArtifact covers artifact object storage migrations. @@ -427,6 +429,7 @@ func backendMigrationLag(sourceCount, targetCount int64) int64 { func (r BackendMigrationResource) valid() bool { switch r { case BackendMigrationResourceSession, + BackendMigrationResourceSummary, BackendMigrationResourceMemory, BackendMigrationResourceArtifact, BackendMigrationResourceKnowledge, diff --git a/platform/backend_migration_status_test.go b/platform/backend_migration_status_test.go index bbfc5f15db..984862c0f5 100644 --- a/platform/backend_migration_status_test.go +++ b/platform/backend_migration_status_test.go @@ -342,6 +342,7 @@ func TestNewBackendMigrationStatusReportEnforcesStatusGates(t *testing.T) { func TestNewBackendMigrationStatusReportSupportsAcceptanceResources(t *testing.T) { for _, resource := range []BackendMigrationResource{ BackendMigrationResourceSession, + BackendMigrationResourceSummary, BackendMigrationResourceMemory, BackendMigrationResourceArtifact, BackendMigrationResourceKnowledge, diff --git a/platform/gateway/lease.go b/platform/gateway/lease.go index f31c671e4c..4348b7590f 100644 --- a/platform/gateway/lease.go +++ b/platform/gateway/lease.go @@ -27,19 +27,22 @@ type SessionLeaseKey struct { // SessionLease releases one acquired session execution slot. type SessionLease interface { + FencingToken() int64 Release(ctx context.Context) error } // InMemorySessionLeaseStore is a process-local lease store for tests and demos. type InMemorySessionLeaseStore struct { - mu sync.Mutex - leases map[SessionLeaseKey]struct{} + mu sync.Mutex + leases map[SessionLeaseKey]int64 + counter map[SessionLeaseKey]int64 } // NewInMemorySessionLeaseStore creates an empty process-local session lease store. func NewInMemorySessionLeaseStore() *InMemorySessionLeaseStore { return &InMemorySessionLeaseStore{ - leases: make(map[SessionLeaseKey]struct{}), + leases: make(map[SessionLeaseKey]int64), + counter: make(map[SessionLeaseKey]int64), } } @@ -56,17 +59,28 @@ func (s *InMemorySessionLeaseStore) Acquire( if _, ok := s.leases[key]; ok { return nil, false, nil } - s.leases[key] = struct{}{} + token := s.counter[key] + 1 + s.counter[key] = token + s.leases[key] = token return &inMemorySessionLease{ - store: s, - key: key, + store: s, + key: key, + fencingToken: token, }, true, nil } type inMemorySessionLease struct { - store *InMemorySessionLeaseStore - key SessionLeaseKey - once sync.Once + store *InMemorySessionLeaseStore + key SessionLeaseKey + fencingToken int64 + once sync.Once +} + +func (l *inMemorySessionLease) FencingToken() int64 { + if l == nil { + return 0 + } + return l.fencingToken } func (l *inMemorySessionLease) Release(ctx context.Context) error { diff --git a/platform/gateway/lease_test.go b/platform/gateway/lease_test.go new file mode 100644 index 0000000000..932b9fa7f3 --- /dev/null +++ b/platform/gateway/lease_test.go @@ -0,0 +1,53 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package gateway + +import ( + "context" + "testing" +) + +func TestInMemorySessionLeaseStoreIssuesMonotonicFencingTokens(t *testing.T) { + ctx := context.Background() + store := NewInMemorySessionLeaseStore() + key := SessionLeaseKey{TenantID: "tenant", AppID: "app", SessionID: "session"} + + first, acquired, err := store.Acquire(ctx, key) + if err != nil { + t.Fatalf("acquire first: %v", err) + } + if !acquired { + t.Fatalf("first acquire should succeed") + } + if got := first.FencingToken(); got != 1 { + t.Fatalf("expected first fencing token 1, got %d", got) + } + + _, acquired, err = store.Acquire(ctx, key) + if err != nil { + t.Fatalf("acquire busy: %v", err) + } + if acquired { + t.Fatalf("same session should not acquire while held") + } + if err := first.Release(ctx); err != nil { + t.Fatalf("release first: %v", err) + } + + second, acquired, err := store.Acquire(ctx, key) + if err != nil { + t.Fatalf("acquire second: %v", err) + } + if !acquired { + t.Fatalf("second acquire should succeed after release") + } + if got := second.FencingToken(); got != 2 { + t.Fatalf("expected second fencing token 2, got %d", got) + } +} diff --git a/platform/gateway/service.go b/platform/gateway/service.go index 9fa202fd52..d970dec2c4 100644 --- a/platform/gateway/service.go +++ b/platform/gateway/service.go @@ -179,6 +179,7 @@ func (s *Service) HandleInbound( InternalUserID: internalUserID, RequestID: requestID, Key: key, + FencingToken: record.SessionLease.FencingToken(), Start: start, }, ) @@ -195,6 +196,7 @@ type inboundRunInput struct { InternalUserID string RequestID string Key string + FencingToken int64 Start time.Time } @@ -385,7 +387,11 @@ func (s *Service) runGatewayRunner( ) (string, error) { runnerCtx, runnerSpan := telemetrytrace.Tracer.Start(routeCtx, "runner.run") defer runnerSpan.End() + runnerCtx = platform.ContextWithStorageFencingToken(runnerCtx, input.FencingToken) setInboundTraceAttributes(runnerSpan, msg, input.SessionID, input.RequestID, input.InternalUserID) + if input.FencingToken > 0 { + runnerSpan.SetAttributes(attribute.Int64("storage.fencing_token", input.FencingToken)) + } ch, err := runtime.Runner.Run( runnerCtx, input.InternalUserID, diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go index cd653b7095..7702a86512 100644 --- a/platform/gateway/service_test.go +++ b/platform/gateway/service_test.go @@ -583,6 +583,26 @@ func TestServiceHandleInboundReleaseIgnoresCanceledRequestContext(t *testing.T) require.NoError(t, lease.ctxErr) } +func TestServiceHandleInboundPropagatesLeaseFencingToken(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "ok"} + registerRuntime(t, registry, "tenant-a", r) + lease := &recordingLease{token: 42} + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + WithSessionLeaseStore(&recordingLeaseStore{lease: lease}), + ) + + _, err := svc.HandleInbound(ctx, inbound("tenant-a", "msg-1", "user-1", "hello")) + + require.NoError(t, err) + require.Len(t, r.calls, 1) + assert.Equal(t, int64(42), r.calls[0].fencingToken) +} + func TestServiceHandleInboundCancellationDuringEventCollectionReleasesSessionLease(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) registry := NewInMemoryRegistry() @@ -1206,11 +1226,12 @@ type runnerStub interface { } type runnerCall struct { - userID string - sessionID string - message model.Message - requestID string - runOptions agent.RunOptions + userID string + sessionID string + message model.Message + requestID string + fencingToken int64 + runOptions agent.RunOptions } type recordingRunner struct { @@ -1251,12 +1272,14 @@ func (r *recordingRunner) Run( return nil, r.runErr } runOptions := runOptionsFromOptions(runOpts...) + fencingToken, _ := platform.StorageFencingTokenFromContext(ctx) r.calls = append(r.calls, runnerCall{ - userID: userID, - sessionID: sessionID, - message: message, - requestID: runOptions.RequestID, - runOptions: runOptions, + userID: userID, + sessionID: sessionID, + message: message, + requestID: runOptions.RequestID, + fencingToken: fencingToken, + runOptions: runOptions, }) out := make(chan *event.Event, 2) go func() { @@ -1432,6 +1455,14 @@ func (s *recordingLeaseStore) Acquire( type recordingLease struct { released bool ctxErr error + token int64 +} + +func (l *recordingLease) FencingToken() int64 { + if l.token == 0 { + return 1 + } + return l.token } func (l *recordingLease) Release(ctx context.Context) error { diff --git a/platform/migration_test.go b/platform/migration_test.go index 0b4a390e2b..afd1752a96 100644 --- a/platform/migration_test.go +++ b/platform/migration_test.go @@ -51,6 +51,7 @@ func TestStorageProfileValidateRejectsInvalidMigrationMode(t *testing.T) { profile := StorageProfile{ TenantID: "tenant", ProfileID: "profile", + Namespace: "tenant/tenant", MigrationMode: "dual-read", } if err := profile.Validate(); err == nil { @@ -58,6 +59,36 @@ func TestStorageProfileValidateRejectsInvalidMigrationMode(t *testing.T) { } } +func TestStorageProfileValidateRequiresTenantScopedNamespace(t *testing.T) { + valid := StorageProfile{ + TenantID: "tenant-a", + ProfileID: "profile", + Namespace: "tenant/tenant-a/profile/profile", + } + if err := valid.Validate(); err != nil { + t.Fatalf("expected tenant-scoped namespace to pass, got %v", err) + } + + tests := []struct { + name string + namespace string + }{ + {name: "missing", namespace: ""}, + {name: "other_tenant", namespace: "tenant/tenant-b/profile/profile"}, + {name: "shared", namespace: "shared/profile"}, + {name: "whitespace", namespace: " tenant/tenant-a/profile/profile "}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + profile := valid + profile.Namespace = tt.namespace + if err := profile.Validate(); err == nil { + t.Fatalf("expected namespace %q to fail", tt.namespace) + } + }) + } +} + func TestIsActiveStorageMigrationMode(t *testing.T) { if IsActiveStorageMigrationMode(StorageMigrationModeNormal) { t.Fatalf("normal mode should not be active migration") diff --git a/platform/storage_context.go b/platform/storage_context.go new file mode 100644 index 0000000000..18bc6838d3 --- /dev/null +++ b/platform/storage_context.go @@ -0,0 +1,27 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import "context" + +type storageFencingTokenContextKey struct{} + +// ContextWithStorageFencingToken returns a child context carrying a storage fencing token. +func ContextWithStorageFencingToken(ctx context.Context, token int64) context.Context { + if token <= 0 { + return ctx + } + return context.WithValue(ctx, storageFencingTokenContextKey{}, token) +} + +// StorageFencingTokenFromContext returns the storage fencing token carried by ctx. +func StorageFencingTokenFromContext(ctx context.Context) (int64, bool) { + token, ok := ctx.Value(storageFencingTokenContextKey{}).(int64) + return token, ok && token > 0 +} diff --git a/platform/storagerouter/adapter.go b/platform/storagerouter/adapter.go new file mode 100644 index 0000000000..d9da52549b --- /dev/null +++ b/platform/storagerouter/adapter.go @@ -0,0 +1,575 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package storagerouter + +import ( + "context" + "errors" + "strings" + + "trpc.group/trpc-go/trpc-agent-go/artifact" + "trpc.group/trpc-go/trpc-agent-go/event" + "trpc.group/trpc-go/trpc-agent-go/knowledge" + "trpc.group/trpc-go/trpc-agent-go/memory" + "trpc.group/trpc-go/trpc-agent-go/platform" + "trpc.group/trpc-go/trpc-agent-go/session" +) + +const tenantMetadataKey = "tenant_id" + +// StorageAdapter is a tenant/profile-bound storage facade. +type StorageAdapter interface { + Scope() StorageScope + Route(ctx context.Context, resource platform.BackendMigrationResource) (RouteBinding, error) + Session(ctx context.Context) (session.Service, error) + Summary(ctx context.Context) (SummaryStore, error) + Memory(ctx context.Context) (memory.Service, error) + Artifact(ctx context.Context) (artifact.Service, error) + Knowledge(ctx context.Context) (knowledge.Knowledge, error) + Audit(ctx context.Context) (platform.AuditSink, error) +} + +// StorageScope describes the tenant boundary applied by a StorageAdapter. +type StorageScope struct { + TenantID string + ProfileID string + Namespace string +} + +// ScopedAppName returns an app name prefixed with the tenant-scoped storage namespace. +func (s StorageScope) ScopedAppName(appName string) string { + prefix := s.namespacePrefix() + appName = strings.Trim(strings.TrimSpace(appName), `/\|:`) + if appName == "" { + return strings.TrimSuffix(prefix, "/") + } + if strings.HasPrefix(appName, prefix) { + return appName + } + return prefix + appName +} + +func (s StorageScope) namespacePrefix() string { + namespace := strings.TrimRight(strings.TrimSpace(s.Namespace), `/\|:`) + if namespace == "" { + return "" + } + return namespace + "/" +} + +func (s StorageScope) validateAppName(appName string) error { + if strings.TrimSpace(appName) == "" || strings.TrimSpace(appName) != appName { + return ErrKeyOutsideTenantScope + } + prefix := s.namespacePrefix() + if prefix == "" || !strings.HasPrefix(appName, prefix) { + return ErrKeyOutsideTenantScope + } + if strings.TrimSpace(strings.TrimPrefix(appName, prefix)) == "" { + return ErrKeyOutsideTenantScope + } + return nil +} + +func (s StorageScope) validateSessionKey(key session.Key) error { + if err := key.CheckSessionKey(); err != nil { + return err + } + return s.validateAppName(key.AppName) +} + +func (s StorageScope) validateSessionUserKey(key session.UserKey) error { + if err := key.CheckUserKey(); err != nil { + return err + } + return s.validateAppName(key.AppName) +} + +func (s StorageScope) validateSession(sess *session.Session) error { + if sess == nil { + return session.ErrNilSession + } + return s.validateAppName(sess.AppName) +} + +func (s StorageScope) validateMemoryKey(key memory.Key) error { + if err := key.CheckMemoryKey(); err != nil { + return err + } + return s.validateAppName(key.AppName) +} + +func (s StorageScope) validateMemoryUserKey(key memory.UserKey) error { + if err := key.CheckUserKey(); err != nil { + return err + } + return s.validateAppName(key.AppName) +} + +func (s StorageScope) validateArtifactSessionInfo(info artifact.SessionInfo) error { + if strings.TrimSpace(info.UserID) == "" || strings.TrimSpace(info.SessionID) == "" { + return ErrKeyOutsideTenantScope + } + return s.validateAppName(info.AppName) +} + +type tenantStorageAdapter struct { + router *InMemoryRouter + scope StorageScope +} + +func (a *tenantStorageAdapter) Scope() StorageScope { + return a.scope +} + +func (a *tenantStorageAdapter) Route( + ctx context.Context, + resource platform.BackendMigrationResource, +) (RouteBinding, error) { + return a.router.Route(ctx, a.scope.TenantID, a.scope.ProfileID, resource) +} + +func (a *tenantStorageAdapter) Session(ctx context.Context) (session.Service, error) { + service, err := a.router.Session(ctx, a.scope.TenantID, a.scope.ProfileID) + if err != nil { + return nil, err + } + return &scopedSessionService{Service: service, scope: a.scope}, nil +} + +func (a *tenantStorageAdapter) Summary(ctx context.Context) (SummaryStore, error) { + store, err := a.router.Summary(ctx, a.scope.TenantID, a.scope.ProfileID) + if err != nil { + return nil, err + } + return &scopedSummaryStore{SummaryStore: store, scope: a.scope}, nil +} + +func (a *tenantStorageAdapter) Memory(ctx context.Context) (memory.Service, error) { + service, err := a.router.Memory(ctx, a.scope.TenantID, a.scope.ProfileID) + if err != nil { + return nil, err + } + return &scopedMemoryService{Service: service, scope: a.scope}, nil +} + +func (a *tenantStorageAdapter) Artifact(ctx context.Context) (artifact.Service, error) { + service, err := a.router.Artifact(ctx, a.scope.TenantID, a.scope.ProfileID) + if err != nil { + return nil, err + } + return &scopedArtifactService{Service: service, scope: a.scope}, nil +} + +func (a *tenantStorageAdapter) Knowledge(ctx context.Context) (knowledge.Knowledge, error) { + service, err := a.router.Knowledge(ctx, a.scope.TenantID, a.scope.ProfileID) + if err != nil { + return nil, err + } + return &scopedKnowledge{Knowledge: service, scope: a.scope}, nil +} + +func (a *tenantStorageAdapter) Audit(ctx context.Context) (platform.AuditSink, error) { + sink, err := a.router.Audit(ctx, a.scope.TenantID, a.scope.ProfileID) + if err != nil { + return nil, err + } + return &scopedAuditSink{AuditSink: sink, scope: a.scope}, nil +} + +type scopedSessionService struct { + session.Service + scope StorageScope +} + +func (s *scopedSessionService) CreateSession( + ctx context.Context, + key session.Key, + state session.StateMap, + options ...session.Option, +) (*session.Session, error) { + if err := s.scope.validateSessionKey(key); err != nil { + return nil, err + } + return s.Service.CreateSession(ctx, key, state, options...) +} + +func (s *scopedSessionService) GetSession( + ctx context.Context, + key session.Key, + options ...session.Option, +) (*session.Session, error) { + if err := s.scope.validateSessionKey(key); err != nil { + return nil, err + } + return s.Service.GetSession(ctx, key, options...) +} + +func (s *scopedSessionService) ListSessions( + ctx context.Context, + userKey session.UserKey, + options ...session.Option, +) ([]*session.Session, error) { + if err := s.scope.validateSessionUserKey(userKey); err != nil { + return nil, err + } + return s.Service.ListSessions(ctx, userKey, options...) +} + +func (s *scopedSessionService) DeleteSession( + ctx context.Context, + key session.Key, + options ...session.Option, +) error { + if err := s.scope.validateSessionKey(key); err != nil { + return err + } + return s.Service.DeleteSession(ctx, key, options...) +} + +func (s *scopedSessionService) UpdateAppState( + ctx context.Context, + appName string, + state session.StateMap, +) error { + if err := s.scope.validateAppName(appName); err != nil { + return err + } + return s.Service.UpdateAppState(ctx, appName, state) +} + +func (s *scopedSessionService) DeleteAppState(ctx context.Context, appName string, key string) error { + if err := s.scope.validateAppName(appName); err != nil { + return err + } + return s.Service.DeleteAppState(ctx, appName, key) +} + +func (s *scopedSessionService) ListAppStates(ctx context.Context, appName string) (session.StateMap, error) { + if err := s.scope.validateAppName(appName); err != nil { + return nil, err + } + return s.Service.ListAppStates(ctx, appName) +} + +func (s *scopedSessionService) UpdateUserState( + ctx context.Context, + userKey session.UserKey, + state session.StateMap, +) error { + if err := s.scope.validateSessionUserKey(userKey); err != nil { + return err + } + return s.Service.UpdateUserState(ctx, userKey, state) +} + +func (s *scopedSessionService) ListUserStates( + ctx context.Context, + userKey session.UserKey, +) (session.StateMap, error) { + if err := s.scope.validateSessionUserKey(userKey); err != nil { + return nil, err + } + return s.Service.ListUserStates(ctx, userKey) +} + +func (s *scopedSessionService) DeleteUserState( + ctx context.Context, + userKey session.UserKey, + key string, +) error { + if err := s.scope.validateSessionUserKey(userKey); err != nil { + return err + } + return s.Service.DeleteUserState(ctx, userKey, key) +} + +func (s *scopedSessionService) UpdateSessionState( + ctx context.Context, + key session.Key, + state session.StateMap, +) error { + if err := s.scope.validateSessionKey(key); err != nil { + return err + } + return s.Service.UpdateSessionState(ctx, key, state) +} + +func (s *scopedSessionService) AppendEvent( + ctx context.Context, + sess *session.Session, + event *event.Event, + options ...session.Option, +) error { + if err := s.scope.validateSession(sess); err != nil { + return err + } + return s.Service.AppendEvent(ctx, sess, event, options...) +} + +func (s *scopedSessionService) CreateSessionSummary( + ctx context.Context, + sess *session.Session, + filterKey string, + force bool, +) error { + if err := s.scope.validateSession(sess); err != nil { + return err + } + return s.Service.CreateSessionSummary(ctx, sess, filterKey, force) +} + +func (s *scopedSessionService) EnqueueSummaryJob( + ctx context.Context, + sess *session.Session, + filterKey string, + force bool, +) error { + if err := s.scope.validateSession(sess); err != nil { + return err + } + return s.Service.EnqueueSummaryJob(ctx, sess, filterKey, force) +} + +func (s *scopedSessionService) GetSessionSummaryText( + ctx context.Context, + sess *session.Session, + opts ...session.SummaryOption, +) (string, bool) { + if err := s.scope.validateSession(sess); err != nil { + return "", false + } + return s.Service.GetSessionSummaryText(ctx, sess, opts...) +} + +type scopedSummaryStore struct { + SummaryStore + scope StorageScope +} + +func (s *scopedSummaryStore) CreateSessionSummary( + ctx context.Context, + sess *session.Session, + filterKey string, + force bool, +) error { + if err := s.scope.validateSession(sess); err != nil { + return err + } + return s.SummaryStore.CreateSessionSummary(ctx, sess, filterKey, force) +} + +func (s *scopedSummaryStore) EnqueueSummaryJob( + ctx context.Context, + sess *session.Session, + filterKey string, + force bool, +) error { + if err := s.scope.validateSession(sess); err != nil { + return err + } + return s.SummaryStore.EnqueueSummaryJob(ctx, sess, filterKey, force) +} + +func (s *scopedSummaryStore) GetSessionSummaryText( + ctx context.Context, + sess *session.Session, + opts ...session.SummaryOption, +) (string, bool) { + if err := s.scope.validateSession(sess); err != nil { + return "", false + } + return s.SummaryStore.GetSessionSummaryText(ctx, sess, opts...) +} + +type scopedMemoryService struct { + memory.Service + scope StorageScope +} + +func (s *scopedMemoryService) ReadMemories( + ctx context.Context, + userKey memory.UserKey, + limit int, +) ([]*memory.Entry, error) { + if err := s.scope.validateMemoryUserKey(userKey); err != nil { + return nil, err + } + return s.Service.ReadMemories(ctx, userKey, limit) +} + +func (s *scopedMemoryService) SearchMemories( + ctx context.Context, + userKey memory.UserKey, + query string, + opts ...memory.SearchOption, +) ([]*memory.Entry, error) { + if err := s.scope.validateMemoryUserKey(userKey); err != nil { + return nil, err + } + return s.Service.SearchMemories(ctx, userKey, query, opts...) +} + +func (s *scopedMemoryService) AddMemory( + ctx context.Context, + userKey memory.UserKey, + mem string, + topics []string, + opts ...memory.AddOption, +) error { + if err := s.scope.validateMemoryUserKey(userKey); err != nil { + return err + } + return s.Service.AddMemory(ctx, userKey, mem, topics, opts...) +} + +func (s *scopedMemoryService) UpdateMemory( + ctx context.Context, + memoryKey memory.Key, + mem string, + topics []string, + opts ...memory.UpdateOption, +) error { + if err := s.scope.validateMemoryKey(memoryKey); err != nil { + return err + } + return s.Service.UpdateMemory(ctx, memoryKey, mem, topics, opts...) +} + +func (s *scopedMemoryService) DeleteMemory(ctx context.Context, memoryKey memory.Key) error { + if err := s.scope.validateMemoryKey(memoryKey); err != nil { + return err + } + return s.Service.DeleteMemory(ctx, memoryKey) +} + +func (s *scopedMemoryService) ClearMemories(ctx context.Context, userKey memory.UserKey) error { + if err := s.scope.validateMemoryUserKey(userKey); err != nil { + return err + } + return s.Service.ClearMemories(ctx, userKey) +} + +func (s *scopedMemoryService) EnqueueAutoMemoryJob( + ctx context.Context, + sess *session.Session, +) error { + if err := s.scope.validateSession(sess); err != nil { + return err + } + return s.Service.EnqueueAutoMemoryJob(ctx, sess) +} + +type scopedArtifactService struct { + artifact.Service + scope StorageScope +} + +func (s *scopedArtifactService) SaveArtifact( + ctx context.Context, + sessionInfo artifact.SessionInfo, + filename string, + artifactValue *artifact.Artifact, +) (int, error) { + if err := s.scope.validateArtifactSessionInfo(sessionInfo); err != nil { + return 0, err + } + return s.Service.SaveArtifact(ctx, sessionInfo, filename, artifactValue) +} + +func (s *scopedArtifactService) LoadArtifact( + ctx context.Context, + sessionInfo artifact.SessionInfo, + filename string, + version *int, +) (*artifact.Artifact, error) { + if err := s.scope.validateArtifactSessionInfo(sessionInfo); err != nil { + return nil, err + } + return s.Service.LoadArtifact(ctx, sessionInfo, filename, version) +} + +func (s *scopedArtifactService) ListArtifactKeys( + ctx context.Context, + sessionInfo artifact.SessionInfo, +) ([]string, error) { + if err := s.scope.validateArtifactSessionInfo(sessionInfo); err != nil { + return nil, err + } + return s.Service.ListArtifactKeys(ctx, sessionInfo) +} + +func (s *scopedArtifactService) DeleteArtifact( + ctx context.Context, + sessionInfo artifact.SessionInfo, + filename string, +) error { + if err := s.scope.validateArtifactSessionInfo(sessionInfo); err != nil { + return err + } + return s.Service.DeleteArtifact(ctx, sessionInfo, filename) +} + +func (s *scopedArtifactService) ListVersions( + ctx context.Context, + sessionInfo artifact.SessionInfo, + filename string, +) ([]int, error) { + if err := s.scope.validateArtifactSessionInfo(sessionInfo); err != nil { + return nil, err + } + return s.Service.ListVersions(ctx, sessionInfo, filename) +} + +type scopedKnowledge struct { + knowledge.Knowledge + scope StorageScope +} + +func (s *scopedKnowledge) Search( + ctx context.Context, + req *knowledge.SearchRequest, +) (*knowledge.SearchResult, error) { + if req == nil { + return nil, errors.New("knowledge search request is required") + } + scopedReq := *req + if req.SearchFilter == nil { + scopedReq.SearchFilter = &knowledge.SearchFilter{} + } else { + filter := *req.SearchFilter + scopedReq.SearchFilter = &filter + } + if scopedReq.SearchFilter.Metadata == nil { + scopedReq.SearchFilter.Metadata = make(map[string]any, 1) + } else { + metadata := make(map[string]any, len(scopedReq.SearchFilter.Metadata)+1) + for key, value := range scopedReq.SearchFilter.Metadata { + metadata[key] = value + } + scopedReq.SearchFilter.Metadata = metadata + } + if tenantID, ok := scopedReq.SearchFilter.Metadata[tenantMetadataKey]; ok && tenantID != s.scope.TenantID { + return nil, ErrKeyOutsideTenantScope + } + scopedReq.SearchFilter.Metadata[tenantMetadataKey] = s.scope.TenantID + return s.Knowledge.Search(ctx, &scopedReq) +} + +type scopedAuditSink struct { + platform.AuditSink + scope StorageScope +} + +func (s *scopedAuditSink) WriteAudit(ctx context.Context, record platform.AuditRecord) error { + if strings.TrimSpace(record.TenantID) != s.scope.TenantID { + return ErrKeyOutsideTenantScope + } + return s.AuditSink.WriteAudit(ctx, record) +} diff --git a/platform/storagerouter/errors.go b/platform/storagerouter/errors.go index 66e9371d07..45dbb2a4ea 100644 --- a/platform/storagerouter/errors.go +++ b/platform/storagerouter/errors.go @@ -21,4 +21,6 @@ var ( ErrBackendIDRequired = errors.New("storage router backend id required") // ErrBackendTenantMismatch indicates that a registered backend belongs to another tenant. ErrBackendTenantMismatch = errors.New("storage router backend tenant mismatch") + // ErrKeyOutsideTenantScope indicates that a storage key is not scoped to the tenant namespace. + ErrKeyOutsideTenantScope = errors.New("storage router key outside tenant scope") ) diff --git a/platform/storagerouter/router.go b/platform/storagerouter/router.go index b5643bd285..e317520066 100644 --- a/platform/storagerouter/router.go +++ b/platform/storagerouter/router.go @@ -10,6 +10,7 @@ package storagerouter import ( "context" + "fmt" "strings" "sync" @@ -25,16 +26,38 @@ type BackendSet struct { TenantID string BackendID string Session session.Service + Summary SummaryStore Memory memory.Service Artifact artifact.Service Knowledge knowledge.Knowledge Audit platform.AuditSink } +// SummaryStore is the summary-specific storage surface selected by SummaryBackend. +type SummaryStore interface { + CreateSessionSummary(ctx context.Context, sess *session.Session, filterKey string, force bool) error + EnqueueSummaryJob(ctx context.Context, sess *session.Session, filterKey string, force bool) error + GetSessionSummaryText(ctx context.Context, sess *session.Session, opts ...session.SummaryOption) (string, bool) +} + +// RouteBinding describes the concrete backend route selected for one resource. +type RouteBinding struct { + TenantID string + ProfileID string + Resource platform.BackendMigrationResource + BackendID string + Namespace string + MigrationMode platform.StorageMigrationMode + IsMigrating bool +} + // Router resolves tenant/app storage services from platform storage profiles. type Router interface { Profile(ctx context.Context, tenantID string, profileID string) (platform.StorageProfile, error) + Adapter(ctx context.Context, tenantID string, profileID string) (StorageAdapter, error) + Route(ctx context.Context, tenantID string, profileID string, resource platform.BackendMigrationResource) (RouteBinding, error) Session(ctx context.Context, tenantID string, profileID string) (session.Service, error) + Summary(ctx context.Context, tenantID string, profileID string) (SummaryStore, error) Memory(ctx context.Context, tenantID string, profileID string) (memory.Service, error) Artifact(ctx context.Context, tenantID string, profileID string) (artifact.Service, error) Knowledge(ctx context.Context, tenantID string, profileID string) (knowledge.Knowledge, error) @@ -120,6 +143,72 @@ func (r *InMemoryRouter) Profile( return profile, nil } +// Adapter returns a tenant/profile-bound storage adapter. +func (r *InMemoryRouter) Adapter( + ctx context.Context, + tenantID string, + profileID string, +) (StorageAdapter, error) { + profile, err := r.Profile(ctx, tenantID, profileID) + if err != nil { + return nil, err + } + return &tenantStorageAdapter{ + router: r, + scope: StorageScope{ + TenantID: profile.TenantID, + ProfileID: profile.ProfileID, + Namespace: profile.Namespace, + }, + }, nil +} + +// Route resolves the concrete tenant-scoped backend route for one resource. +func (r *InMemoryRouter) Route( + ctx context.Context, + tenantID string, + profileID string, + resource platform.BackendMigrationResource, +) (RouteBinding, error) { + profile, err := r.Profile(ctx, tenantID, profileID) + if err != nil { + return RouteBinding{}, err + } + kind, err := resourceKindFor(resource) + if err != nil { + return RouteBinding{}, err + } + backendID := backendIDFor(profile, kind) + if strings.TrimSpace(backendID) == "" { + return RouteBinding{}, ErrBackendNotFound + } + mode, err := platform.NormalizeStorageMigrationMode(profile.MigrationMode) + if err != nil { + return RouteBinding{}, err + } + r.mu.RLock() + backend, ok := r.backends[backendKey{tenantID: tenantID, backendID: backendID}] + r.mu.RUnlock() + if !ok { + return RouteBinding{}, ErrBackendNotFound + } + if backend.TenantID != tenantID { + return RouteBinding{}, ErrBackendTenantMismatch + } + if !backendHasResource(backend, kind) { + return RouteBinding{}, ErrBackendNotFound + } + return RouteBinding{ + TenantID: profile.TenantID, + ProfileID: profile.ProfileID, + Resource: resource, + BackendID: strings.TrimSpace(backendID), + Namespace: profile.Namespace, + MigrationMode: mode, + IsMigrating: platform.IsActiveStorageMigrationMode(mode), + }, nil +} + // Session resolves the session service selected by a tenant storage profile. func (r *InMemoryRouter) Session( ctx context.Context, @@ -136,6 +225,22 @@ func (r *InMemoryRouter) Session( return backend.Session, nil } +// Summary resolves the summary store selected by a tenant storage profile. +func (r *InMemoryRouter) Summary( + ctx context.Context, + tenantID string, + profileID string, +) (SummaryStore, error) { + backend, err := r.backend(ctx, tenantID, profileID, resourceSummary) + if err != nil { + return nil, err + } + if backend.Summary == nil { + return nil, ErrBackendNotFound + } + return backend.Summary, nil +} + // Memory resolves the memory service selected by a tenant storage profile. func (r *InMemoryRouter) Memory( ctx context.Context, @@ -204,6 +309,7 @@ type resourceKind string const ( resourceSession resourceKind = "session" + resourceSummary resourceKind = "summary" resourceMemory resourceKind = "memory" resourceArtifact resourceKind = "artifact" resourceKnowledge resourceKind = "knowledge" @@ -241,6 +347,8 @@ func backendIDFor(profile platform.StorageProfile, kind resourceKind) string { switch kind { case resourceSession: return strings.TrimSpace(profile.SessionBackend) + case resourceSummary: + return strings.TrimSpace(profile.SummaryBackend) case resourceMemory: return strings.TrimSpace(profile.MemoryBackend) case resourceArtifact: @@ -253,3 +361,22 @@ func backendIDFor(profile platform.StorageProfile, kind resourceKind) string { return "" } } + +func resourceKindFor(resource platform.BackendMigrationResource) (resourceKind, error) { + switch resource { + case platform.BackendMigrationResourceSession: + return resourceSession, nil + case platform.BackendMigrationResourceSummary: + return resourceSummary, nil + case platform.BackendMigrationResourceMemory: + return resourceMemory, nil + case platform.BackendMigrationResourceArtifact: + return resourceArtifact, nil + case platform.BackendMigrationResourceKnowledge: + return resourceKnowledge, nil + case platform.BackendMigrationResourceAudit: + return resourceAudit, nil + default: + return "", fmt.Errorf("unsupported storage resource %q", resource) + } +} diff --git a/platform/storagerouter/router_test.go b/platform/storagerouter/router_test.go index 1cba3b6c7a..289954203d 100644 --- a/platform/storagerouter/router_test.go +++ b/platform/storagerouter/router_test.go @@ -18,8 +18,10 @@ import ( artifactmemory "trpc.group/trpc-go/trpc-agent-go/artifact/inmemory" "trpc.group/trpc-go/trpc-agent-go/knowledge" + "trpc.group/trpc-go/trpc-agent-go/memory" memoryinmemory "trpc.group/trpc-go/trpc-agent-go/memory/inmemory" "trpc.group/trpc-go/trpc-agent-go/platform" + "trpc.group/trpc-go/trpc-agent-go/session" sessioninmemory "trpc.group/trpc-go/trpc-agent-go/session/inmemory" ) @@ -29,11 +31,14 @@ func TestRouterResolvesTenantStorageServices(t *testing.T) { ctx := context.Background() router := NewInMemoryRouter() sessionSvc := sessioninmemory.NewSessionService() + summarySvc := sessioninmemory.NewSessionService() memorySvc := memoryinmemory.NewMemoryService() artifactSvc := artifactmemory.NewService() knowledgeSvc := &stubKnowledge{} auditSink := platform.NewInMemoryAuditSink() - require.NoError(t, router.RegisterProfile(profile("tenant-a", "profile-a", "hot"))) + p := profile("tenant-a", "profile-a", "hot") + p.SummaryBackend = "summary-hot" + require.NoError(t, router.RegisterProfile(p)) require.NoError(t, router.RegisterBackend(BackendSet{ TenantID: "tenant-a", BackendID: "hot", @@ -43,9 +48,16 @@ func TestRouterResolvesTenantStorageServices(t *testing.T) { Knowledge: knowledgeSvc, Audit: auditSink, })) + require.NoError(t, router.RegisterBackend(BackendSet{ + TenantID: "tenant-a", + BackendID: "summary-hot", + Summary: summarySvc, + })) gotSession, err := router.Session(ctx, "tenant-a", "profile-a") require.NoError(t, err) + gotSummary, err := router.Summary(ctx, "tenant-a", "profile-a") + require.NoError(t, err) gotMemory, err := router.Memory(ctx, "tenant-a", "profile-a") require.NoError(t, err) gotArtifact, err := router.Artifact(ctx, "tenant-a", "profile-a") @@ -56,6 +68,7 @@ func TestRouterResolvesTenantStorageServices(t *testing.T) { require.NoError(t, err) assert.Same(t, sessionSvc, gotSession) + assert.Same(t, summarySvc, gotSummary) assert.Same(t, memorySvc, gotMemory) assert.Same(t, artifactSvc, gotArtifact) assert.Same(t, knowledgeSvc, gotKnowledge) @@ -109,6 +122,144 @@ func TestRouterRejectsMissingResourceService(t *testing.T) { require.ErrorIs(t, err, ErrBackendNotFound) } +func TestRouterRouteReturnsTenantScopedMetadata(t *testing.T) { + ctx := context.Background() + router := NewInMemoryRouter() + p := profile("tenant-a", "profile-a", "hot") + p.SessionBackend = "session-hot" + p.MigrationMode = string(platform.StorageMigrationModeDualWrite) + require.NoError(t, router.RegisterProfile(p)) + require.NoError(t, router.RegisterBackend(BackendSet{ + TenantID: "tenant-a", + BackendID: "session-hot", + Session: sessioninmemory.NewSessionService(), + })) + + route, err := router.Route(ctx, "tenant-a", "profile-a", platform.BackendMigrationResourceSession) + require.NoError(t, err) + + assert.Equal(t, "tenant-a", route.TenantID) + assert.Equal(t, "profile-a", route.ProfileID) + assert.Equal(t, platform.BackendMigrationResourceSession, route.Resource) + assert.Equal(t, "session-hot", route.BackendID) + assert.Equal(t, "tenant/tenant-a", route.Namespace) + assert.Equal(t, platform.StorageMigrationModeDualWrite, route.MigrationMode) + assert.True(t, route.IsMigrating) +} + +func TestRouterAdapterReturnsScopedStores(t *testing.T) { + ctx := context.Background() + router := NewInMemoryRouter() + sessionSvc := sessioninmemory.NewSessionService() + p := profile("tenant-a", "profile-a", "hot") + p.SessionBackend = "session-hot" + require.NoError(t, router.RegisterProfile(p)) + require.NoError(t, router.RegisterBackend(BackendSet{ + TenantID: "tenant-a", + BackendID: "session-hot", + Session: sessionSvc, + })) + + adapter, err := router.Adapter(ctx, "tenant-a", "profile-a") + require.NoError(t, err) + sessionStore, err := adapter.Session(ctx) + require.NoError(t, err) + scopedApp := adapter.Scope().ScopedAppName("app-a") + + created, err := sessionStore.CreateSession(ctx, session.Key{ + AppName: scopedApp, + UserID: "user-a", + SessionID: "session-a", + }, nil) + require.NoError(t, err) + assert.Equal(t, "tenant/tenant-a/app-a", created.AppName) + + route, err := adapter.Route(ctx, platform.BackendMigrationResourceSession) + require.NoError(t, err) + assert.Equal(t, "tenant-a", route.TenantID) + assert.Equal(t, "profile-a", route.ProfileID) + assert.Equal(t, "tenant/tenant-a", route.Namespace) +} + +func TestRouterAdapterRejectsUnscopedKeys(t *testing.T) { + ctx := context.Background() + router := NewInMemoryRouter() + p := profile("tenant-a", "profile-a", "hot") + require.NoError(t, router.RegisterProfile(p)) + require.NoError(t, router.RegisterBackend(BackendSet{ + TenantID: "tenant-a", + BackendID: "hot", + Session: sessioninmemory.NewSessionService(), + Memory: memoryinmemory.NewMemoryService(), + })) + + adapter, err := router.Adapter(ctx, "tenant-a", "profile-a") + require.NoError(t, err) + sessionStore, err := adapter.Session(ctx) + require.NoError(t, err) + memoryStore, err := adapter.Memory(ctx) + require.NoError(t, err) + + _, err = sessionStore.CreateSession(ctx, session.Key{ + AppName: "app-a", + UserID: "user-a", + SessionID: "session-a", + }, nil) + require.ErrorIs(t, err, ErrKeyOutsideTenantScope) + + err = memoryStore.AddMemory(ctx, memory.UserKey{ + AppName: "app-a", + UserID: "user-a", + }, "prefers tea", []string{"preference"}) + require.ErrorIs(t, err, ErrKeyOutsideTenantScope) +} + +func TestRouterAdapterScopesKnowledgeQueries(t *testing.T) { + ctx := context.Background() + router := NewInMemoryRouter() + knowledgeSvc := &capturingKnowledge{} + p := profile("tenant-a", "profile-a", "hot") + require.NoError(t, router.RegisterProfile(p)) + require.NoError(t, router.RegisterBackend(BackendSet{ + TenantID: "tenant-a", + BackendID: "hot", + Knowledge: knowledgeSvc, + })) + adapter, err := router.Adapter(ctx, "tenant-a", "profile-a") + require.NoError(t, err) + knowledgeStore, err := adapter.Knowledge(ctx) + require.NoError(t, err) + req := &knowledge.SearchRequest{ + Query: "deployment", + SearchFilter: &knowledge.SearchFilter{ + Metadata: map[string]any{"category": "runbook"}, + }, + } + + _, err = knowledgeStore.Search(ctx, req) + require.NoError(t, err) + + assert.Equal(t, "tenant-a", knowledgeSvc.last.SearchFilter.Metadata["tenant_id"]) + assert.Equal(t, "runbook", knowledgeSvc.last.SearchFilter.Metadata["category"]) + assert.NotContains(t, req.SearchFilter.Metadata, "tenant_id") + + _, err = knowledgeStore.Search(ctx, &knowledge.SearchRequest{ + Query: "deployment", + SearchFilter: &knowledge.SearchFilter{ + Metadata: map[string]any{"tenant_id": "tenant-b"}, + }, + }) + require.ErrorIs(t, err, ErrKeyOutsideTenantScope) + + _, err = knowledgeStore.Search(ctx, &knowledge.SearchRequest{ + Query: "deployment", + SearchFilter: &knowledge.SearchFilter{ + Metadata: map[string]any{"tenant_id": []string{"tenant-a"}}, + }, + }) + require.ErrorIs(t, err, ErrKeyOutsideTenantScope) +} + func TestRegisterProfileValidatesSecretRefs(t *testing.T) { router := NewInMemoryRouter() p := profile("tenant-a", "profile-a", "hot") @@ -135,6 +286,7 @@ func profile(tenantID string, profileID string, backendID string) platform.Stora TenantID: tenantID, ProfileID: profileID, SessionBackend: backendID, + SummaryBackend: backendID, MemoryBackend: backendID, ArtifactBackend: backendID, KnowledgeBackend: backendID, @@ -155,3 +307,18 @@ func (s *stubKnowledge) Search( } return &knowledge.SearchResult{}, nil } + +type capturingKnowledge struct { + last knowledge.SearchRequest +} + +func (s *capturingKnowledge) Search( + ctx context.Context, + req *knowledge.SearchRequest, +) (*knowledge.SearchResult, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + s.last = *req + return &knowledge.SearchResult{}, nil +} diff --git a/platform/storagerouter/status.go b/platform/storagerouter/status.go index 12c9d0c415..a1c369c467 100644 --- a/platform/storagerouter/status.go +++ b/platform/storagerouter/status.go @@ -83,6 +83,7 @@ func (r *InMemoryRouter) Status( resource platform.BackendMigrationResource }{ {kind: resourceSession, resource: platform.BackendMigrationResourceSession}, + {kind: resourceSummary, resource: platform.BackendMigrationResourceSummary}, {kind: resourceMemory, resource: platform.BackendMigrationResourceMemory}, {kind: resourceArtifact, resource: platform.BackendMigrationResourceArtifact}, {kind: resourceKnowledge, resource: platform.BackendMigrationResourceKnowledge}, @@ -151,6 +152,8 @@ func backendHasResource(backend BackendSet, kind resourceKind) bool { switch kind { case resourceSession: return backend.Session != nil + case resourceSummary: + return backend.Summary != nil case resourceMemory: return backend.Memory != nil case resourceArtifact: diff --git a/platform/storagerouter/status_test.go b/platform/storagerouter/status_test.go index 0c706e6673..17b112eb78 100644 --- a/platform/storagerouter/status_test.go +++ b/platform/storagerouter/status_test.go @@ -32,6 +32,7 @@ func TestRouterStatusReportsAllResourcesReady(t *testing.T) { TenantID: "tenant-a", BackendID: "hot", Session: sessioninmemory.NewSessionService(), + Summary: sessioninmemory.NewSessionService(), Memory: memoryinmemory.NewMemoryService(), Artifact: artifactmemory.NewService(), Knowledge: &stubKnowledge{}, @@ -45,9 +46,9 @@ func TestRouterStatusReportsAllResourcesReady(t *testing.T) { assert.Equal(t, "profile-a", summary.ProfileID) assert.Equal(t, platform.StorageMigrationModeDualWrite, summary.MigrationMode) assert.True(t, summary.IsMigrating) - assert.Equal(t, 5, summary.ReadyCount) + assert.Equal(t, 6, summary.ReadyCount) assert.Equal(t, 0, summary.MissingCount) - require.Len(t, summary.Resources, 5) + require.Len(t, summary.Resources, 6) for _, resource := range summary.Resources { assert.Equal(t, "hot", resource.BackendID) assert.Equal(t, ResourceStatusReady, resource.Status) @@ -74,8 +75,9 @@ func TestRouterStatusReportsMissingBackendAndService(t *testing.T) { assert.False(t, summary.IsMigrating) assert.Equal(t, 2, summary.ReadyCount) - assert.Equal(t, 3, summary.MissingCount) + assert.Equal(t, 4, summary.MissingCount) assertResourceStatus(t, summary, platform.BackendMigrationResourceSession, "hot", ResourceStatusReady) + assertResourceStatus(t, summary, platform.BackendMigrationResourceSummary, "hot", ResourceStatusServiceMissing) assertResourceStatus(t, summary, platform.BackendMigrationResourceMemory, "missing", ResourceStatusBackendMissing) assertResourceStatus(t, summary, platform.BackendMigrationResourceArtifact, "", ResourceStatusBackendMissing) assertResourceStatus(t, summary, platform.BackendMigrationResourceKnowledge, "hot", ResourceStatusServiceMissing) diff --git a/platform/validation.go b/platform/validation.go index 3804b6e5cb..ed9610639a 100644 --- a/platform/validation.go +++ b/platform/validation.go @@ -298,6 +298,9 @@ func (p StorageProfile) Validate() error { if err := validateRoutingIdentifier("profile_id", p.ProfileID, fmt.Errorf("profile_id is required")); err != nil { return err } + if err := validateStorageNamespace(p.TenantID, p.Namespace); err != nil { + return err + } if err := validateSecretReference("dsn_ref", p.DSNRef); err != nil { return err } @@ -307,6 +310,35 @@ func (p StorageProfile) Validate() error { return nil } +func validateStorageNamespace(tenantID, namespace string) error { + if err := validateRoutingIdentifier("namespace", namespace, fmt.Errorf("namespace is required")); err != nil { + return err + } + if err := validateAuditRedactedText("namespace", namespace); err != nil { + return err + } + if !namespaceContainsSegment(namespace, tenantID) { + return fmt.Errorf("namespace must include tenant_id") + } + return nil +} + +func namespaceContainsSegment(namespace, tenantID string) bool { + for _, segment := range strings.FieldsFunc(namespace, func(r rune) bool { + switch r { + case '/', '\\', ':', '|': + return true + default: + return false + } + }) { + if segment == tenantID { + return true + } + } + return false +} + // Validate checks that audit retention and sampling policy is safe to use. func (p AuditPolicy) Validate() error { if err := validateRoutingIdentifier("tenant_id", p.TenantID, ErrTenantIDRequired); err != nil { From 1b0147e99efb52c4abe6b0f679fd59c4f5f53c66 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Fri, 10 Jul 2026 15:32:36 +0800 Subject: [PATCH 54/95] memory: add tenant-scoped memory knowledge abstraction --- platform/memoryknowledge/doc.go | 16 ++ platform/memoryknowledge/errors.go | 24 ++ platform/memoryknowledge/service.go | 324 +++++++++++++++++++++++ platform/memoryknowledge/service_test.go | 228 ++++++++++++++++ 4 files changed, 592 insertions(+) create mode 100644 platform/memoryknowledge/doc.go create mode 100644 platform/memoryknowledge/errors.go create mode 100644 platform/memoryknowledge/service.go create mode 100644 platform/memoryknowledge/service_test.go diff --git a/platform/memoryknowledge/doc.go b/platform/memoryknowledge/doc.go new file mode 100644 index 0000000000..0d4fb1218f --- /dev/null +++ b/platform/memoryknowledge/doc.go @@ -0,0 +1,16 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +// Package memoryknowledge provides tenant-scoped memory and knowledge facades. +// +// The facade keeps the platform boundary in front of concrete backends: +// memory writes are accepted with eventual consistency because vector indexing +// or remote memory providers may lag, while knowledge reads always inject +// tenant_id and internal_user_id filters. SearchRequest MaxResults and MinScore +// remain caller-controlled latency, recall, and cost knobs. +package memoryknowledge diff --git a/platform/memoryknowledge/errors.go b/platform/memoryknowledge/errors.go new file mode 100644 index 0000000000..678b58eabe --- /dev/null +++ b/platform/memoryknowledge/errors.go @@ -0,0 +1,24 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package memoryknowledge + +import "errors" + +var ( + // ErrMemoryBackendRequired indicates that the facade has no memory backend. + ErrMemoryBackendRequired = errors.New("memory backend is required") + // ErrKnowledgeBackendRequired indicates that the facade has no knowledge backend. + ErrKnowledgeBackendRequired = errors.New("knowledge backend is required") + // ErrInternalUserIDRequired indicates that retrieval lacks the internal user boundary. + ErrInternalUserIDRequired = errors.New("internal_user_id is required") + // ErrNamespaceRequired indicates that the storage namespace is missing. + ErrNamespaceRequired = errors.New("namespace is required") + // ErrFilterOutsideScope indicates that caller-supplied retrieval filters escape the scope. + ErrFilterOutsideScope = errors.New("memoryknowledge filter outside scope") +) diff --git a/platform/memoryknowledge/service.go b/platform/memoryknowledge/service.go new file mode 100644 index 0000000000..35bb58805e --- /dev/null +++ b/platform/memoryknowledge/service.go @@ -0,0 +1,324 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package memoryknowledge + +import ( + "context" + "fmt" + "strings" + "unicode" + + "trpc.group/trpc-go/trpc-agent-go/knowledge" + "trpc.group/trpc-go/trpc-agent-go/memory" + "trpc.group/trpc-go/trpc-agent-go/platform" +) + +const ( + // MetadataTenantID is the metadata key used to enforce tenant search scope. + MetadataTenantID = "tenant_id" + // MetadataAppID is the metadata key used to enforce app search scope. + MetadataAppID = "app_id" + // MetadataInternalUserID is the metadata key used to enforce internal user search scope. + MetadataInternalUserID = "internal_user_id" + // MetadataUserIDHash is the metadata key used to carry privacy-safe user identity. + MetadataUserIDHash = "user_id_hash" +) + +// Consistency describes when a memory write should become visible to retrieval. +type Consistency string + +const ( + // ConsistencyEventual means the write was accepted but downstream vector/index + // visibility may lag behind durable storage. + ConsistencyEventual Consistency = "eventual" +) + +// MemoryBackend is the long-term memory surface used by the facade. +type MemoryBackend interface { + memory.Service +} + +// KnowledgeBackend is the knowledge retrieval surface used by the facade. +// Implementations must honor SearchFilter.Metadata as mandatory filters. +type KnowledgeBackend interface { + knowledge.Knowledge +} + +// ServiceConfig wires concrete memory and knowledge backends. Tests can provide +// in-memory or mock backends, while production callers can wire vector stores. +type ServiceConfig struct { + Memory MemoryBackend + Knowledge KnowledgeBackend +} + +// Scope carries the tenant and privacy-safe user boundary for memory and RAG. +type Scope struct { + TenantID string + AppID string + InternalUserID string + UserIDHash string + Namespace string +} + +// Validate checks that the scope is strong enough for tenant/user isolation. +func (s Scope) Validate() error { + if err := validateIdentifier(MetadataTenantID, s.TenantID, platform.ErrTenantIDRequired); err != nil { + return err + } + if err := validateIdentifier(MetadataAppID, s.AppID, platform.ErrAppIDRequired); err != nil { + return err + } + if err := validateIdentifier(MetadataInternalUserID, s.InternalUserID, ErrInternalUserIDRequired); err != nil { + return err + } + if err := validateIdentifier("namespace", s.Namespace, ErrNamespaceRequired); err != nil { + return err + } + if s.UserIDHash != "" { + if err := validateIdentifier(MetadataUserIDHash, s.UserIDHash, nil); err != nil { + return err + } + } + if !namespaceContainsSegment(s.Namespace, s.TenantID) { + return fmt.Errorf("namespace must include tenant_id") + } + return nil +} + +// ScopedAppName returns the memory app key inside the tenant namespace. +func (s Scope) ScopedAppName() string { + namespace := strings.TrimRight(strings.TrimSpace(s.Namespace), `/\|:`) + appID := strings.Trim(strings.TrimSpace(s.AppID), `/\|:`) + if namespace == "" { + return appID + } + if appID == "" { + return namespace + } + return namespace + "/" + appID +} + +func (s Scope) memoryUserKey() memory.UserKey { + return memory.UserKey{ + AppName: s.ScopedAppName(), + UserID: s.InternalUserID, + } +} + +// MemoryWriteRequest writes one long-term memory in a scoped backend. +type MemoryWriteRequest struct { + Scope Scope + Memory string + Topics []string + Metadata *memory.Metadata +} + +// MemoryWriteReceipt confirms acceptance without promising immediate retrieval visibility. +type MemoryWriteReceipt struct { + TenantID string + AppID string + InternalUserID string + UserIDHash string + AppName string + Accepted bool + Consistency Consistency +} + +// KnowledgeSearchRequest wraps a knowledge request with mandatory tenant/user scope. +type KnowledgeSearchRequest struct { + Scope Scope + Request *knowledge.SearchRequest +} + +// Service enforces tenant/internal-user scope across memory writes and retrieval. +type Service struct { + memory MemoryBackend + knowledge KnowledgeBackend +} + +// New creates a scoped memory and knowledge service facade. +func New(config ServiceConfig) (*Service, error) { + if config.Memory == nil { + return nil, ErrMemoryBackendRequired + } + if config.Knowledge == nil { + return nil, ErrKnowledgeBackendRequired + } + return &Service{ + memory: config.Memory, + knowledge: config.Knowledge, + }, nil +} + +// AddMemory accepts a scoped memory write. The receipt is intentionally eventual: +// callers should not assume vector/search visibility before a later retrieval cycle. +func (s *Service) AddMemory( + ctx context.Context, + req MemoryWriteRequest, +) (MemoryWriteReceipt, error) { + if err := ctx.Err(); err != nil { + return MemoryWriteReceipt{}, err + } + if err := req.Scope.Validate(); err != nil { + return MemoryWriteReceipt{}, err + } + opts := make([]memory.AddOption, 0, 1) + if req.Metadata != nil { + opts = append(opts, memory.WithMetadata(req.Metadata)) + } + topics := append([]string(nil), req.Topics...) + if err := s.memory.AddMemory(ctx, req.Scope.memoryUserKey(), req.Memory, topics, opts...); err != nil { + return MemoryWriteReceipt{}, err + } + return MemoryWriteReceipt{ + TenantID: req.Scope.TenantID, + AppID: req.Scope.AppID, + InternalUserID: req.Scope.InternalUserID, + UserIDHash: req.Scope.UserIDHash, + AppName: req.Scope.ScopedAppName(), + Accepted: true, + Consistency: ConsistencyEventual, + }, nil +} + +// ReadMemories reads memories inside the tenant/internal-user boundary. +func (s *Service) ReadMemories( + ctx context.Context, + scope Scope, + limit int, +) ([]*memory.Entry, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + if err := scope.Validate(); err != nil { + return nil, err + } + return s.memory.ReadMemories(ctx, scope.memoryUserKey(), limit) +} + +// SearchMemories searches memories inside the tenant/internal-user boundary. +func (s *Service) SearchMemories( + ctx context.Context, + scope Scope, + query string, + opts ...memory.SearchOption, +) ([]*memory.Entry, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + if err := scope.Validate(); err != nil { + return nil, err + } + return s.memory.SearchMemories(ctx, scope.memoryUserKey(), query, opts...) +} + +// SearchKnowledge injects tenant/internal-user filters into a cloned request. +// MaxResults and MinScore remain caller-controlled knobs for latency, cost, and +// recall tradeoffs; the scope filters are mandatory regardless of those choices. +func (s *Service) SearchKnowledge( + ctx context.Context, + req KnowledgeSearchRequest, +) (*knowledge.SearchResult, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + if err := req.Scope.Validate(); err != nil { + return nil, err + } + if req.Request == nil { + return nil, fmt.Errorf("knowledge search request is required") + } + scopedReq := cloneKnowledgeRequest(req.Request) + if strings.TrimSpace(scopedReq.UserID) != "" && scopedReq.UserID != req.Scope.InternalUserID { + return nil, ErrFilterOutsideScope + } + scopedReq.UserID = req.Scope.InternalUserID + metadata := scopedReq.SearchFilter.Metadata + for key, value := range map[string]string{ + MetadataTenantID: req.Scope.TenantID, + MetadataAppID: req.Scope.AppID, + MetadataInternalUserID: req.Scope.InternalUserID, + } { + if err := enforceMetadata(metadata, key, value); err != nil { + return nil, err + } + } + if req.Scope.UserIDHash != "" { + if err := enforceMetadata(metadata, MetadataUserIDHash, req.Scope.UserIDHash); err != nil { + return nil, err + } + } + return s.knowledge.Search(ctx, &scopedReq) +} + +func cloneKnowledgeRequest(req *knowledge.SearchRequest) knowledge.SearchRequest { + scopedReq := *req + if req.SearchFilter == nil { + scopedReq.SearchFilter = &knowledge.SearchFilter{ + Metadata: make(map[string]any, 4), + } + return scopedReq + } + filter := *req.SearchFilter + if req.SearchFilter.DocumentIDs != nil { + filter.DocumentIDs = append([]string(nil), req.SearchFilter.DocumentIDs...) + } + filter.Metadata = make(map[string]any, len(req.SearchFilter.Metadata)+4) + for key, value := range req.SearchFilter.Metadata { + filter.Metadata[key] = value + } + scopedReq.SearchFilter = &filter + return scopedReq +} + +func enforceMetadata(metadata map[string]any, key string, value string) error { + if existing, ok := metadata[key]; ok { + existingText, ok := existing.(string) + if !ok || existingText != value { + return ErrFilterOutsideScope + } + } + metadata[key] = value + return nil +} + +func validateIdentifier(field string, value string, requiredErr error) error { + trimmed := strings.TrimSpace(value) + if trimmed == "" { + if requiredErr == nil { + return fmt.Errorf("%s must not be blank", field) + } + return requiredErr + } + if trimmed != value { + return fmt.Errorf("%s must not contain leading or trailing whitespace", field) + } + for _, r := range value { + if unicode.IsControl(r) { + return fmt.Errorf("%s must not contain control characters", field) + } + } + return nil +} + +func namespaceContainsSegment(namespace, tenantID string) bool { + for _, segment := range strings.FieldsFunc(namespace, func(r rune) bool { + switch r { + case '/', '\\', ':', '|': + return true + default: + return false + } + }) { + if segment == tenantID { + return true + } + } + return false +} diff --git a/platform/memoryknowledge/service_test.go b/platform/memoryknowledge/service_test.go new file mode 100644 index 0000000000..861028dc98 --- /dev/null +++ b/platform/memoryknowledge/service_test.go @@ -0,0 +1,228 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package memoryknowledge + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "trpc.group/trpc-go/trpc-agent-go/knowledge" + memoryinmemory "trpc.group/trpc-go/trpc-agent-go/memory/inmemory" +) + +func TestServiceAcceptsEventualMemoryWriteAndScopesByTenantUser(t *testing.T) { + ctx := context.Background() + service, err := New(ServiceConfig{ + Memory: memoryinmemory.NewMemoryService(), + Knowledge: &capturingKnowledge{}, + }) + require.NoError(t, err) + scope := Scope{ + TenantID: "tenant-a", + AppID: "app-a", + InternalUserID: "internal-user-a", + UserIDHash: "hash-a", + Namespace: "tenant/tenant-a", + } + + receipt, err := service.AddMemory(ctx, MemoryWriteRequest{ + Scope: scope, + Memory: "Prefers concise deployment runbooks.", + Topics: []string{"preference", "runbook"}, + }) + require.NoError(t, err) + + assert.True(t, receipt.Accepted) + assert.Equal(t, ConsistencyEventual, receipt.Consistency) + assert.Equal(t, "tenant/tenant-a/app-a", receipt.AppName) + assert.Equal(t, "internal-user-a", receipt.InternalUserID) + assert.Equal(t, "hash-a", receipt.UserIDHash) + + entries, err := service.SearchMemories(ctx, scope, "deployment") + require.NoError(t, err) + require.Len(t, entries, 1) + assert.Equal(t, "tenant/tenant-a/app-a", entries[0].AppName) + assert.Equal(t, "internal-user-a", entries[0].UserID) + + otherTenant := scope + otherTenant.TenantID = "tenant-b" + otherTenant.Namespace = "tenant/tenant-b" + entries, err = service.SearchMemories(ctx, otherTenant, "deployment") + require.NoError(t, err) + assert.Empty(t, entries) + + otherUser := scope + otherUser.InternalUserID = "internal-user-b" + entries, err = service.SearchMemories(ctx, otherUser, "deployment") + require.NoError(t, err) + assert.Empty(t, entries) +} + +func TestServiceRejectsMemoryReadsWithoutInternalUserScope(t *testing.T) { + ctx := context.Background() + service, err := New(ServiceConfig{ + Memory: memoryinmemory.NewMemoryService(), + Knowledge: &capturingKnowledge{}, + }) + require.NoError(t, err) + + _, err = service.ReadMemories(ctx, Scope{ + TenantID: "tenant-a", + AppID: "app-a", + Namespace: "tenant/tenant-a", + }, 10) + + require.ErrorIs(t, err, ErrInternalUserIDRequired) +} + +func TestServiceRejectsUnsafeUserIDHashScope(t *testing.T) { + ctx := context.Background() + service, err := New(ServiceConfig{ + Memory: memoryinmemory.NewMemoryService(), + Knowledge: &capturingKnowledge{}, + }) + require.NoError(t, err) + + _, err = service.SearchKnowledge(ctx, KnowledgeSearchRequest{ + Scope: Scope{ + TenantID: "tenant-a", + AppID: "app-a", + InternalUserID: "internal-user-a", + UserIDHash: " hash-a ", + Namespace: "tenant/tenant-a", + }, + Request: &knowledge.SearchRequest{Query: "deployment"}, + }) + + require.Error(t, err) + assert.Contains(t, err.Error(), MetadataUserIDHash) +} + +func TestServiceSearchKnowledgeInjectsTenantAndInternalUserFilters(t *testing.T) { + ctx := context.Background() + knowledgeBackend := &capturingKnowledge{} + service, err := New(ServiceConfig{ + Memory: memoryinmemory.NewMemoryService(), + Knowledge: knowledgeBackend, + }) + require.NoError(t, err) + scope := Scope{ + TenantID: "tenant-a", + AppID: "app-a", + InternalUserID: "internal-user-a", + UserIDHash: "hash-a", + Namespace: "tenant/tenant-a", + } + req := &knowledge.SearchRequest{ + Query: "deployment runbook", + SearchFilter: &knowledge.SearchFilter{ + Metadata: map[string]any{"category": "runbook"}, + }, + } + + _, err = service.SearchKnowledge(ctx, KnowledgeSearchRequest{ + Scope: scope, + Request: req, + }) + require.NoError(t, err) + + assert.Equal(t, "internal-user-a", knowledgeBackend.last.UserID) + assert.Equal(t, "tenant-a", knowledgeBackend.last.SearchFilter.Metadata[MetadataTenantID]) + assert.Equal(t, "app-a", knowledgeBackend.last.SearchFilter.Metadata[MetadataAppID]) + assert.Equal(t, "internal-user-a", knowledgeBackend.last.SearchFilter.Metadata[MetadataInternalUserID]) + assert.Equal(t, "hash-a", knowledgeBackend.last.SearchFilter.Metadata[MetadataUserIDHash]) + assert.Equal(t, "runbook", knowledgeBackend.last.SearchFilter.Metadata["category"]) + assert.NotContains(t, req.SearchFilter.Metadata, MetadataTenantID) + assert.NotContains(t, req.SearchFilter.Metadata, MetadataInternalUserID) +} + +func TestServiceRejectsKnowledgeFilterOutsideScope(t *testing.T) { + ctx := context.Background() + service, err := New(ServiceConfig{ + Memory: memoryinmemory.NewMemoryService(), + Knowledge: &capturingKnowledge{}, + }) + require.NoError(t, err) + scope := Scope{ + TenantID: "tenant-a", + AppID: "app-a", + InternalUserID: "internal-user-a", + UserIDHash: "hash-a", + Namespace: "tenant/tenant-a", + } + + tests := []struct { + name string + req *knowledge.SearchRequest + }{ + { + name: "conflicting tenant metadata", + req: &knowledge.SearchRequest{ + Query: "deployment", + SearchFilter: &knowledge.SearchFilter{ + Metadata: map[string]any{MetadataTenantID: "tenant-b"}, + }, + }, + }, + { + name: "conflicting internal user metadata", + req: &knowledge.SearchRequest{ + Query: "deployment", + SearchFilter: &knowledge.SearchFilter{ + Metadata: map[string]any{MetadataInternalUserID: "internal-user-b"}, + }, + }, + }, + { + name: "conflicting knowledge user", + req: &knowledge.SearchRequest{ + Query: "deployment", + UserID: "internal-user-b", + }, + }, + { + name: "non-string tenant metadata", + req: &knowledge.SearchRequest{ + Query: "deployment", + SearchFilter: &knowledge.SearchFilter{ + Metadata: map[string]any{MetadataTenantID: []string{"tenant-a"}}, + }, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := service.SearchKnowledge(ctx, KnowledgeSearchRequest{ + Scope: scope, + Request: tt.req, + }) + + require.ErrorIs(t, err, ErrFilterOutsideScope) + }) + } +} + +type capturingKnowledge struct { + last knowledge.SearchRequest +} + +func (c *capturingKnowledge) Search( + ctx context.Context, + req *knowledge.SearchRequest, +) (*knowledge.SearchResult, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + c.last = *req + return &knowledge.SearchResult{}, nil +} From 177d60951a02c04b53596d00f122948dabdccf61 Mon Sep 17 00:00:00 2001 From: Nene7ko_ <1604009816@qq.com> Date: Fri, 10 Jul 2026 17:13:57 +0800 Subject: [PATCH 55/95] artifact: add tenant-scoped object store Separate queryable artifact metadata from object content, add pending/deleting lifecycle states for retryable cleanup, enforce tenant and user scope, and cover upload, read, retry, deletion, and concurrency behavior. --- platform/artifactstore/doc.go | 11 + platform/artifactstore/errors.go | 40 ++ platform/artifactstore/inmemory.go | 310 +++++++++++ platform/artifactstore/service.go | 568 +++++++++++++++++++ platform/artifactstore/service_test.go | 719 +++++++++++++++++++++++++ platform/artifactstore/types.go | 107 ++++ 6 files changed, 1755 insertions(+) create mode 100644 platform/artifactstore/doc.go create mode 100644 platform/artifactstore/errors.go create mode 100644 platform/artifactstore/inmemory.go create mode 100644 platform/artifactstore/service.go create mode 100644 platform/artifactstore/service_test.go create mode 100644 platform/artifactstore/types.go diff --git a/platform/artifactstore/doc.go b/platform/artifactstore/doc.go new file mode 100644 index 0000000000..3d41dca6a3 --- /dev/null +++ b/platform/artifactstore/doc.go @@ -0,0 +1,11 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +// Package artifactstore provides a tenant-scoped artifact service that keeps +// queryable artifact metadata separate from object content bytes. +package artifactstore diff --git a/platform/artifactstore/errors.go b/platform/artifactstore/errors.go new file mode 100644 index 0000000000..d55cd633b5 --- /dev/null +++ b/platform/artifactstore/errors.go @@ -0,0 +1,40 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package artifactstore + +import "errors" + +var ( + // ErrTenantIDRequired indicates that the service lacks a tenant boundary. + ErrTenantIDRequired = errors.New("artifactstore tenant_id is required") + // ErrNamespaceRequired indicates that the storage namespace is missing. + ErrNamespaceRequired = errors.New("artifactstore namespace is required") + // ErrMetadataStoreRequired indicates that metadata storage was not configured. + ErrMetadataStoreRequired = errors.New("artifactstore metadata store is required") + // ErrObjectStoreRequired indicates that object storage was not configured. + ErrObjectStoreRequired = errors.New("artifactstore object store is required") + // ErrOutsideTenantScope indicates that a key or query escapes the tenant scope. + ErrOutsideTenantScope = errors.New("artifactstore key outside tenant scope") + // ErrEmptySessionInfo indicates that required session fields are missing. + ErrEmptySessionInfo = errors.New("artifactstore session info fields cannot be empty") + // ErrEmptyFilename indicates that the filename is empty. + ErrEmptyFilename = errors.New("artifactstore filename cannot be empty") + // ErrInvalidFilename indicates that the filename contains unsafe path data. + ErrInvalidFilename = errors.New("artifactstore filename contains invalid characters") + // ErrNilArtifact indicates that the artifact payload is nil. + ErrNilArtifact = errors.New("artifactstore artifact cannot be nil") + // ErrObjectNotFound indicates that object content is missing. + ErrObjectNotFound = errors.New("artifactstore object not found") + // ErrVersionConflict indicates that another writer committed the same version. + ErrVersionConflict = errors.New("artifactstore version conflict") + // ErrMetadataReservationNotFound indicates that a pending upload record is missing. + ErrMetadataReservationNotFound = errors.New("artifactstore metadata reservation not found") + // ErrArtifactWriteInProgress indicates that deletion raced with a pending upload. + ErrArtifactWriteInProgress = errors.New("artifactstore write in progress") +) diff --git a/platform/artifactstore/inmemory.go b/platform/artifactstore/inmemory.go new file mode 100644 index 0000000000..04124e8427 --- /dev/null +++ b/platform/artifactstore/inmemory.go @@ -0,0 +1,310 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package artifactstore + +import ( + "context" + "sort" + "sync" +) + +var ( + _ MetadataStore = (*InMemoryMetadataStore)(nil) + _ ObjectStore = (*InMemoryObjectStore)(nil) +) + +// InMemoryMetadataStore stores metadata records in memory for tests and local runs. +type InMemoryMetadataStore struct { + mu sync.RWMutex + records []MetadataRecord +} + +// NewInMemoryMetadataStore creates an empty in-memory metadata store. +func NewInMemoryMetadataStore() *InMemoryMetadataStore { + return &InMemoryMetadataStore{} +} + +// Put inserts or replaces one metadata record. +func (s *InMemoryMetadataStore) Put(ctx context.Context, record MetadataRecord) error { + if err := ctx.Err(); err != nil { + return err + } + if record.Status == "" { + record.Status = MetadataStatusActive + } + s.mu.Lock() + defer s.mu.Unlock() + for _, existing := range s.records { + if sameVersion(existing, record) { + return ErrVersionConflict + } + } + s.records = append(s.records, record) + return nil +} + +// Query returns records matching all non-empty query fields. +func (s *InMemoryMetadataStore) Query(ctx context.Context, query MetadataQuery) ([]MetadataRecord, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + s.mu.RLock() + defer s.mu.RUnlock() + records := make([]MetadataRecord, 0) + for _, record := range s.records { + if !query.IncludePending && record.Status == MetadataStatusPending { + continue + } + if !query.IncludeDeleting && record.Status == MetadataStatusDeleting { + continue + } + if !matchMetadata(record, query) { + continue + } + records = append(records, record) + } + sortMetadata(records) + return records, nil +} + +// Activate publishes one pending metadata reservation. +func (s *InMemoryMetadataStore) Activate( + ctx context.Context, + query MetadataQuery, + objectID string, +) error { + if err := ctx.Err(); err != nil { + return err + } + s.mu.Lock() + defer s.mu.Unlock() + for index := range s.records { + record := &s.records[index] + if !matchMetadata(*record, query) || record.ObjectID != objectID { + continue + } + if record.Status == MetadataStatusActive { + return nil + } + if record.Status != MetadataStatusPending { + return ErrMetadataReservationNotFound + } + record.Status = MetadataStatusActive + return nil + } + return ErrMetadataReservationNotFound +} + +// MarkDeleting atomically hides matching records and returns cleanup tombstones. +func (s *InMemoryMetadataStore) MarkDeleting( + ctx context.Context, + query MetadataQuery, +) ([]MetadataRecord, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + s.mu.Lock() + defer s.mu.Unlock() + indexes := make([]int, 0) + for index := range s.records { + if !matchMetadata(s.records[index], query) { + continue + } + if s.records[index].Status == MetadataStatusPending && !query.AllowPendingTransition { + return nil, ErrArtifactWriteInProgress + } + indexes = append(indexes, index) + } + records := make([]MetadataRecord, 0, len(indexes)) + for _, index := range indexes { + s.records[index].Status = MetadataStatusDeleting + records = append(records, s.records[index]) + } + sortMetadata(records) + return records, nil +} + +// Delete removes records matching all non-empty query fields and returns them. +func (s *InMemoryMetadataStore) Delete(ctx context.Context, query MetadataQuery) ([]MetadataRecord, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + s.mu.Lock() + defer s.mu.Unlock() + deleted := make([]MetadataRecord, 0) + kept := s.records[:0] + for _, record := range s.records { + if matchMetadata(record, query) { + deleted = append(deleted, record) + continue + } + kept = append(kept, record) + } + s.records = kept + sortMetadata(deleted) + return deleted, nil +} + +// HasInlineContent reports whether metadata contains embedded object bytes. +func (s *InMemoryMetadataStore) HasInlineContent(artifactID string) bool { + return false +} + +// InMemoryObjectStore stores object bytes in memory for tests and local runs. +type InMemoryObjectStore struct { + mu sync.RWMutex + objects map[string]objectValue + failNextPut []error + putAttempts int +} + +type objectValue struct { + data []byte + key string +} + +// NewInMemoryObjectStore creates an empty in-memory object store. +func NewInMemoryObjectStore() *InMemoryObjectStore { + return &InMemoryObjectStore{ + objects: make(map[string]objectValue), + } +} + +// Put stores object bytes by opaque object ID. +func (s *InMemoryObjectStore) Put(ctx context.Context, object ObjectRecord) error { + if err := ctx.Err(); err != nil { + return err + } + s.mu.Lock() + defer s.mu.Unlock() + s.putAttempts++ + if len(s.failNextPut) > 0 { + err := s.failNextPut[0] + s.failNextPut = s.failNextPut[1:] + return err + } + s.objects[object.ObjectID] = objectValue{ + data: append([]byte(nil), object.Data...), + key: "objects/" + object.TenantID + "/" + object.ObjectID, + } + return nil +} + +// Get returns a copy of object bytes. +func (s *InMemoryObjectStore) Get(ctx context.Context, objectID string) ([]byte, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + s.mu.RLock() + defer s.mu.RUnlock() + object, ok := s.objects[objectID] + if !ok { + return nil, ErrObjectNotFound + } + return append([]byte(nil), object.data...), nil +} + +// Delete removes object bytes. +func (s *InMemoryObjectStore) Delete(ctx context.Context, objectID string) error { + if err := ctx.Err(); err != nil { + return err + } + s.mu.Lock() + defer s.mu.Unlock() + delete(s.objects, objectID) + return nil +} + +// FailNextPut makes future Put calls fail with the supplied error in FIFO order. +func (s *InMemoryObjectStore) FailNextPut(err error) { + if err == nil { + return + } + s.mu.Lock() + defer s.mu.Unlock() + s.failNextPut = append(s.failNextPut, err) +} + +// PutAttempts returns the number of Put attempts made. +func (s *InMemoryObjectStore) PutAttempts() int { + s.mu.RLock() + defer s.mu.RUnlock() + return s.putAttempts +} + +// RawKey returns the internal object key for tests that check content-ref leakage. +func (s *InMemoryObjectStore) RawKey(objectID string) string { + s.mu.RLock() + defer s.mu.RUnlock() + if object, ok := s.objects[objectID]; ok { + return object.key + } + return "objects/" + objectID +} + +// ObjectIDs returns all currently stored object IDs. +func (s *InMemoryObjectStore) ObjectIDs() []string { + s.mu.RLock() + defer s.mu.RUnlock() + ids := make([]string, 0, len(s.objects)) + for id := range s.objects { + ids = append(ids, id) + } + sort.Strings(ids) + return ids +} + +type testingT interface { + Helper() + Fatalf(format string, args ...any) +} + +// MustData returns object data or fails the test. +func (s *InMemoryObjectStore) MustData(t testingT, objectID string) []byte { + t.Helper() + data, err := s.Get(context.Background(), objectID) + if err != nil { + t.Fatalf("object %q not found: %v", objectID, err) + } + return data +} + +func sameVersion(left MetadataRecord, right MetadataRecord) bool { + return left.TenantID == right.TenantID && + left.AppName == right.AppName && + left.UserID == right.UserID && + left.SessionID == right.SessionID && + left.Filename == right.Filename && + left.Version == right.Version +} + +func matchMetadata(record MetadataRecord, query MetadataQuery) bool { + if query.TenantID != "" && record.TenantID != query.TenantID { + return false + } + if query.AppName != "" && record.AppName != query.AppName { + return false + } + if query.UserID != "" && record.UserID != query.UserID { + return false + } + if query.SessionID != "" && record.SessionID != query.SessionID { + return false + } + if query.Filename != "" && record.Filename != query.Filename { + return false + } + if query.ObjectID != "" && record.ObjectID != query.ObjectID { + return false + } + if query.Version != nil && record.Version != *query.Version { + return false + } + return true +} diff --git a/platform/artifactstore/service.go b/platform/artifactstore/service.go new file mode 100644 index 0000000000..b49ef65289 --- /dev/null +++ b/platform/artifactstore/service.go @@ -0,0 +1,568 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package artifactstore + +import ( + "context" + "crypto/rand" + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + "sort" + "strings" + "sync" + "unicode" + + "trpc.group/trpc-go/trpc-agent-go/artifact" +) + +const ( + defaultContentType = "application/octet-stream" + userNamespace = "user:" + maxMetadataCommitAttempts = 8 +) + +var _ artifact.Service = (*Service)(nil) + +// Service implements artifact.Service using split metadata and object stores. +type Service struct { + tenantID string + namespace string + metadataStore MetadataStore + objectStore ObjectStore + maxAttempts int + mu sync.Mutex +} + +// New creates a tenant-scoped artifact service. +func New(config ServiceConfig) (*Service, error) { + tenantID := strings.TrimSpace(config.TenantID) + if tenantID == "" { + return nil, ErrTenantIDRequired + } + namespace := strings.TrimRight(strings.TrimSpace(config.Namespace), `/\|:`) + if namespace == "" { + return nil, ErrNamespaceRequired + } + if !namespaceContainsSegment(namespace, tenantID) { + return nil, ErrOutsideTenantScope + } + if config.MetadataStore == nil { + return nil, ErrMetadataStoreRequired + } + if config.ObjectStore == nil { + return nil, ErrObjectStoreRequired + } + maxAttempts := config.MaxAttempts + if maxAttempts == 0 { + maxAttempts = 1 + } + if maxAttempts < 0 { + return nil, fmt.Errorf("artifactstore max attempts cannot be negative") + } + return &Service{ + tenantID: tenantID, + namespace: namespace, + metadataStore: config.MetadataStore, + objectStore: config.ObjectStore, + maxAttempts: maxAttempts, + }, nil +} + +// SaveArtifact reserves metadata, uploads object bytes, then publishes the version. +func (s *Service) SaveArtifact( + ctx context.Context, + sessionInfo artifact.SessionInfo, + filename string, + art *artifact.Artifact, +) (int, error) { + if err := ctx.Err(); err != nil { + return 0, err + } + if err := s.validateSessionInfo(sessionInfo); err != nil { + return 0, err + } + if err := validateFilename(filename); err != nil { + return 0, err + } + if art == nil { + return 0, ErrNilArtifact + } + + s.mu.Lock() + defer s.mu.Unlock() + + data := append([]byte(nil), art.Data...) + digest := sha256.Sum256(data) + sha := hex.EncodeToString(digest[:]) + mimeType := strings.TrimSpace(art.MimeType) + if mimeType == "" { + mimeType = defaultContentType + } + artifactID := makeArtifactID(s.tenantID, sessionInfo, filename) + for commitAttempt := 0; commitAttempt < maxMetadataCommitAttempts; commitAttempt++ { + query := s.metadataQuery(sessionInfo, filename) + query.IncludePending = true + query.IncludeDeleting = true + records, err := s.metadataStore.Query(ctx, query) + if err != nil { + return 0, fmt.Errorf("query artifact metadata: %w", err) + } + version := nextVersion(records) + objectID, err := makeObjectID(artifactID, version, sha) + if err != nil { + return 0, err + } + record := MetadataRecord{ + TenantID: s.tenantID, + AppName: sessionInfo.AppName, + UserID: sessionInfo.UserID, + SessionID: metadataSessionID(sessionInfo, filename), + Filename: filename, + Version: version, + MimeType: mimeType, + SizeBytes: int64(len(data)), + SHA256: sha, + AttachmentKind: attachmentKind(mimeType), + ContentRef: makeContentRef(artifactID, version), + ObjectID: objectID, + ArtifactID: artifactID, + Status: MetadataStatusPending, + } + object := ObjectRecord{ + ObjectID: objectID, + TenantID: s.tenantID, + Data: data, + MimeType: mimeType, + SizeBytes: int64(len(data)), + SHA256: sha, + } + if err := s.metadataStore.Put(ctx, record); err != nil { + if errors.Is(err, ErrVersionConflict) { + continue + } + return 0, fmt.Errorf("reserve artifact metadata: %w", err) + } + if err := s.putObjectWithRetry(ctx, object); err != nil { + return 0, errors.Join(err, s.cleanupReservedArtifact(ctx, record)) + } + activateQuery := s.metadataQuery(sessionInfo, filename) + activateQuery.Version = &version + activateQuery.IncludePending = true + if err := s.metadataStore.Activate(ctx, activateQuery, objectID); err != nil { + return 0, errors.Join( + fmt.Errorf("activate artifact metadata: %w", err), + s.cleanupReservedArtifact(ctx, record), + ) + } + return version, nil + } + return 0, ErrVersionConflict +} + +// LoadArtifact loads object bytes using metadata as the authority. +func (s *Service) LoadArtifact( + ctx context.Context, + sessionInfo artifact.SessionInfo, + filename string, + version *int, +) (*artifact.Artifact, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + record, err := s.Metadata(ctx, sessionInfo, filename, version) + if err != nil || record == nil { + return nil, err + } + data, err := s.objectStore.Get(ctx, record.ObjectID) + if err != nil { + return nil, fmt.Errorf("get artifact object: %w", err) + } + return &artifact.Artifact{ + Data: data, + MimeType: record.MimeType, + Name: filename, + }, nil +} + +// Metadata returns one metadata record for the requested artifact version. +func (s *Service) Metadata( + ctx context.Context, + sessionInfo artifact.SessionInfo, + filename string, + version *int, +) (*MetadataRecord, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + if err := s.validateSessionInfo(sessionInfo); err != nil { + return nil, err + } + if err := validateFilename(filename); err != nil { + return nil, err + } + query := s.metadataQuery(sessionInfo, filename) + query.Version = version + records, err := s.metadataStore.Query(ctx, query) + if err != nil { + return nil, fmt.Errorf("query artifact metadata: %w", err) + } + if len(records) == 0 { + return nil, nil + } + sortMetadata(records) + record := records[len(records)-1] + return &record, nil +} + +// ListArtifactKeys lists artifact filenames within the session boundary. +func (s *Service) ListArtifactKeys( + ctx context.Context, + sessionInfo artifact.SessionInfo, +) ([]string, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + if err := s.validateSessionInfo(sessionInfo); err != nil { + return nil, err + } + records, err := s.metadataStore.Query(ctx, MetadataQuery{ + TenantID: s.tenantID, + AppName: sessionInfo.AppName, + UserID: sessionInfo.UserID, + SessionID: sessionInfo.SessionID, + }) + if err != nil { + return nil, fmt.Errorf("query artifact metadata: %w", err) + } + userRecords, err := s.metadataStore.Query(ctx, MetadataQuery{ + TenantID: s.tenantID, + AppName: sessionInfo.AppName, + UserID: sessionInfo.UserID, + SessionID: userArtifactSessionID, + }) + if err != nil { + return nil, fmt.Errorf("query user artifact metadata: %w", err) + } + records = append(records, userRecords...) + names := make(map[string]struct{}, len(records)) + for _, record := range records { + names[record.Filename] = struct{}{} + } + filenames := make([]string, 0, len(names)) + for filename := range names { + filenames = append(filenames, filename) + } + sort.Strings(filenames) + return filenames, nil +} + +// DeleteArtifact removes all metadata and object versions for an artifact. +func (s *Service) DeleteArtifact( + ctx context.Context, + sessionInfo artifact.SessionInfo, + filename string, +) error { + if err := ctx.Err(); err != nil { + return err + } + if err := s.validateSessionInfo(sessionInfo); err != nil { + return err + } + if err := validateFilename(filename); err != nil { + return err + } + s.mu.Lock() + defer s.mu.Unlock() + + query := s.metadataQuery(sessionInfo, filename) + records, err := s.metadataStore.MarkDeleting(ctx, query) + if err != nil { + return fmt.Errorf("mark artifact metadata deleting: %w", err) + } + if len(records) == 0 { + return nil + } + return s.cleanupMetadataRecords(ctx, query, records) +} + +func (s *Service) cleanupReservedArtifact( + ctx context.Context, + record MetadataRecord, +) error { + cleanupCtx := context.WithoutCancel(ctx) + version := record.Version + query := MetadataQuery{ + TenantID: record.TenantID, + AppName: record.AppName, + UserID: record.UserID, + SessionID: record.SessionID, + Filename: record.Filename, + Version: &version, + ObjectID: record.ObjectID, + IncludePending: true, + IncludeDeleting: true, + AllowPendingTransition: true, + } + records, err := s.metadataStore.MarkDeleting(cleanupCtx, query) + if err != nil { + return fmt.Errorf("mark reserved artifact deleting: %w", err) + } + if len(records) == 0 { + return ErrMetadataReservationNotFound + } + return s.cleanupMetadataRecords(cleanupCtx, query, records) +} + +func (s *Service) cleanupMetadataRecords( + ctx context.Context, + query MetadataQuery, + records []MetadataRecord, +) error { + var cleanupErrs []error + for _, record := range records { + if err := s.objectStore.Delete(ctx, record.ObjectID); err != nil { + cleanupErrs = append(cleanupErrs, fmt.Errorf( + "delete artifact object %q: %w", + record.ObjectID, + err, + )) + continue + } + version := record.Version + deleteQuery := query + deleteQuery.Version = &version + deleteQuery.ObjectID = record.ObjectID + deleteQuery.IncludePending = true + deleteQuery.IncludeDeleting = true + if _, err := s.metadataStore.Delete(ctx, deleteQuery); err != nil { + cleanupErrs = append(cleanupErrs, fmt.Errorf( + "delete artifact metadata version %d: %w", + version, + err, + )) + } + } + return errors.Join(cleanupErrs...) +} + +// ListVersions lists all versions of an artifact. +func (s *Service) ListVersions( + ctx context.Context, + sessionInfo artifact.SessionInfo, + filename string, +) ([]int, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + if err := s.validateSessionInfo(sessionInfo); err != nil { + return nil, err + } + if err := validateFilename(filename); err != nil { + return nil, err + } + records, err := s.metadataStore.Query(ctx, s.metadataQuery(sessionInfo, filename)) + if err != nil { + return nil, fmt.Errorf("query artifact metadata: %w", err) + } + versionSet := make(map[int]struct{}, len(records)) + for _, record := range records { + versionSet[record.Version] = struct{}{} + } + versions := make([]int, 0, len(versionSet)) + for version := range versionSet { + versions = append(versions, version) + } + sort.Ints(versions) + return versions, nil +} + +func (s *Service) putObjectWithRetry(ctx context.Context, object ObjectRecord) error { + var lastErr error + for attempt := 0; attempt < s.maxAttempts; attempt++ { + if err := ctx.Err(); err != nil { + return err + } + if err := s.objectStore.Put(ctx, object); err != nil { + lastErr = err + continue + } + return nil + } + return fmt.Errorf("put artifact object after %d attempts: %w", s.maxAttempts, lastErr) +} + +func (s *Service) metadataQuery(sessionInfo artifact.SessionInfo, filename string) MetadataQuery { + return MetadataQuery{ + TenantID: s.tenantID, + AppName: sessionInfo.AppName, + UserID: sessionInfo.UserID, + SessionID: metadataSessionID(sessionInfo, filename), + Filename: filename, + } +} + +func (s *Service) validateSessionInfo(info artifact.SessionInfo) error { + if strings.TrimSpace(info.AppName) == "" || + strings.TrimSpace(info.UserID) == "" || + strings.TrimSpace(info.SessionID) == "" { + return ErrEmptySessionInfo + } + if hasLeadingOrTrailingSpace(info.AppName) || + hasLeadingOrTrailingSpace(info.UserID) || + hasLeadingOrTrailingSpace(info.SessionID) { + return ErrEmptySessionInfo + } + if containsControl(info.AppName) || + containsControl(info.UserID) || + containsControl(info.SessionID) { + return ErrEmptySessionInfo + } + prefix := s.namespace + "/" + if !strings.HasPrefix(info.AppName, prefix) || strings.TrimSpace(strings.TrimPrefix(info.AppName, prefix)) == "" { + return ErrOutsideTenantScope + } + return nil +} + +func validateFilename(filename string) error { + if strings.TrimSpace(filename) == "" { + return ErrEmptyFilename + } + if hasLeadingOrTrailingSpace(filename) || + strings.Contains(filename, "\\") || + strings.Contains(filename, "\x00") || + containsControl(filename) { + return ErrInvalidFilename + } + for _, segment := range strings.Split(filename, "/") { + if segment == "" || segment == "." || segment == ".." { + return ErrInvalidFilename + } + } + return nil +} + +const userArtifactSessionID = "user" + +func metadataSessionID(info artifact.SessionInfo, filename string) string { + if strings.HasPrefix(filename, userNamespace) { + return userArtifactSessionID + } + return info.SessionID +} + +func nextVersion(records []MetadataRecord) int { + if len(records) == 0 { + return 0 + } + version := 0 + for _, record := range records { + if record.Version >= version { + version = record.Version + 1 + } + } + return version +} + +func sortMetadata(records []MetadataRecord) { + sort.Slice(records, func(i, j int) bool { + left := records[i] + right := records[j] + if left.TenantID != right.TenantID { + return left.TenantID < right.TenantID + } + if left.AppName != right.AppName { + return left.AppName < right.AppName + } + if left.UserID != right.UserID { + return left.UserID < right.UserID + } + if left.SessionID != right.SessionID { + return left.SessionID < right.SessionID + } + if left.Filename != right.Filename { + return left.Filename < right.Filename + } + return left.Version < right.Version + }) +} + +func attachmentKind(mimeType string) string { + switch { + case strings.HasPrefix(mimeType, "image/"): + return "image" + case strings.HasPrefix(mimeType, "audio/"): + return "audio" + case strings.HasPrefix(mimeType, "video/"): + return "video" + default: + return "file" + } +} + +func makeArtifactID(tenantID string, info artifact.SessionInfo, filename string) string { + return "art_" + scopedHash(tenantID, info.AppName, info.UserID, metadataSessionID(info, filename), filename) +} + +func makeObjectID(artifactID string, version int, sha string) (string, error) { + var nonce [16]byte + if _, err := rand.Read(nonce[:]); err != nil { + return "", fmt.Errorf("generate artifact object id: %w", err) + } + return "obj_" + scopedHash( + artifactID, + fmt.Sprintf("%d", version), + sha, + hex.EncodeToString(nonce[:]), + ), nil +} + +func makeContentRef(artifactID string, version int) string { + return fmt.Sprintf("artifact://%s?version=%d", artifactID, version) +} + +func scopedHash(parts ...string) string { + hash := sha256.New() + for _, part := range parts { + hash.Write([]byte(part)) + hash.Write([]byte{0}) + } + return hex.EncodeToString(hash.Sum(nil))[:32] +} + +func namespaceContainsSegment(namespace, tenantID string) bool { + for _, segment := range strings.FieldsFunc(namespace, func(r rune) bool { + switch r { + case '/', '\\', ':', '|': + return true + default: + return false + } + }) { + if segment == tenantID { + return true + } + } + return false +} + +func hasLeadingOrTrailingSpace(value string) bool { + return strings.TrimSpace(value) != value +} + +func containsControl(value string) bool { + for _, r := range value { + if unicode.IsControl(r) { + return true + } + } + return false +} diff --git a/platform/artifactstore/service_test.go b/platform/artifactstore/service_test.go new file mode 100644 index 0000000000..a742c59534 --- /dev/null +++ b/platform/artifactstore/service_test.go @@ -0,0 +1,719 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package artifactstore + +import ( + "context" + "errors" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "trpc.group/trpc-go/trpc-agent-go/artifact" +) + +func TestServiceStoresMetadataSeparatelyFromObjectContent(t *testing.T) { + ctx := context.Background() + metadata := NewInMemoryMetadataStore() + objects := NewInMemoryObjectStore() + service, err := New(ServiceConfig{ + TenantID: "tenant-a", + Namespace: "tenant/tenant-a", + MetadataStore: metadata, + ObjectStore: objects, + }) + require.NoError(t, err) + sessionInfo := artifact.SessionInfo{ + AppName: "tenant/tenant-a/app-a", + UserID: "internal-user-a", + SessionID: "session-a", + } + + version, err := service.SaveArtifact(ctx, sessionInfo, "diagram.png", &artifact.Artifact{ + Data: []byte("png-bytes"), + MimeType: "image/png", + Name: "diagram.png", + }) + require.NoError(t, err) + assert.Equal(t, 0, version) + + record, err := service.Metadata(ctx, sessionInfo, "diagram.png", &version) + require.NoError(t, err) + require.NotNil(t, record) + assert.Equal(t, "tenant-a", record.TenantID) + assert.Equal(t, "tenant/tenant-a/app-a", record.AppName) + assert.Equal(t, "internal-user-a", record.UserID) + assert.Equal(t, "session-a", record.SessionID) + assert.Equal(t, "diagram.png", record.Filename) + assert.Equal(t, "image/png", record.MimeType) + assert.Equal(t, int64(len("png-bytes")), record.SizeBytes) + assert.NotEmpty(t, record.SHA256) + assert.Equal(t, "image", record.AttachmentKind) + for _, secret := range []string{ + "secret", + record.TenantID, + record.AppName, + record.UserID, + record.SessionID, + record.Filename, + record.ObjectID, + objects.RawKey(record.ObjectID), + } { + assert.NotContains(t, record.ContentRef, secret) + } + assert.False(t, metadata.HasInlineContent(record.ArtifactID)) + assert.Equal(t, []byte("png-bytes"), objects.MustData(t, record.ObjectID)) + + loaded, err := service.LoadArtifact(ctx, sessionInfo, "diagram.png", nil) + require.NoError(t, err) + require.NotNil(t, loaded) + assert.Equal(t, []byte("png-bytes"), loaded.Data) + assert.Equal(t, "image/png", loaded.MimeType) + assert.Equal(t, "diagram.png", loaded.Name) + assert.Empty(t, loaded.URL) + + records, err := metadata.Query(ctx, MetadataQuery{ + TenantID: "tenant-a", + AppName: "tenant/tenant-a/app-a", + SessionID: "session-a", + }) + require.NoError(t, err) + require.Len(t, records, 1) + assert.Equal(t, record.ArtifactID, records[0].ArtifactID) +} + +func TestServiceRejectsCrossTenantAndCrossUserAccess(t *testing.T) { + ctx := context.Background() + metadata := NewInMemoryMetadataStore() + objects := NewInMemoryObjectStore() + service, err := New(ServiceConfig{ + TenantID: "tenant-a", + Namespace: "tenant/tenant-a", + MetadataStore: metadata, + ObjectStore: objects, + }) + require.NoError(t, err) + sessionInfo := artifact.SessionInfo{ + AppName: "tenant/tenant-a/app-a", + UserID: "internal-user-a", + SessionID: "session-a", + } + _, err = service.SaveArtifact(ctx, sessionInfo, "report.pdf", &artifact.Artifact{ + Data: []byte("pdf"), + MimeType: "application/pdf", + Name: "report.pdf", + }) + require.NoError(t, err) + + _, err = service.LoadArtifact(ctx, artifact.SessionInfo{ + AppName: "tenant/tenant-b/app-a", + UserID: "internal-user-a", + SessionID: "session-a", + }, "report.pdf", nil) + require.ErrorIs(t, err, ErrOutsideTenantScope) + + loaded, err := service.LoadArtifact(ctx, artifact.SessionInfo{ + AppName: "tenant/tenant-a/app-a", + UserID: "internal-user-b", + SessionID: "session-a", + }, "report.pdf", nil) + require.NoError(t, err) + assert.Nil(t, loaded) +} + +func TestServiceRetriesTransientObjectUploadFailure(t *testing.T) { + ctx := context.Background() + metadata := NewInMemoryMetadataStore() + objects := NewInMemoryObjectStore() + objects.FailNextPut(errors.New("temporary object store failure")) + service, err := New(ServiceConfig{ + TenantID: "tenant-a", + Namespace: "tenant/tenant-a", + MetadataStore: metadata, + ObjectStore: objects, + MaxAttempts: 2, + }) + require.NoError(t, err) + + version, err := service.SaveArtifact(ctx, artifact.SessionInfo{ + AppName: "tenant/tenant-a/app-a", + UserID: "internal-user-a", + SessionID: "session-a", + }, "attachment.txt", &artifact.Artifact{ + Data: []byte("hello"), + MimeType: "text/plain", + Name: "attachment.txt", + }) + require.NoError(t, err) + assert.Equal(t, 0, version) + assert.Equal(t, 2, objects.PutAttempts()) +} + +func TestServiceDeleteRemovesMetadataAndObjects(t *testing.T) { + ctx := context.Background() + metadata := NewInMemoryMetadataStore() + objects := NewInMemoryObjectStore() + service, err := New(ServiceConfig{ + TenantID: "tenant-a", + Namespace: "tenant/tenant-a", + MetadataStore: metadata, + ObjectStore: objects, + }) + require.NoError(t, err) + sessionInfo := artifact.SessionInfo{ + AppName: "tenant/tenant-a/app-a", + UserID: "internal-user-a", + SessionID: "session-a", + } + for _, content := range []string{"v0", "v1"} { + _, err := service.SaveArtifact(ctx, sessionInfo, "notes.txt", &artifact.Artifact{ + Data: []byte(content), + MimeType: "text/plain", + Name: "notes.txt", + }) + require.NoError(t, err) + } + + err = service.DeleteArtifact(ctx, sessionInfo, "notes.txt") + require.NoError(t, err) + + versions, err := service.ListVersions(ctx, sessionInfo, "notes.txt") + require.NoError(t, err) + assert.Empty(t, versions) + assert.Empty(t, objects.ObjectIDs()) +} + +func TestServiceDeleteKeepsObjectsWhenMarkDeletingFails(t *testing.T) { + ctx := context.Background() + metadata := &failingDeleteMetadataStore{InMemoryMetadataStore: NewInMemoryMetadataStore()} + objects := NewInMemoryObjectStore() + service, err := New(ServiceConfig{ + TenantID: "tenant-a", + Namespace: "tenant/tenant-a", + MetadataStore: metadata, + ObjectStore: objects, + }) + require.NoError(t, err) + sessionInfo := artifact.SessionInfo{ + AppName: "tenant/tenant-a/app-a", + UserID: "internal-user-a", + SessionID: "session-a", + } + version, err := service.SaveArtifact(ctx, sessionInfo, "notes.txt", &artifact.Artifact{ + Data: []byte("v0"), + MimeType: "text/plain", + Name: "notes.txt", + }) + require.NoError(t, err) + record, err := service.Metadata(ctx, sessionInfo, "notes.txt", &version) + require.NoError(t, err) + require.NotNil(t, record) + metadata.failDelete = errors.New("metadata delete unavailable") + + err = service.DeleteArtifact(ctx, sessionInfo, "notes.txt") + require.Error(t, err) + + stillPresent, err := service.Metadata(ctx, sessionInfo, "notes.txt", &version) + require.NoError(t, err) + require.NotNil(t, stillPresent) + assert.Equal(t, []byte("v0"), objects.MustData(t, record.ObjectID)) +} + +func TestServiceDeleteRetriesMetadataCleanupAfterObjectDelete(t *testing.T) { + ctx := context.Background() + metadata := &failingMetadataDeleteStore{ + InMemoryMetadataStore: NewInMemoryMetadataStore(), + failDelete: errors.New("metadata delete unavailable"), + } + objects := NewInMemoryObjectStore() + service, err := New(ServiceConfig{ + TenantID: "tenant-a", + Namespace: "tenant/tenant-a", + MetadataStore: metadata, + ObjectStore: objects, + }) + require.NoError(t, err) + sessionInfo := artifact.SessionInfo{ + AppName: "tenant/tenant-a/app-a", + UserID: "internal-user-a", + SessionID: "session-a", + } + _, err = service.SaveArtifact(ctx, sessionInfo, "notes.txt", &artifact.Artifact{ + Data: []byte("v0"), + MimeType: "text/plain", + Name: "notes.txt", + }) + require.NoError(t, err) + + err = service.DeleteArtifact(ctx, sessionInfo, "notes.txt") + require.Error(t, err) + loaded, err := service.LoadArtifact(ctx, sessionInfo, "notes.txt", nil) + require.NoError(t, err) + assert.Nil(t, loaded) + assert.Empty(t, objects.ObjectIDs()) + pending, err := metadata.Query(ctx, MetadataQuery{ + TenantID: "tenant-a", + AppName: sessionInfo.AppName, + UserID: sessionInfo.UserID, + SessionID: sessionInfo.SessionID, + Filename: "notes.txt", + IncludeDeleting: true, + }) + require.NoError(t, err) + require.Len(t, pending, 1) + assert.Equal(t, MetadataStatusDeleting, pending[0].Status) + + require.NoError(t, service.DeleteArtifact(ctx, sessionInfo, "notes.txt")) + pending, err = metadata.Query(ctx, MetadataQuery{ + TenantID: "tenant-a", + AppName: sessionInfo.AppName, + UserID: sessionInfo.UserID, + SessionID: sessionInfo.SessionID, + Filename: "notes.txt", + IncludeDeleting: true, + }) + require.NoError(t, err) + assert.Empty(t, pending) +} + +func TestServiceSupportsNestedArtifactFilenames(t *testing.T) { + ctx := context.Background() + metadata := NewInMemoryMetadataStore() + objects := NewInMemoryObjectStore() + service, err := New(ServiceConfig{ + TenantID: "tenant-a", + Namespace: "tenant/tenant-a", + MetadataStore: metadata, + ObjectStore: objects, + }) + require.NoError(t, err) + sessionInfo := artifact.SessionInfo{ + AppName: "tenant/tenant-a/app-a", + UserID: "internal-user-a", + SessionID: "session-a", + } + + _, err = service.SaveArtifact(ctx, sessionInfo, "out/site.zip", &artifact.Artifact{ + Data: []byte("zip"), + MimeType: "application/zip", + Name: "site.zip", + }) + require.NoError(t, err) + + loaded, err := service.LoadArtifact(ctx, sessionInfo, "out/site.zip", nil) + require.NoError(t, err) + require.NotNil(t, loaded) + assert.Equal(t, []byte("zip"), loaded.Data) + keys, err := service.ListArtifactKeys(ctx, sessionInfo) + require.NoError(t, err) + assert.Equal(t, []string{"out/site.zip"}, keys) +} + +func TestServiceDeleteRetriesPendingObjectCleanup(t *testing.T) { + ctx := context.Background() + metadata := NewInMemoryMetadataStore() + objects := &failingDeleteObjectStore{ + InMemoryObjectStore: NewInMemoryObjectStore(), + failDelete: errors.New("object delete unavailable"), + } + service, err := New(ServiceConfig{ + TenantID: "tenant-a", + Namespace: "tenant/tenant-a", + MetadataStore: metadata, + ObjectStore: objects, + }) + require.NoError(t, err) + sessionInfo := artifact.SessionInfo{ + AppName: "tenant/tenant-a/app-a", + UserID: "internal-user-a", + SessionID: "session-a", + } + version, err := service.SaveArtifact(ctx, sessionInfo, "notes.txt", &artifact.Artifact{ + Data: []byte("v0"), + MimeType: "text/plain", + Name: "notes.txt", + }) + require.NoError(t, err) + record, err := service.Metadata(ctx, sessionInfo, "notes.txt", &version) + require.NoError(t, err) + require.NotNil(t, record) + + err = service.DeleteArtifact(ctx, sessionInfo, "notes.txt") + require.Error(t, err) + + loaded, err := service.LoadArtifact(ctx, sessionInfo, "notes.txt", nil) + require.NoError(t, err) + assert.Nil(t, loaded) + pending, err := metadata.Query(ctx, MetadataQuery{ + TenantID: "tenant-a", + AppName: sessionInfo.AppName, + UserID: sessionInfo.UserID, + SessionID: sessionInfo.SessionID, + Filename: "notes.txt", + IncludeDeleting: true, + }) + require.NoError(t, err) + require.Len(t, pending, 1) + assert.Equal(t, MetadataStatusDeleting, pending[0].Status) + assert.Equal(t, []byte("v0"), objects.MustData(t, record.ObjectID)) + + err = service.DeleteArtifact(ctx, sessionInfo, "notes.txt") + require.NoError(t, err) + pending, err = metadata.Query(ctx, MetadataQuery{ + TenantID: "tenant-a", + AppName: sessionInfo.AppName, + UserID: sessionInfo.UserID, + SessionID: sessionInfo.SessionID, + Filename: "notes.txt", + IncludeDeleting: true, + }) + require.NoError(t, err) + assert.Empty(t, pending) + assert.Empty(t, objects.ObjectIDs()) +} + +func TestServiceFailedMetadataActivationKeepsRetryableCleanupTombstone(t *testing.T) { + ctx := context.Background() + metadata := &failingActivateMetadataStore{ + InMemoryMetadataStore: NewInMemoryMetadataStore(), + failActivate: errors.New("metadata activation unavailable"), + } + objects := &failingDeleteObjectStore{ + InMemoryObjectStore: NewInMemoryObjectStore(), + failDelete: errors.New("object cleanup unavailable"), + } + service, err := New(ServiceConfig{ + TenantID: "tenant-a", + Namespace: "tenant/tenant-a", + MetadataStore: metadata, + ObjectStore: objects, + }) + require.NoError(t, err) + sessionInfo := artifact.SessionInfo{ + AppName: "tenant/tenant-a/app-a", + UserID: "internal-user-a", + SessionID: "session-a", + } + + _, err = service.SaveArtifact(ctx, sessionInfo, "notes.txt", &artifact.Artifact{ + Data: []byte("v0"), + MimeType: "text/plain", + Name: "notes.txt", + }) + require.Error(t, err) + + loaded, err := service.LoadArtifact(ctx, sessionInfo, "notes.txt", nil) + require.NoError(t, err) + assert.Nil(t, loaded) + require.Len(t, objects.ObjectIDs(), 1) + pending, err := metadata.Query(ctx, MetadataQuery{ + TenantID: "tenant-a", + AppName: sessionInfo.AppName, + UserID: sessionInfo.UserID, + SessionID: sessionInfo.SessionID, + Filename: "notes.txt", + IncludeDeleting: true, + }) + require.NoError(t, err) + require.Len(t, pending, 1) + assert.Equal(t, MetadataStatusDeleting, pending[0].Status) + + require.NoError(t, service.DeleteArtifact(ctx, sessionInfo, "notes.txt")) + pending, err = metadata.Query(ctx, MetadataQuery{ + TenantID: "tenant-a", + AppName: sessionInfo.AppName, + UserID: sessionInfo.UserID, + SessionID: sessionInfo.SessionID, + Filename: "notes.txt", + IncludeDeleting: true, + }) + require.NoError(t, err) + assert.Empty(t, pending) + assert.Empty(t, objects.ObjectIDs()) +} + +func TestServiceFailedActivationCleansUpAfterRequestCancellation(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + metadata := &failingActivateMetadataStore{ + InMemoryMetadataStore: NewInMemoryMetadataStore(), + failActivate: errors.New("metadata activation unavailable"), + cancel: cancel, + } + objects := NewInMemoryObjectStore() + service, err := New(ServiceConfig{ + TenantID: "tenant-a", + Namespace: "tenant/tenant-a", + MetadataStore: metadata, + ObjectStore: objects, + }) + require.NoError(t, err) + sessionInfo := artifact.SessionInfo{ + AppName: "tenant/tenant-a/app-a", + UserID: "internal-user-a", + SessionID: "session-a", + } + + _, err = service.SaveArtifact(ctx, sessionInfo, "notes.txt", &artifact.Artifact{ + Data: []byte("v0"), + MimeType: "text/plain", + Name: "notes.txt", + }) + require.Error(t, err) + require.ErrorIs(t, ctx.Err(), context.Canceled) + assert.Empty(t, objects.ObjectIDs()) + pending, err := metadata.Query(context.Background(), MetadataQuery{ + TenantID: "tenant-a", + AppName: sessionInfo.AppName, + UserID: sessionInfo.UserID, + SessionID: sessionInfo.SessionID, + Filename: "notes.txt", + IncludePending: true, + IncludeDeleting: true, + }) + require.NoError(t, err) + assert.Empty(t, pending) +} + +func TestServiceConcurrentDeleteAndReuploadKeepsNewObject(t *testing.T) { + ctx := context.Background() + metadata := NewInMemoryMetadataStore() + objects := &blockingDeleteObjectStore{ + InMemoryObjectStore: NewInMemoryObjectStore(), + deleteStarted: make(chan struct{}), + continueDelete: make(chan struct{}), + } + deleteService, err := New(ServiceConfig{ + TenantID: "tenant-a", + Namespace: "tenant/tenant-a", + MetadataStore: metadata, + ObjectStore: objects, + }) + require.NoError(t, err) + saveService, err := New(ServiceConfig{ + TenantID: "tenant-a", + Namespace: "tenant/tenant-a", + MetadataStore: metadata, + ObjectStore: objects, + }) + require.NoError(t, err) + sessionInfo := artifact.SessionInfo{ + AppName: "tenant/tenant-a/app-a", + UserID: "internal-user-a", + SessionID: "session-a", + } + content := []byte("same-content") + _, err = saveService.SaveArtifact(ctx, sessionInfo, "notes.txt", &artifact.Artifact{ + Data: content, + MimeType: "text/plain", + Name: "notes.txt", + }) + require.NoError(t, err) + + deleteDone := make(chan error, 1) + go func() { + deleteDone <- deleteService.DeleteArtifact(ctx, sessionInfo, "notes.txt") + }() + <-objects.deleteStarted + + version, err := saveService.SaveArtifact(ctx, sessionInfo, "notes.txt", &artifact.Artifact{ + Data: content, + MimeType: "text/plain", + Name: "notes.txt", + }) + require.NoError(t, err) + assert.Equal(t, 1, version) + close(objects.continueDelete) + require.NoError(t, <-deleteDone) + + loaded, err := saveService.LoadArtifact(ctx, sessionInfo, "notes.txt", nil) + require.NoError(t, err) + require.NotNil(t, loaded) + assert.Equal(t, content, loaded.Data) + versions, err := saveService.ListVersions(ctx, sessionInfo, "notes.txt") + require.NoError(t, err) + assert.Equal(t, []int{1}, versions) +} + +func TestServiceDeleteDoesNotCancelPendingUpload(t *testing.T) { + ctx := context.Background() + metadata := NewInMemoryMetadataStore() + objects := &blockingFirstPutObjectStore{ + InMemoryObjectStore: NewInMemoryObjectStore(), + firstPutStarted: make(chan struct{}), + continueFirstPut: make(chan struct{}), + } + firstService, err := New(ServiceConfig{ + TenantID: "tenant-a", + Namespace: "tenant/tenant-a", + MetadataStore: metadata, + ObjectStore: objects, + }) + require.NoError(t, err) + secondService, err := New(ServiceConfig{ + TenantID: "tenant-a", + Namespace: "tenant/tenant-a", + MetadataStore: metadata, + ObjectStore: objects, + }) + require.NoError(t, err) + sessionInfo := artifact.SessionInfo{ + AppName: "tenant/tenant-a/app-a", + UserID: "internal-user-a", + SessionID: "session-a", + } + + firstDone := make(chan struct { + version int + err error + }, 1) + go func() { + version, saveErr := firstService.SaveArtifact(ctx, sessionInfo, "notes.txt", &artifact.Artifact{ + Data: []byte("first"), + MimeType: "text/plain", + Name: "notes.txt", + }) + firstDone <- struct { + version int + err error + }{version: version, err: saveErr} + }() + <-objects.firstPutStarted + + err = secondService.DeleteArtifact(ctx, sessionInfo, "notes.txt") + require.ErrorIs(t, err, ErrArtifactWriteInProgress) + + secondVersion, err := secondService.SaveArtifact(ctx, sessionInfo, "notes.txt", &artifact.Artifact{ + Data: []byte("second"), + MimeType: "text/plain", + Name: "notes.txt", + }) + require.NoError(t, err) + assert.Equal(t, 1, secondVersion) + + close(objects.continueFirstPut) + firstResult := <-firstDone + require.NoError(t, firstResult.err) + assert.Equal(t, 0, firstResult.version) + + loaded, err := secondService.LoadArtifact(ctx, sessionInfo, "notes.txt", nil) + require.NoError(t, err) + require.NotNil(t, loaded) + assert.Equal(t, []byte("second"), loaded.Data) + versions, err := secondService.ListVersions(ctx, sessionInfo, "notes.txt") + require.NoError(t, err) + assert.Equal(t, []int{0, 1}, versions) + assert.Len(t, objects.ObjectIDs(), 2) +} + +type failingDeleteMetadataStore struct { + *InMemoryMetadataStore + failDelete error +} + +type failingMetadataDeleteStore struct { + *InMemoryMetadataStore + failDelete error +} + +func (s *failingMetadataDeleteStore) Delete( + ctx context.Context, + query MetadataQuery, +) ([]MetadataRecord, error) { + if s.failDelete != nil { + err := s.failDelete + s.failDelete = nil + return nil, err + } + return s.InMemoryMetadataStore.Delete(ctx, query) +} + +type failingActivateMetadataStore struct { + *InMemoryMetadataStore + failActivate error + cancel context.CancelFunc +} + +func (s *failingActivateMetadataStore) Activate( + ctx context.Context, + query MetadataQuery, + objectID string, +) error { + if s.failActivate != nil { + err := s.failActivate + s.failActivate = nil + if s.cancel != nil { + s.cancel() + } + return err + } + return s.InMemoryMetadataStore.Activate(ctx, query, objectID) +} + +type failingDeleteObjectStore struct { + *InMemoryObjectStore + failDelete error +} + +func (s *failingDeleteObjectStore) Delete(ctx context.Context, objectID string) error { + if s.failDelete != nil { + err := s.failDelete + s.failDelete = nil + return err + } + return s.InMemoryObjectStore.Delete(ctx, objectID) +} + +type blockingDeleteObjectStore struct { + *InMemoryObjectStore + deleteStarted chan struct{} + continueDelete chan struct{} + once sync.Once +} + +func (s *blockingDeleteObjectStore) Delete(ctx context.Context, objectID string) error { + s.once.Do(func() { + close(s.deleteStarted) + <-s.continueDelete + }) + return s.InMemoryObjectStore.Delete(ctx, objectID) +} + +type blockingFirstPutObjectStore struct { + *InMemoryObjectStore + mu sync.Mutex + putCalls int + firstPutStarted chan struct{} + continueFirstPut chan struct{} +} + +func (s *blockingFirstPutObjectStore) Put(ctx context.Context, object ObjectRecord) error { + s.mu.Lock() + s.putCalls++ + call := s.putCalls + s.mu.Unlock() + if call == 1 { + close(s.firstPutStarted) + <-s.continueFirstPut + } + return s.InMemoryObjectStore.Put(ctx, object) +} + +func (s *failingDeleteMetadataStore) MarkDeleting( + ctx context.Context, + query MetadataQuery, +) ([]MetadataRecord, error) { + if s.failDelete != nil { + err := s.failDelete + s.failDelete = nil + return nil, err + } + return s.InMemoryMetadataStore.MarkDeleting(ctx, query) +} diff --git a/platform/artifactstore/types.go b/platform/artifactstore/types.go new file mode 100644 index 0000000000..55f9e62a2c --- /dev/null +++ b/platform/artifactstore/types.go @@ -0,0 +1,107 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package artifactstore + +import "context" + +// MetadataStore stores queryable artifact metadata without embedding object bytes. +type MetadataStore interface { + // Put reserves one scoped version and returns ErrVersionConflict when the + // scoped version already exists. + Put(ctx context.Context, record MetadataRecord) error + // Query hides pending uploads and deleting tombstones unless requested. + Query(ctx context.Context, query MetadataQuery) ([]MetadataRecord, error) + // Activate publishes a pending version after its object upload succeeds. + Activate(ctx context.Context, query MetadataQuery, objectID string) error + // MarkDeleting atomically hides matching records. Repeated calls must also + // return existing tombstones so failed object cleanup can be retried. It + // returns ErrArtifactWriteInProgress rather than changing pending uploads, + // unless AllowPendingTransition is set for an exact owner cleanup. + MarkDeleting(ctx context.Context, query MetadataQuery) ([]MetadataRecord, error) + // Delete permanently removes matching metadata after object cleanup. + Delete(ctx context.Context, query MetadataQuery) ([]MetadataRecord, error) +} + +// ObjectStore stores artifact object content addressed by opaque object IDs. +type ObjectStore interface { + Put(ctx context.Context, object ObjectRecord) error + Get(ctx context.Context, objectID string) ([]byte, error) + // Delete must be idempotent and treat an already-missing object as success. + Delete(ctx context.Context, objectID string) error +} + +// MetadataStatus describes whether an artifact version is visible or pending cleanup. +type MetadataStatus string + +const ( + // MetadataStatusPending reserves a version while its object upload commits. + MetadataStatusPending MetadataStatus = "pending" + // MetadataStatusActive makes the artifact version visible to normal queries. + MetadataStatusActive MetadataStatus = "active" + // MetadataStatusDeleting hides the version while object cleanup is pending. + MetadataStatusDeleting MetadataStatus = "deleting" +) + +// ServiceConfig wires the metadata and object stores for one tenant namespace. +type ServiceConfig struct { + TenantID string + Namespace string + MetadataStore MetadataStore + ObjectStore ObjectStore + MaxAttempts int +} + +// MetadataRecord describes one artifact version. +type MetadataRecord struct { + TenantID string + AppName string + UserID string + SessionID string + Filename string + Version int + MimeType string + SizeBytes int64 + SHA256 string + AttachmentKind string + ContentRef string + // ObjectID is an opaque backend identifier, not a raw object key. + // Store implementations must not encode secrets or credentials in it. + ObjectID string + ArtifactID string + Status MetadataStatus +} + +// MetadataQuery filters artifact metadata records. Empty string fields are not +// applied, while Version filters only when non-nil. +type MetadataQuery struct { + TenantID string + AppName string + UserID string + SessionID string + Filename string + Version *int + ObjectID string + // IncludePending includes upload reservations retained for safe cleanup. + IncludePending bool + // IncludeDeleting includes tombstones retained for retryable object cleanup. + IncludeDeleting bool + // AllowPendingTransition permits an exact ObjectID owner cleanup to cancel + // its own pending upload. + AllowPendingTransition bool +} + +// ObjectRecord contains the bytes written to object storage. +type ObjectRecord struct { + ObjectID string + TenantID string + Data []byte + MimeType string + SizeBytes int64 + SHA256 string +} From fd77bf695a1a2e7563b1e8d049b5f40d47df586e Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 06:14:12 +0800 Subject: [PATCH 56/95] feat(platform): wire tenant storage into runtimes --- platform/gateway/registry.go | 1 + platform/gateway/service.go | 104 +++++- platform/gateway/service_test.go | 114 ++++++ platform/worker/builder.go | 219 +++++++++++ platform/worker/builder_test.go | 604 +++++++++++++++++++++++++++++++ platform/worker/doc.go | 10 + platform/worker/errors.go | 30 ++ 7 files changed, 1070 insertions(+), 12 deletions(-) create mode 100644 platform/worker/builder.go create mode 100644 platform/worker/builder_test.go create mode 100644 platform/worker/doc.go create mode 100644 platform/worker/errors.go diff --git a/platform/gateway/registry.go b/platform/gateway/registry.go index c3eab62739..d370596827 100644 --- a/platform/gateway/registry.go +++ b/platform/gateway/registry.go @@ -22,6 +22,7 @@ type Runtime struct { App platform.AgentApp Binding platform.ChannelBinding Runner runner.Runner + Audit platform.AuditSink } // Validate checks that the runtime can process inbound messages. diff --git a/platform/gateway/service.go b/platform/gateway/service.go index d970dec2c4..ef12dd2fb6 100644 --- a/platform/gateway/service.go +++ b/platform/gateway/service.go @@ -12,6 +12,7 @@ import ( "context" "errors" "fmt" + "reflect" "strings" "time" @@ -133,7 +134,8 @@ func (s *Service) HandleInbound( if err != nil { return Result{}, err } - text, err := s.validateInboundContent(ctx, routeSpan, msg, start) + auditSink := s.auditSinkForRuntime(runtime) + text, err := s.validateInboundContent(ctx, routeSpan, msg, start, auditSink) if err != nil { return Result{}, err } @@ -172,6 +174,7 @@ func (s *Service) HandleInbound( routeCtx, ctx, runtime, + auditSink, msg, inboundRunInput{ Text: text, @@ -223,6 +226,17 @@ func (s *Service) lookupRuntime( recordSpanError(routeSpan, err) return Runtime{}, err } + if err := authorizeBinding(runtime.Binding, msg); err != nil { + s.writeRejectAuditTo( + auditCtx, + s.auditSinkForRuntime(runtime), + msg, + start, + err, + ) + recordSpanError(routeSpan, err) + return Runtime{}, err + } return runtime, nil } @@ -233,7 +247,7 @@ func validateRuntimeForMessage(runtime Runtime, msg platform.InboundMessage) err if !runtime.matchesInbound(msg) { return ErrRuntimeMismatch } - return authorizeBinding(runtime.Binding, msg) + return nil } func (s *Service) validateInboundContent( @@ -241,10 +255,11 @@ func (s *Service) validateInboundContent( routeSpan oteltrace.Span, msg platform.InboundMessage, start time.Time, + auditSink platform.AuditSink, ) (string, error) { text, err := inboundText(msg) if err != nil { - s.writeRejectAudit(ctx, msg, start, err) + s.writeRejectAuditTo(ctx, auditSink, msg, start, err) recordSpanError(routeSpan, err) return "", err } @@ -368,20 +383,37 @@ func (s *Service) runAndReply( routeCtx context.Context, auditCtx context.Context, runtime Runtime, + auditSink platform.AuditSink, msg platform.InboundMessage, input inboundRunInput, ) (Result, error) { - content, err := s.runGatewayRunner(routeCtx, auditCtx, runtime, msg, input) + content, err := s.runGatewayRunner( + routeCtx, + auditCtx, + runtime, + auditSink, + msg, + input, + ) if err != nil { return Result{}, err } - return s.writeReply(routeCtx, auditCtx, runtime, msg, input, content) + return s.writeReply( + routeCtx, + auditCtx, + runtime, + auditSink, + msg, + input, + content, + ) } func (s *Service) runGatewayRunner( routeCtx context.Context, auditCtx context.Context, runtime Runtime, + auditSink platform.AuditSink, msg platform.InboundMessage, input inboundRunInput, ) (string, error) { @@ -402,13 +434,13 @@ func (s *Service) runGatewayRunner( agent.WithLatencyDiagnosticsEvents(false), ) if err != nil { - s.writeAudit(auditCtx, auditFromMessage(msg, input.SessionID, input.InternalUserID, "runner_error", err.Error(), input.Start, err)) + s.writeAuditTo(auditCtx, auditSink, auditFromMessage(msg, input.SessionID, input.InternalUserID, "runner_error", err.Error(), input.Start, err)) recordSpanError(runnerSpan, err) return "", err } content, err := collectAssistantText(auditCtx, ch) if err != nil { - s.writeAudit(auditCtx, auditFromMessage(msg, input.SessionID, input.InternalUserID, "runner_error", err.Error(), input.Start, err)) + s.writeAuditTo(auditCtx, auditSink, auditFromMessage(msg, input.SessionID, input.InternalUserID, "runner_error", err.Error(), input.Start, err)) recordSpanError(runnerSpan, err) return "", err } @@ -419,6 +451,7 @@ func (s *Service) writeReply( routeCtx context.Context, auditCtx context.Context, runtime Runtime, + auditSink platform.AuditSink, msg platform.InboundMessage, input inboundRunInput, content string, @@ -440,7 +473,7 @@ func (s *Service) writeReply( defer replySpan.End() setInboundTraceAttributes(replySpan, msg, input.SessionID, input.RequestID, input.InternalUserID) if err := s.outboundStore.Save(replyCtx, reply.ResultRef, outbound); err != nil { - s.writeAudit(auditCtx, auditFromMessage(msg, input.SessionID, input.InternalUserID, "outbound_error", err.Error(), input.Start, err)) + s.writeAuditTo(auditCtx, auditSink, auditFromMessage(msg, input.SessionID, input.InternalUserID, "outbound_error", err.Error(), input.Start, err)) recordSpanError(replySpan, err) return Result{}, err } @@ -453,7 +486,7 @@ func (s *Service) writeReply( recordSpanError(replySpan, markErr) return Result{}, markErr } - s.writeAudit(auditCtx, auditFromMessage(msg, input.SessionID, input.InternalUserID, "outbound_error", err.Error(), input.Start, err)) + s.writeAuditTo(auditCtx, auditSink, auditFromMessage(msg, input.SessionID, input.InternalUserID, "outbound_error", err.Error(), input.Start, err)) recordSpanError(replySpan, err) return Result{}, err } @@ -464,7 +497,7 @@ func (s *Service) writeReply( } s.writeMessageEvent(auditCtx, messageEventFromInbound(msg, input.SessionID, input.Key, input.RequestID, reply.InboundSequence, input.Start)) s.writeMessageEvent(auditCtx, messageEventFromAssistant(msg, input.SessionID, reply.ResultRef, input.RequestID, reply.AssistantSequence, s.now())) - s.writeAudit(auditCtx, auditFromMessage(msg, input.SessionID, input.InternalUserID, "completed", "", input.Start, nil)) + s.writeAuditTo(auditCtx, auditSink, auditFromMessage(msg, input.SessionID, input.InternalUserID, "completed", "", input.Start, nil)) return Result{ RequestID: input.RequestID, SessionID: input.SessionID, @@ -502,6 +535,20 @@ func (s *Service) writeRejectAudit( s.writeAudit(ctx, auditFromMessage(msg, "", "", "reject", err.Error(), start, err)) } +func (s *Service) writeRejectAuditTo( + ctx context.Context, + auditSink platform.AuditSink, + msg platform.InboundMessage, + start time.Time, + err error, +) { + s.writeAuditTo( + ctx, + auditSink, + auditFromMessage(msg, "", "", "reject", err.Error(), start, err), + ) +} + func (s *Service) validateService() error { if s.registry == nil { return fmt.Errorf("gateway registry is required") @@ -703,10 +750,43 @@ func redactAuditReason(reason string) string { } func (s *Service) writeAudit(ctx context.Context, record platform.AuditRecord) { - if s.auditSink == nil { + s.writeAuditTo(ctx, s.auditSink, record) +} + +func (s *Service) writeAuditTo( + ctx context.Context, + auditSink platform.AuditSink, + record platform.AuditRecord, +) { + if isNilAuditSink(auditSink) { return } - _ = s.auditSink.WriteAudit(ctx, record) + _ = auditSink.WriteAudit(ctx, record) +} + +func (s *Service) auditSinkForRuntime(runtime Runtime) platform.AuditSink { + if !isNilAuditSink(runtime.Audit) { + return runtime.Audit + } + return s.auditSink +} + +func isNilAuditSink(auditSink platform.AuditSink) bool { + if auditSink == nil { + return true + } + reflected := reflect.ValueOf(auditSink) + switch reflected.Kind() { + case reflect.Chan, + reflect.Func, + reflect.Interface, + reflect.Map, + reflect.Pointer, + reflect.Slice: + return reflected.IsNil() + default: + return false + } } func (s *Service) writeMessageEvent(ctx context.Context, event platform.MessageEvent) { diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go index 7702a86512..985dd06453 100644 --- a/platform/gateway/service_test.go +++ b/platform/gateway/service_test.go @@ -351,6 +351,120 @@ func TestServiceHandleInboundCoversMinimumLoopAcceptance(t *testing.T) { assert.Equal(t, platform.IdempotencyStatusCompleted, groupResult.Status) } +func TestServiceHandleInboundPrefersRuntimeAuditSink(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + runtimeAudit := platform.NewInMemoryAuditSink() + fallbackAudit := platform.NewInMemoryAuditSink() + runtime := validRuntime( + "tenant-a", + &recordingRunner{response: "runtime audit reply"}, + ) + runtime.Audit = runtimeAudit + require.NoError(t, registry.Register(runtime)) + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + WithAuditSink(fallbackAudit), + ) + + result, err := svc.HandleInbound( + ctx, + inbound("tenant-a", "msg-runtime-audit", "user-a", "hello"), + ) + require.NoError(t, err) + assert.Equal(t, "runtime audit reply", result.Outbound.Content) + require.Len(t, runtimeAudit.Records(), 1) + assert.Equal(t, "completed", runtimeAudit.Records()[0].Decision) + assert.Empty(t, fallbackAudit.Records()) +} + +func TestServiceHandleInboundUsesFallbackAuditForInvalidRuntime(t *testing.T) { + ctx := context.Background() + runtimeAudit := platform.NewInMemoryAuditSink() + fallbackAudit := platform.NewInMemoryAuditSink() + runtime := validRuntime( + "tenant-a", + &recordingRunner{response: "unused"}, + ) + runtime.App.AppID = "other-app" + runtime.Binding.AppID = "other-app" + runtime.Audit = runtimeAudit + svc := NewService( + staticRegistry{runtime: runtime}, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + WithAuditSink(fallbackAudit), + ) + + _, err := svc.HandleInbound( + ctx, + inbound("tenant-a", "msg-runtime-mismatch", "user-a", "hello"), + ) + require.ErrorIs(t, err, ErrRuntimeMismatch) + assert.Empty(t, runtimeAudit.Records()) + require.Len(t, fallbackAudit.Records(), 1) + assert.Equal(t, "reject", fallbackAudit.Records()[0].Decision) +} + +func TestServiceHandleInboundFallsBackFromTypedNilRuntimeAudit(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + fallbackAudit := platform.NewInMemoryAuditSink() + runtime := validRuntime( + "tenant-a", + &recordingRunner{response: "fallback audit reply"}, + ) + var typedNilAudit *platform.InMemoryAuditSink + runtime.Audit = typedNilAudit + require.NoError(t, registry.Register(runtime)) + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + WithAuditSink(fallbackAudit), + ) + + result, err := svc.HandleInbound( + ctx, + inbound("tenant-a", "msg-typed-nil-audit", "user-a", "hello"), + ) + require.NoError(t, err) + assert.Equal(t, "fallback audit reply", result.Outbound.Content) + require.Len(t, fallbackAudit.Records(), 1) + assert.Equal(t, "completed", fallbackAudit.Records()[0].Decision) +} + +func TestServiceHandleInboundUsesRuntimeAuditForBindingRejection(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + runtimeAudit := platform.NewInMemoryAuditSink() + fallbackAudit := platform.NewInMemoryAuditSink() + runtime := validRuntime( + "tenant-a", + &recordingRunner{response: "unused"}, + ) + runtime.Binding.AllowedUsers = []string{"allowed-user"} + runtime.Audit = runtimeAudit + require.NoError(t, registry.Register(runtime)) + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + WithAuditSink(fallbackAudit), + ) + + _, err := svc.HandleInbound( + ctx, + inbound("tenant-a", "msg-binding-reject", "denied-user", "hello"), + ) + require.ErrorIs(t, err, ErrBindingAccessDenied) + require.Len(t, runtimeAudit.Records(), 1) + assert.Equal(t, "reject", runtimeAudit.Records()[0].Decision) + assert.Empty(t, fallbackAudit.Records()) +} + func TestServiceHandleInboundDuplicateReusesOutboxBackedResult(t *testing.T) { ctx := context.Background() registry := NewInMemoryRegistry() diff --git a/platform/worker/builder.go b/platform/worker/builder.go new file mode 100644 index 0000000000..2576f4efe5 --- /dev/null +++ b/platform/worker/builder.go @@ -0,0 +1,219 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package worker + +import ( + "context" + "fmt" + "reflect" + "strings" + + "trpc.group/trpc-go/trpc-agent-go/agent" + "trpc.group/trpc-go/trpc-agent-go/artifact" + "trpc.group/trpc-go/trpc-agent-go/knowledge" + "trpc.group/trpc-go/trpc-agent-go/memory" + "trpc.group/trpc-go/trpc-agent-go/platform" + "trpc.group/trpc-go/trpc-agent-go/platform/gateway" + "trpc.group/trpc-go/trpc-agent-go/platform/storagerouter" + "trpc.group/trpc-go/trpc-agent-go/runner" + "trpc.group/trpc-go/trpc-agent-go/session" +) + +// AgentDependencies contains tenant-scoped services available while building +// one runtime agent. +type AgentDependencies struct { + Tenant platform.Tenant + App platform.AgentApp + Binding platform.ChannelBinding + Storage storagerouter.StorageAdapter + Session session.Service + Memory memory.Service + Artifact artifact.Service + Knowledge knowledge.Knowledge + Audit platform.AuditSink +} + +// AgentFactory builds an agent for one tenant app runtime. +type AgentFactory interface { + BuildAgent(ctx context.Context, dependencies AgentDependencies) (agent.Agent, error) +} + +// AgentFactoryFunc adapts a function into an AgentFactory. +type AgentFactoryFunc func(context.Context, AgentDependencies) (agent.Agent, error) + +// BuildAgent implements AgentFactory. +func (f AgentFactoryFunc) BuildAgent( + ctx context.Context, + dependencies AgentDependencies, +) (agent.Agent, error) { + return f(ctx, dependencies) +} + +// RuntimeBuilder assembles gateway runtimes from tenant storage profiles. +type RuntimeBuilder struct { + router storagerouter.Router + factory AgentFactory +} + +// NewRuntimeBuilder creates a runtime builder. +func NewRuntimeBuilder( + router storagerouter.Router, + factory AgentFactory, +) (*RuntimeBuilder, error) { + if isNilDependency(router) { + return nil, ErrStorageRouterRequired + } + if isNilDependency(factory) { + return nil, ErrAgentFactoryRequired + } + return &RuntimeBuilder{ + router: router, + factory: factory, + }, nil +} + +// Build resolves tenant-scoped storage services, builds the configured agent, +// and injects Session, Memory, and Artifact services into a Runner. +func (b *RuntimeBuilder) Build( + ctx context.Context, + tenant platform.Tenant, + app platform.AgentApp, + binding platform.ChannelBinding, +) (gateway.Runtime, error) { + if err := ctx.Err(); err != nil { + return gateway.Runtime{}, err + } + if err := validateRuntimeConfig(tenant, app, binding); err != nil { + return gateway.Runtime{}, err + } + + storage, err := b.router.Adapter(ctx, tenant.TenantID, app.StorageProfileID) + if err != nil { + return gateway.Runtime{}, fmt.Errorf("resolve storage adapter: %w", err) + } + sessionService, err := storage.Session(ctx) + if err != nil { + return gateway.Runtime{}, fmt.Errorf("resolve session service: %w", err) + } + memoryService, err := storage.Memory(ctx) + if err != nil { + return gateway.Runtime{}, fmt.Errorf("resolve memory service: %w", err) + } + artifactService, err := storage.Artifact(ctx) + if err != nil { + return gateway.Runtime{}, fmt.Errorf("resolve artifact service: %w", err) + } + knowledgeService, err := storage.Knowledge(ctx) + if err != nil { + return gateway.Runtime{}, fmt.Errorf("resolve knowledge service: %w", err) + } + auditSink, err := storage.Audit(ctx) + if err != nil { + return gateway.Runtime{}, fmt.Errorf("resolve audit sink: %w", err) + } + + dependencies := AgentDependencies{ + Tenant: tenant, + App: app, + Binding: binding, + Storage: storage, + Session: sessionService, + Memory: memoryService, + Artifact: artifactService, + Knowledge: knowledgeService, + Audit: auditSink, + } + ag, err := b.factory.BuildAgent(ctx, dependencies) + if err != nil { + return gateway.Runtime{}, fmt.Errorf("build agent: %w", err) + } + if isNilDependency(ag) { + return gateway.Runtime{}, ErrAgentRequired + } + if ag.Info().Name != app.AgentName { + return gateway.Runtime{}, ErrAgentNameMismatch + } + + runtime := gateway.Runtime{ + Tenant: tenant, + App: app, + Binding: binding, + Runner: runner.NewRunner( + storage.Scope().ScopedAppName(app.AppID), + ag, + runner.WithSessionService(sessionService), + runner.WithMemoryService(memoryService), + runner.WithArtifactService(artifactService), + ), + Audit: auditSink, + } + if err := runtime.Validate(); err != nil { + _ = runtime.Runner.Close() + return gateway.Runtime{}, err + } + return runtime, nil +} + +func validateRuntimeConfig( + tenant platform.Tenant, + app platform.AgentApp, + binding platform.ChannelBinding, +) error { + if err := tenant.Validate(); err != nil { + return err + } + if err := app.Validate(); err != nil { + return err + } + if err := binding.Validate(); err != nil { + return err + } + if app.TenantID != tenant.TenantID || + binding.TenantID != tenant.TenantID || + binding.AppID != app.AppID { + return ErrRuntimeIdentityMismatch + } + if tenant.Status != "" && tenant.Status != platform.TenantStatusActive { + return gateway.ErrRuntimeInactive + } + if app.Status != "" && app.Status != platform.AppStatusActive { + return gateway.ErrRuntimeInactive + } + if binding.Status != "" && binding.Status != platform.BindingStatusActive { + return gateway.ErrRuntimeInactive + } + if strings.TrimSpace(app.AppName) == "" { + return ErrAppNameRequired + } + if strings.TrimSpace(app.AgentName) == "" { + return ErrAgentNameRequired + } + if strings.TrimSpace(app.StorageProfileID) == "" { + return ErrStorageProfileIDRequired + } + return nil +} + +func isNilDependency(value any) bool { + if value == nil { + return true + } + reflected := reflect.ValueOf(value) + switch reflected.Kind() { + case reflect.Chan, + reflect.Func, + reflect.Interface, + reflect.Map, + reflect.Pointer, + reflect.Slice: + return reflected.IsNil() + default: + return false + } +} diff --git a/platform/worker/builder_test.go b/platform/worker/builder_test.go new file mode 100644 index 0000000000..a810353e8f --- /dev/null +++ b/platform/worker/builder_test.go @@ -0,0 +1,604 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package worker + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "trpc.group/trpc-go/trpc-agent-go/agent" + "trpc.group/trpc-go/trpc-agent-go/artifact" + "trpc.group/trpc-go/trpc-agent-go/event" + "trpc.group/trpc-go/trpc-agent-go/knowledge" + "trpc.group/trpc-go/trpc-agent-go/memory" + memoryinmemory "trpc.group/trpc-go/trpc-agent-go/memory/inmemory" + "trpc.group/trpc-go/trpc-agent-go/model" + "trpc.group/trpc-go/trpc-agent-go/platform" + "trpc.group/trpc-go/trpc-agent-go/platform/artifactstore" + "trpc.group/trpc-go/trpc-agent-go/platform/gateway" + "trpc.group/trpc-go/trpc-agent-go/platform/storagerouter" + "trpc.group/trpc-go/trpc-agent-go/session" + sessioninmemory "trpc.group/trpc-go/trpc-agent-go/session/inmemory" + "trpc.group/trpc-go/trpc-agent-go/tool" +) + +func TestRuntimeBuilderClosesGatewayStorageLoop(t *testing.T) { + ctx := context.Background() + tenantID := "tenant-a" + appID := "support-app" + profileID := "storage-a" + namespace := "tenant/" + tenantID + backendID := "backend-a" + + sessionService := sessioninmemory.NewSessionService() + t.Cleanup(func() { + require.NoError(t, sessionService.Close()) + }) + memoryService := memoryinmemory.NewMemoryService() + t.Cleanup(func() { + require.NoError(t, memoryService.Close()) + }) + metadataStore := artifactstore.NewInMemoryMetadataStore() + objectStore := artifactstore.NewInMemoryObjectStore() + artifactService, err := artifactstore.New(artifactstore.ServiceConfig{ + TenantID: tenantID, + Namespace: namespace, + MetadataStore: metadataStore, + ObjectStore: objectStore, + MaxAttempts: 2, + }) + require.NoError(t, err) + knowledgeService := &stubKnowledge{} + auditSink := platform.NewInMemoryAuditSink() + + router := storagerouter.NewInMemoryRouter() + require.NoError(t, router.RegisterBackend(storagerouter.BackendSet{ + TenantID: tenantID, + BackendID: backendID, + Session: sessionService, + Summary: sessionService, + Memory: memoryService, + Artifact: artifactService, + Knowledge: knowledgeService, + Audit: auditSink, + })) + require.NoError(t, router.RegisterProfile(platform.StorageProfile{ + TenantID: tenantID, + ProfileID: profileID, + SessionBackend: backendID, + MemoryBackend: backendID, + SummaryBackend: backendID, + ArtifactBackend: backendID, + KnowledgeBackend: backendID, + AuditBackend: backendID, + DSNRef: "secret://storage/" + tenantID, + Namespace: namespace, + })) + + tenant := platform.Tenant{ + TenantID: tenantID, + Status: platform.TenantStatusActive, + } + app := platform.AgentApp{ + TenantID: tenantID, + AppID: appID, + AppName: "support", + AgentName: "storage-probe", + StorageProfileID: profileID, + Status: platform.AppStatusActive, + } + binding := platform.ChannelBinding{ + TenantID: tenantID, + AppID: appID, + BindingID: "binding-a", + Channel: "wecom", + AccountID: "account-a", + WebhookPath: "/channels/wecom/binding-a/callback", + TokenRef: "secret://channel/token-a", + SecretRef: "secret://channel/secret-a", + Status: platform.BindingStatusActive, + ChannelLimits: platform.ChannelLimits{MaxTextLength: 4096}, + } + + var captured AgentDependencies + builder, err := NewRuntimeBuilder( + router, + AgentFactoryFunc(func( + _ context.Context, + dependencies AgentDependencies, + ) (agent.Agent, error) { + captured = dependencies + return &storageProbeAgent{name: app.AgentName}, nil + }), + ) + require.NoError(t, err) + runtime, err := builder.Build(ctx, tenant, app, binding) + require.NoError(t, err) + t.Cleanup(func() { + require.NoError(t, runtime.Runner.Close()) + }) + + registry := gateway.NewInMemoryRegistry() + require.NoError(t, registry.Register(runtime)) + service := gateway.NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + gateway.NewInMemoryOutboundStore(), + ) + inbound := platform.InboundMessage{ + TenantID: tenantID, + AppID: appID, + BindingID: binding.BindingID, + Channel: binding.Channel, + ChannelAccountID: binding.AccountID, + PlatformMessageID: "message-a", + ExternalUserID: "external-user", + ConversationType: platform.ConversationTypeDM, + MessageType: platform.MessageTypeText, + ContentParts: []platform.ContentPart{ + {Type: platform.ContentPartTypeText, Text: "persist this"}, + }, + ReceivedAt: time.Unix(100, 0), + } + result, err := service.HandleInbound(ctx, inbound) + require.NoError(t, err) + assert.Equal(t, "stored", result.Outbound.Content) + auditRecords := auditSink.Records() + require.Len(t, auditRecords, 1) + assert.Equal(t, tenantID, auditRecords[0].TenantID) + assert.Equal(t, "completed", auditRecords[0].Decision) + + internalUserID := platform.InternalUserID( + tenantID, + binding.Channel, + inbound.ExternalUserID, + ) + scopedAppName := namespace + "/" + app.AppID + storedSession, err := sessionService.GetSession(ctx, session.Key{ + AppName: scopedAppName, + UserID: internalUserID, + SessionID: result.SessionID, + }) + require.NoError(t, err) + require.NotNil(t, storedSession) + assert.NotEmpty(t, storedSession.Events) + + memories, err := memoryService.ReadMemories(ctx, memory.UserKey{ + AppName: scopedAppName, + UserID: internalUserID, + }, 10) + require.NoError(t, err) + require.Len(t, memories, 1) + require.NotNil(t, memories[0].Memory) + assert.Equal(t, "persist this", memories[0].Memory.Memory) + + loadedArtifact, err := artifactService.LoadArtifact(ctx, artifact.SessionInfo{ + AppName: scopedAppName, + UserID: internalUserID, + SessionID: result.SessionID, + }, "result.txt", nil) + require.NoError(t, err) + require.NotNil(t, loadedArtifact) + assert.Equal(t, []byte("persist this"), loadedArtifact.Data) + + assert.Equal(t, tenantID, captured.Storage.Scope().TenantID) + assert.Equal(t, profileID, captured.Storage.Scope().ProfileID) + require.NotNil(t, captured.Session) + require.NotNil(t, captured.Memory) + require.NotNil(t, captured.Artifact) + require.NotNil(t, captured.Knowledge) + require.NotNil(t, captured.Audit) + + _, err = captured.Session.CreateSession(ctx, session.Key{ + AppName: app.AppName, + UserID: internalUserID, + SessionID: "unscoped-session", + }, nil) + assert.ErrorIs(t, err, storagerouter.ErrKeyOutsideTenantScope) + + err = captured.Memory.AddMemory(ctx, memory.UserKey{ + AppName: app.AppName, + UserID: internalUserID, + }, "unscoped", nil) + assert.ErrorIs(t, err, storagerouter.ErrKeyOutsideTenantScope) + + _, err = captured.Artifact.LoadArtifact(ctx, artifact.SessionInfo{ + AppName: app.AppName, + UserID: internalUserID, + SessionID: result.SessionID, + }, "result.txt", nil) + assert.ErrorIs(t, err, storagerouter.ErrKeyOutsideTenantScope) + + _, err = captured.Knowledge.Search(ctx, &knowledge.SearchRequest{ + Query: "deployment", + }) + require.NoError(t, err) + require.NotNil(t, knowledgeService.last) + assert.Equal( + t, + tenantID, + knowledgeService.last.SearchFilter.Metadata["tenant_id"], + ) + + err = captured.Audit.WriteAudit(ctx, platform.AuditRecord{ + TenantID: "tenant-b", + }) + assert.ErrorIs(t, err, storagerouter.ErrKeyOutsideTenantScope) +} + +func TestRuntimeBuilderRejectsIdentityMismatch(t *testing.T) { + router := storagerouter.NewInMemoryRouter() + builder, err := NewRuntimeBuilder( + router, + AgentFactoryFunc(func( + context.Context, + AgentDependencies, + ) (agent.Agent, error) { + return &storageProbeAgent{name: "wrong-agent"}, nil + }), + ) + require.NoError(t, err) + + tenant := platform.Tenant{ + TenantID: "tenant-a", + Status: platform.TenantStatusActive, + } + app := platform.AgentApp{ + TenantID: tenant.TenantID, + AppID: "app-a", + AppName: "app", + AgentName: "expected-agent", + StorageProfileID: "profile-a", + Status: platform.AppStatusActive, + } + binding := platform.ChannelBinding{ + TenantID: "tenant-b", + AppID: app.AppID, + BindingID: "binding-a", + Channel: "wecom", + AccountID: "account-a", + WebhookPath: "/callback", + TokenRef: "secret://token", + SecretRef: "secret://secret", + Status: platform.BindingStatusActive, + } + + _, err = builder.Build(context.Background(), tenant, app, binding) + assert.ErrorIs(t, err, ErrRuntimeIdentityMismatch) +} + +func TestRuntimeBuilderRejectsInactiveConfigBeforeFactory(t *testing.T) { + tests := []struct { + name string + mutate func(*platform.Tenant, *platform.AgentApp, *platform.ChannelBinding) + }{ + { + name: "tenant suspended", + mutate: func( + tenant *platform.Tenant, + _ *platform.AgentApp, + _ *platform.ChannelBinding, + ) { + tenant.Status = platform.TenantStatusSuspended + }, + }, + { + name: "app suspended", + mutate: func( + _ *platform.Tenant, + app *platform.AgentApp, + _ *platform.ChannelBinding, + ) { + app.Status = platform.AppStatusSuspended + }, + }, + { + name: "binding disabled", + mutate: func( + _ *platform.Tenant, + _ *platform.AgentApp, + binding *platform.ChannelBinding, + ) { + binding.Status = platform.BindingStatusDisabled + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + factoryCalled := false + router := &countingRouter{ + Router: storagerouter.NewInMemoryRouter(), + } + builder, err := NewRuntimeBuilder( + router, + AgentFactoryFunc(func( + context.Context, + AgentDependencies, + ) (agent.Agent, error) { + factoryCalled = true + return &storageProbeAgent{name: "agent"}, nil + }), + ) + require.NoError(t, err) + + tenant := platform.Tenant{ + TenantID: "tenant-a", + Status: platform.TenantStatusActive, + } + app := platform.AgentApp{ + TenantID: tenant.TenantID, + AppID: "app-a", + AppName: "app", + AgentName: "agent", + StorageProfileID: "profile-a", + Status: platform.AppStatusActive, + } + binding := platform.ChannelBinding{ + TenantID: tenant.TenantID, + AppID: app.AppID, + BindingID: "binding-a", + Channel: "wecom", + AccountID: "account-a", + WebhookPath: "/callback", + TokenRef: "secret://token", + SecretRef: "secret://secret", + Status: platform.BindingStatusActive, + } + tt.mutate(&tenant, &app, &binding) + + _, err = builder.Build(context.Background(), tenant, app, binding) + assert.ErrorIs(t, err, gateway.ErrRuntimeInactive) + assert.Zero(t, router.adapterCalls) + assert.False(t, factoryCalled) + }) + } +} + +func TestNewRuntimeBuilderRejectsTypedNilFactory(t *testing.T) { + var factory AgentFactoryFunc + + _, err := NewRuntimeBuilder(storagerouter.NewInMemoryRouter(), factory) + + assert.ErrorIs(t, err, ErrAgentFactoryRequired) +} + +func TestNewRuntimeBuilderRejectsTypedNilRouter(t *testing.T) { + var router *storagerouter.InMemoryRouter + + _, err := NewRuntimeBuilder( + router, + AgentFactoryFunc(func( + context.Context, + AgentDependencies, + ) (agent.Agent, error) { + return &storageProbeAgent{name: "agent"}, nil + }), + ) + + assert.ErrorIs(t, err, ErrStorageRouterRequired) +} + +func TestRuntimeBuilderRejectsInvalidAgent(t *testing.T) { + ctx := context.Background() + tenantID := "tenant-a" + profileID := "profile-a" + backendID := "backend-a" + namespace := "tenant/" + tenantID + + sessionService := sessioninmemory.NewSessionService() + t.Cleanup(func() { + require.NoError(t, sessionService.Close()) + }) + memoryService := memoryinmemory.NewMemoryService() + t.Cleanup(func() { + require.NoError(t, memoryService.Close()) + }) + artifactService, err := artifactstore.New(artifactstore.ServiceConfig{ + TenantID: tenantID, + Namespace: namespace, + MetadataStore: artifactstore.NewInMemoryMetadataStore(), + ObjectStore: artifactstore.NewInMemoryObjectStore(), + MaxAttempts: 2, + }) + require.NoError(t, err) + + router := storagerouter.NewInMemoryRouter() + require.NoError(t, router.RegisterBackend(storagerouter.BackendSet{ + TenantID: tenantID, + BackendID: backendID, + Session: sessionService, + Summary: sessionService, + Memory: memoryService, + Artifact: artifactService, + Knowledge: &stubKnowledge{}, + Audit: platform.NewInMemoryAuditSink(), + })) + require.NoError(t, router.RegisterProfile(platform.StorageProfile{ + TenantID: tenantID, + ProfileID: profileID, + SessionBackend: backendID, + MemoryBackend: backendID, + SummaryBackend: backendID, + ArtifactBackend: backendID, + KnowledgeBackend: backendID, + AuditBackend: backendID, + DSNRef: "secret://storage/" + tenantID, + Namespace: namespace, + })) + + tenant := platform.Tenant{ + TenantID: tenantID, + Status: platform.TenantStatusActive, + } + app := platform.AgentApp{ + TenantID: tenantID, + AppID: "app-a", + AppName: "app", + AgentName: "expected-agent", + StorageProfileID: profileID, + Status: platform.AppStatusActive, + } + binding := platform.ChannelBinding{ + TenantID: tenantID, + AppID: app.AppID, + BindingID: "binding-a", + Channel: "wecom", + AccountID: "account-a", + WebhookPath: "/callback", + TokenRef: "secret://token", + SecretRef: "secret://secret", + Status: platform.BindingStatusActive, + } + + t.Run("name mismatch", func(t *testing.T) { + builder, err := NewRuntimeBuilder( + router, + AgentFactoryFunc(func( + context.Context, + AgentDependencies, + ) (agent.Agent, error) { + return &storageProbeAgent{name: "wrong-agent"}, nil + }), + ) + require.NoError(t, err) + + _, err = builder.Build(ctx, tenant, app, binding) + assert.ErrorIs(t, err, ErrAgentNameMismatch) + }) + + t.Run("typed nil", func(t *testing.T) { + var nilAgent *storageProbeAgent + builder, err := NewRuntimeBuilder( + router, + AgentFactoryFunc(func( + context.Context, + AgentDependencies, + ) (agent.Agent, error) { + return nilAgent, nil + }), + ) + require.NoError(t, err) + + _, err = builder.Build(ctx, tenant, app, binding) + assert.ErrorIs(t, err, ErrAgentRequired) + }) +} + +type storageProbeAgent struct { + name string +} + +func (a *storageProbeAgent) Run( + ctx context.Context, + invocation *agent.Invocation, +) (<-chan *event.Event, error) { + userKey := memory.UserKey{ + AppName: invocation.Session.AppName, + UserID: invocation.Session.UserID, + } + if err := invocation.MemoryService.AddMemory( + ctx, + userKey, + invocation.Message.Content, + []string{"gateway"}, + ); err != nil { + return nil, err + } + if _, err := invocation.ArtifactService.SaveArtifact( + ctx, + artifact.SessionInfo{ + AppName: invocation.Session.AppName, + UserID: invocation.Session.UserID, + SessionID: invocation.Session.ID, + }, + "result.txt", + &artifact.Artifact{ + Data: []byte(invocation.Message.Content), + MimeType: "text/plain", + Name: "result.txt", + }, + ); err != nil { + return nil, err + } + + out := make(chan *event.Event, 1) + out <- event.NewResponseEvent( + invocation.InvocationID, + a.name, + &model.Response{ + ID: "storage-probe-response", + Object: model.ObjectTypeChatCompletion, + Done: true, + Choices: []model.Choice{ + { + Index: 0, + Message: model.Message{ + Role: model.RoleAssistant, + Content: "stored", + }, + }, + }, + }, + ) + close(out) + return out, nil +} + +func (a *storageProbeAgent) Tools() []tool.Tool { + return nil +} + +func (a *storageProbeAgent) Info() agent.Info { + return agent.Info{Name: a.name} +} + +func (a *storageProbeAgent) SubAgents() []agent.Agent { + return nil +} + +func (a *storageProbeAgent) FindSubAgent(string) agent.Agent { + return nil +} + +type stubKnowledge struct { + last *knowledge.SearchRequest +} + +func (s *stubKnowledge) Search( + ctx context.Context, + req *knowledge.SearchRequest, +) (*knowledge.SearchResult, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + if req != nil { + copied := *req + s.last = &copied + } + return &knowledge.SearchResult{}, nil +} + +type countingRouter struct { + storagerouter.Router + adapterCalls int +} + +func (r *countingRouter) Adapter( + ctx context.Context, + tenantID string, + profileID string, +) (storagerouter.StorageAdapter, error) { + r.adapterCalls++ + return r.Router.Adapter(ctx, tenantID, profileID) +} diff --git a/platform/worker/doc.go b/platform/worker/doc.go new file mode 100644 index 0000000000..1b1e3efd9e --- /dev/null +++ b/platform/worker/doc.go @@ -0,0 +1,10 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +// Package worker assembles tenant-scoped platform runtimes. +package worker diff --git a/platform/worker/errors.go b/platform/worker/errors.go new file mode 100644 index 0000000000..b650e54c64 --- /dev/null +++ b/platform/worker/errors.go @@ -0,0 +1,30 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package worker + +import "errors" + +var ( + // ErrStorageRouterRequired indicates that runtime storage cannot be resolved. + ErrStorageRouterRequired = errors.New("worker storage router is required") + // ErrAgentFactoryRequired indicates that no tenant agent factory was configured. + ErrAgentFactoryRequired = errors.New("worker agent factory is required") + // ErrAppNameRequired indicates that the runtime app has no storage app name. + ErrAppNameRequired = errors.New("worker app_name is required") + // ErrAgentNameRequired indicates that the runtime app has no agent identity. + ErrAgentNameRequired = errors.New("worker agent_name is required") + // ErrStorageProfileIDRequired indicates that the app has no storage profile. + ErrStorageProfileIDRequired = errors.New("worker storage_profile_id is required") + // ErrRuntimeIdentityMismatch indicates that tenant, app, and binding disagree. + ErrRuntimeIdentityMismatch = errors.New("worker runtime identity mismatch") + // ErrAgentRequired indicates that the factory returned no agent. + ErrAgentRequired = errors.New("worker agent is required") + // ErrAgentNameMismatch indicates that the built agent does not match app config. + ErrAgentNameMismatch = errors.New("worker agent name does not match app config") +) From f9a509b5c1f5bfd04909ee667feef39267deddf3 Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 08:10:33 +0800 Subject: [PATCH 57/95] feat(platform): enforce runtime tool governance --- agent/invocation.go | 86 +++- agent/invocation_surface.go | 18 + agent/invocation_test.go | 68 +++ agent/llmagent/llm_agent.go | 8 +- agent/llmagent/surface_runtime_test.go | 67 +++ agent/llmagent/tool_activation.go | 17 + graph/state_graph.go | 113 ++++- graph/state_graph_test.go | 197 ++++++++- graph/surface_runtime_test.go | 85 ++++ internal/flow/llmflow/llmflow.go | 51 ++- internal/flow/llmflow/llmflow_test.go | 91 ++++ internal/flow/llmflow/tool_filter_test.go | 96 +++++ internal/flow/processor/functioncall.go | 56 ++- internal/flow/processor/functioncall_test.go | 135 ++++++ internal/toolsurface/toolsurface.go | 104 ++++- internal/toolsurface/toolsurface_test.go | 166 ++++++++ platform/gateway/errors.go | 2 + platform/gateway/registry.go | 10 + platform/gateway/registry_test.go | 34 ++ platform/gateway/service.go | 33 +- platform/toolpolicy/filter.go | 57 +++ platform/toolpolicy/filter_test.go | 80 ++++ platform/toolpolicy/policy.go | 13 + platform/toolpolicy/policy_test.go | 30 ++ platform/worker/builder.go | 70 +++- platform/worker/errors.go | 4 + platform/worker/governance.go | 104 +++++ platform/worker/governance_test.go | 420 +++++++++++++++++++ tool/agent/agent_tool.go | 18 +- tool/agent/agent_tool_test.go | 22 + tool/agent/dynamic_tool.go | 5 +- tool/agent/dynamic_tool_test.go | 305 +++++++++++++- tool/dynamicworkflow/tool.go | 63 ++- tool/dynamicworkflow/tool_test.go | 110 ++++- 34 files changed, 2669 insertions(+), 69 deletions(-) create mode 100644 platform/toolpolicy/filter.go create mode 100644 platform/toolpolicy/filter_test.go create mode 100644 platform/worker/governance.go create mode 100644 platform/worker/governance_test.go diff --git a/agent/invocation.go b/agent/invocation.go index d31291063b..311ec47800 100644 --- a/agent/invocation.go +++ b/agent/invocation.go @@ -922,6 +922,16 @@ func WithToolFilter(filter tool.FilterFunc) RunOption { } } +// WithMandatoryToolFilter sets a non-negotiable tool visibility boundary for +// this run. Unlike WithToolFilter, it applies to the complete invocation tool +// surface, including framework-managed tools, and is preserved across derived +// child invocations. +func WithMandatoryToolFilter(filter tool.FilterFunc) RunOption { + return func(opts *RunOptions) { + opts.MandatoryToolFilter = filter + } +} + // WithAdditionalTools appends tools that are visible only for this run. // // Additional tools are treated as user tools, so WithToolFilter can still @@ -997,6 +1007,71 @@ func WithToolPermissionPolicyFunc(fn tool.PermissionPolicyFunc) RunOption { return WithToolPermissionPolicy(fn) } +// WithMandatoryToolPermissionPolicy sets a non-negotiable permission policy +// that derived child invocations must preserve. It is checked before the +// ordinary per-run ToolPermissionPolicy. +func WithMandatoryToolPermissionPolicy(policy tool.PermissionPolicy) RunOption { + return func(opts *RunOptions) { + opts.MandatoryToolPermissionPolicy = policy + } +} + +// WithMandatoryToolPermissionPolicyFunc adapts fn into a mandatory per-run +// tool permission policy. +func WithMandatoryToolPermissionPolicyFunc(fn tool.PermissionPolicyFunc) RunOption { + return WithMandatoryToolPermissionPolicy(fn) +} + +// CheckToolPermission applies the non-negotiable policy followed by the +// ordinary per-run policy. The first non-allow decision terminates the chain. +func (opts *RunOptions) CheckToolPermission( + ctx context.Context, + req *tool.PermissionRequest, +) (tool.PermissionDecision, error) { + if opts == nil { + return tool.AllowPermission(), nil + } + policies := [...]tool.PermissionPolicy{ + opts.MandatoryToolPermissionPolicy, + opts.ToolPermissionPolicy, + } + for _, policy := range policies { + if isNilToolPermissionPolicy(policy) { + continue + } + decision, err := policy.CheckToolPermission(ctx, req) + if err != nil { + return tool.PermissionDecision{}, err + } + decision, err = tool.NormalizePermissionDecision(decision) + if err != nil { + return tool.PermissionDecision{}, err + } + if decision.Action != tool.PermissionActionAllow { + return decision, nil + } + } + return tool.AllowPermission(), nil +} + +func isNilToolPermissionPolicy(policy tool.PermissionPolicy) bool { + if policy == nil { + return true + } + value := reflect.ValueOf(policy) + switch value.Kind() { + case reflect.Chan, + reflect.Func, + reflect.Interface, + reflect.Map, + reflect.Pointer, + reflect.Slice: + return value.IsNil() + default: + return false + } +} + func appendRunTools(opts *RunOptions, tools []tool.Tool) { if opts == nil || len(tools) == 0 { return @@ -1336,7 +1411,12 @@ type RunOptions struct { // StructuredOutputType is the Go type to unmarshal the final JSON into for this run. StructuredOutputType reflect.Type - // ToolFilter is a custom function to filter tools for this run. + // MandatoryToolFilter is a non-negotiable visibility boundary applied to + // the complete invocation tool surface, including framework-managed tools. + // Derived child invocations must preserve it. + MandatoryToolFilter tool.FilterFunc + + // ToolFilter is a custom function to filter user tools for this run. // If set, only tools for which the filter returns true will be available to the model. // If nil, all registered tools will be available (default behavior). // @@ -1389,6 +1469,10 @@ type RunOptions struct { // externally and later provide tool results (RoleTool messages). ToolExecutionFilter tool.FilterFunc + // MandatoryToolPermissionPolicy is checked before ToolPermissionPolicy and + // is preserved across derived child invocations. + MandatoryToolPermissionPolicy tool.PermissionPolicy + // ToolPermissionPolicy checks whether a tool call may run after the model // has requested it, after argument repair, and after before-tool callbacks // have finalized arguments. diff --git a/agent/invocation_surface.go b/agent/invocation_surface.go index fe3ec659b4..b23b8f12e4 100644 --- a/agent/invocation_surface.go +++ b/agent/invocation_surface.go @@ -36,6 +36,24 @@ type InvocationToolSurfaceProvider interface { ) ([]tool.Tool, map[string]bool) } +// InvocationToolActivationProvider is an optional interface implemented by +// agents that apply invocation-scoped activation after run-option tools have +// been appended to the base surface. +// +// The provider must return the activated tool surface together with updated +// user and external tool classifications. Callers provide private slice/map +// copies, so implementations may mutate the inputs without affecting the +// invocation's configured surface. +type InvocationToolActivationProvider interface { + ApplyInvocationToolActivation( + ctx context.Context, + inv *Invocation, + tools []tool.Tool, + userToolNames map[string]bool, + externalToolNames map[string]bool, + ) ([]tool.Tool, map[string]bool, map[string]bool) +} + // InvocationSkillRepositoryProvider is an optional interface implemented by // agents that can expose the effective, invocation-scoped skill repository. // diff --git a/agent/invocation_test.go b/agent/invocation_test.go index ffa7e24e1c..e1aa11abbe 100644 --- a/agent/invocation_test.go +++ b/agent/invocation_test.go @@ -1229,6 +1229,74 @@ func TestWithToolPermissionPolicy(t *testing.T) { require.Equal(t, tool.PermissionActionDeny, decision.Action) } +func TestRunOptionsCheckToolPermissionAppliesMandatoryPolicyFirst( + t *testing.T, +) { + var calls []string + opts := NewRunOptions( + WithMandatoryToolPermissionPolicyFunc( + func( + context.Context, + *tool.PermissionRequest, + ) (tool.PermissionDecision, error) { + calls = append(calls, "mandatory") + return tool.DenyPermission("tenant policy"), nil + }, + ), + WithToolPermissionPolicyFunc( + func( + context.Context, + *tool.PermissionRequest, + ) (tool.PermissionDecision, error) { + calls = append(calls, "ordinary") + return tool.AllowPermission(), nil + }, + ), + ) + + decision, err := opts.CheckToolPermission( + context.Background(), + &tool.PermissionRequest{ToolName: "shell"}, + ) + require.NoError(t, err) + require.Equal(t, tool.PermissionActionDeny, decision.Action) + require.Equal(t, []string{"mandatory"}, calls) +} + +func TestRunOptionsCheckToolPermissionAllowsOrdinaryPolicyToTighten( + t *testing.T, +) { + var calls []string + opts := NewRunOptions( + WithMandatoryToolPermissionPolicyFunc( + func( + context.Context, + *tool.PermissionRequest, + ) (tool.PermissionDecision, error) { + calls = append(calls, "mandatory") + return tool.AllowPermission(), nil + }, + ), + WithToolPermissionPolicyFunc( + func( + context.Context, + *tool.PermissionRequest, + ) (tool.PermissionDecision, error) { + calls = append(calls, "ordinary") + return tool.DenyPermission("child policy"), nil + }, + ), + ) + + decision, err := opts.CheckToolPermission( + context.Background(), + &tool.PermissionRequest{ToolName: "shell"}, + ) + require.NoError(t, err) + require.Equal(t, tool.PermissionActionDeny, decision.Action) + require.Equal(t, []string{"mandatory", "ordinary"}, calls) +} + func TestWithInstruction(t *testing.T) { opts := &RunOptions{} WithInstruction(testRunInstruction)(opts) diff --git a/agent/llmagent/llm_agent.go b/agent/llmagent/llm_agent.go index ce7ec50266..f4c9d5fc42 100644 --- a/agent/llmagent/llm_agent.go +++ b/agent/llmagent/llm_agent.go @@ -1645,8 +1645,12 @@ func (a *LLMAgent) resolveBaseModel(inv *agent.Invocation) baseModelResolution { // setupInvocation sets up the invocation. func (a *LLMAgent) setupInvocation(invocation *agent.Invocation) { // Set agent identity before resolving node-scoped surfaces. - invocation.Agent = a - invocation.AgentName = a.name + if invocation.Agent != a { + invocation.Agent = a + } + if invocation.AgentName != a.name { + invocation.AgentName = a.name + } // Set the base model once for compatibility with existing callbacks. resolution := a.resolveBaseModel(invocation) diff --git a/agent/llmagent/surface_runtime_test.go b/agent/llmagent/surface_runtime_test.go index 120e01cf24..22b782beba 100644 --- a/agent/llmagent/surface_runtime_test.go +++ b/agent/llmagent/surface_runtime_test.go @@ -670,6 +670,73 @@ func TestLLMAgent_Run_AgentToolFilterStillAppliesWithInvocationToolSurface( require.Contains(t, m.got.Tools, testTransferToolName) } +func TestLLMAgent_Run_MandatoryToolFilterAppliesToFrameworkTools( + t *testing.T, +) { + m := &captureModel{} + agt := New( + "test-agent", + WithModel(m), + WithTools([]tool.Tool{ + dummyTool{decl: &tool.Declaration{Name: "allowed_user_tool"}}, + }), + WithSubAgents([]agent.Agent{&mockAgent{name: "child"}}), + WithAwaitUserReplyTool(true), + ) + inv := agent.NewInvocation( + agent.WithInvocationMessage(model.NewUserMessage("hello")), + agent.WithInvocationRunOptions(agent.NewRunOptions( + agent.WithMandatoryToolFilter( + func(_ context.Context, tl tool.Tool) bool { + switch tl.Declaration().Name { + case testTransferToolName, testAwaitReplyToolName: + return false + default: + return true + } + }, + ), + )), + ) + + ch, err := agt.Run(context.Background(), inv) + require.NoError(t, err) + for range ch { + } + + require.NotNil(t, m.got) + require.Contains(t, m.got.Tools, "allowed_user_tool") + require.NotContains(t, m.got.Tools, testTransferToolName) + require.NotContains(t, m.got.Tools, testAwaitReplyToolName) +} + +func TestLLMAgent_SetupInvocationDoesNotRewritePreinitializedIdentity( + t *testing.T, +) { + agt := New("test-agent", WithModel(&captureModel{})) + inv := agent.NewInvocation(agent.WithInvocationAgent(agt)) + const iterations = 10000 + start := make(chan struct{}) + done := make(chan struct{}) + go func() { + <-start + for i := 0; i < iterations; i++ { + _ = inv.Agent + _ = inv.AgentName + } + close(done) + }() + + close(start) + for i := 0; i < iterations; i++ { + agt.setupInvocation(inv) + } + <-done + + require.Same(t, agt, inv.Agent) + require.Equal(t, "test-agent", inv.AgentName) +} + func TestLLMAgent_Run_SurfacePatch_OverridesToolDeclarations(t *testing.T) { m := &captureModel{} agt := New( diff --git a/agent/llmagent/tool_activation.go b/agent/llmagent/tool_activation.go index d7c848bd7b..c6c842e58c 100644 --- a/agent/llmagent/tool_activation.go +++ b/agent/llmagent/tool_activation.go @@ -332,6 +332,23 @@ func (a *LLMAgent) applyToolActivation( ) } +// ApplyInvocationToolActivation implements agent.InvocationToolActivationProvider. +func (a *LLMAgent) ApplyInvocationToolActivation( + ctx context.Context, + inv *agent.Invocation, + tools []tool.Tool, + userToolNames map[string]bool, + externalToolNames map[string]bool, +) ([]tool.Tool, map[string]bool, map[string]bool) { + return a.applyToolActivation( + ctx, + inv, + tools, + userToolNames, + externalToolNames, + ) +} + func (a *LLMAgent) toolActivationInputs() ( []tool.ToolSet, []toolActivationRule, diff --git a/graph/state_graph.go b/graph/state_graph.go index 90e7ae4864..e5450ba16e 100644 --- a/graph/state_graph.go +++ b/graph/state_graph.go @@ -1599,6 +1599,7 @@ func (r *llmRunner) executeModel( Tools: tools, GenerationConfig: r.generationConfig, } + applyMandatoryRequestToolFilter(ctx, callInvocation, request) // Sanitize invalid tool calls in history to avoid poisoning future requests. request.Messages = toolcall.SanitizeMessagesWithTools(ctx, request.Messages, request.Tools) applyInvocationRequestOverrides(request, callInvocation, nodeID) @@ -2283,6 +2284,14 @@ func runModelStream( } return ctx, singleResponseStream(customResponse), nil } + applyMandatoryRequestToolFilter(ctx, invocation, request) + if request != nil { + request.Messages = toolcall.SanitizeMessagesWithTools( + ctx, + request.Messages, + request.Tools, + ) + } if beforeGenerate != nil { beforeGenerate(ctx) } @@ -2290,6 +2299,39 @@ func runModelStream( return ctx, stream, err } +func applyMandatoryRequestToolFilter( + ctx context.Context, + invocation *agent.Invocation, + request *model.Request, +) { + if invocation == nil || + invocation.RunOptions.MandatoryToolFilter == nil || + request == nil || + len(request.Tools) == 0 { + return + } + for name, candidate := range request.Tools { + if graphToolName(candidate) == "" || + !invocation.RunOptions.MandatoryToolFilter( + ctx, + itool.ResolveDeclaration(candidate), + ) { + delete(request.Tools, name) + } + } +} + +func graphToolName(tl tool.Tool) string { + if tl == nil { + return "" + } + decl := tl.Declaration() + if decl == nil { + return "" + } + return decl.Name +} + // runModel preserves the pre-refactor test-facing helper signature by // adapting iterator-based model streams back to the legacy channel form. func runModel( @@ -4364,6 +4406,24 @@ func runToolWithEventContexts( } decl := t.Declaration() startInvocation := invocationFromContextOrFallback(ctx, nil) + var mandatoryToolPermissionPolicy tool.PermissionPolicy + if startInvocation != nil { + mandatoryToolPermissionPolicy = + startInvocation.RunOptions.MandatoryToolPermissionPolicy + } + visibilityResult, err := checkMandatoryToolVisibility( + ctx, + startInvocation, + toolCall, + t, + decl, + ) + if err != nil { + return ctx, startInvocation, ctx, startInvocation, nil, toolCall.Function.Arguments, err + } + if visibilityResult != nil { + return ctx, startInvocation, ctx, startInvocation, visibilityResult, toolCall.Function.Arguments, nil + } ctx, toolCall, customResult, err := runBeforeToolPluginCallbacks( ctx, @@ -4395,6 +4455,7 @@ func runToolWithEventContexts( } permissionResult, err := checkToolPermission( ctx, + mandatoryToolPermissionPolicy, startInvocation, toolCall, t, @@ -4561,8 +4622,45 @@ func callToolWithRetry( return runResult.Result, runResult.Error } +func checkMandatoryToolVisibility( + ctx context.Context, + invocation *agent.Invocation, + toolCall model.ToolCall, + t tool.Tool, + decl *tool.Declaration, +) (*tool.PermissionResult, error) { + if invocation == nil || invocation.RunOptions.MandatoryToolFilter == nil { + return nil, nil + } + if invocation.RunOptions.MandatoryToolFilter( + ctx, + itool.ResolveDeclaration(t), + ) { + return nil, nil + } + req := &tool.PermissionRequest{ + Tool: t, + ToolName: toolCall.Function.Name, + ToolCallID: toolCall.ID, + Declaration: decl, + Arguments: toolCall.Function.Arguments, + Metadata: tool.MetadataOf(itool.ResolveSemantic(t)), + } + return normalizeToolPermissionResult( + req, + tool.DenyPermission( + fmt.Sprintf( + "tool %q is hidden by mandatory tool filter", + req.ToolName, + ), + ), + nil, + ) +} + func checkToolPermission( ctx context.Context, + mandatoryPolicy tool.PermissionPolicy, invocation *agent.Invocation, toolCall model.ToolCall, t tool.Tool, @@ -4583,10 +4681,21 @@ func checkToolPermission( return result, err } } - if invocation == nil || invocation.RunOptions.ToolPermissionPolicy == nil { + mandatoryOpts := agent.RunOptions{ + MandatoryToolPermissionPolicy: mandatoryPolicy, + } + decision, err := mandatoryOpts.CheckToolPermission(ctx, req) + result, err := normalizeToolPermissionResult(req, decision, err) + if result != nil || err != nil { + return result, err + } + if invocation == nil { return nil, nil } - decision, err := invocation.RunOptions.ToolPermissionPolicy.CheckToolPermission(ctx, req) + ordinaryOpts := agent.RunOptions{ + ToolPermissionPolicy: invocation.RunOptions.ToolPermissionPolicy, + } + decision, err = ordinaryOpts.CheckToolPermission(ctx, req) return normalizeToolPermissionResult(req, decision, err) } diff --git a/graph/state_graph_test.go b/graph/state_graph_test.go index 569d43ae91..644d8a8b93 100644 --- a/graph/state_graph_test.go +++ b/graph/state_graph_test.go @@ -1788,7 +1788,7 @@ func TestRunToolWithEventContexts_OrdinaryToolIgnoresAgentToolInterruptState(t * require.Equal(t, toolCall.Function.Arguments, modifiedArgs) } -func TestRunToolWithEventContexts_ToolPermissionPolicyDenySkipsExecution( +func TestRunToolWithEventContexts_MandatoryToolPermissionPolicyDenySkipsExecution( t *testing.T, ) { const ( @@ -1800,9 +1800,10 @@ func TestRunToolWithEventContexts_ToolPermissionPolicyDenySkipsExecution( ) var ( - beforeCalled bool - afterCalled bool - policyCalled bool + beforeCalled bool + afterCalled bool + mandatoryCalled bool + ordinaryCalled bool ) callbacks := tool.NewCallbacks() callbacks.RegisterBeforeTool(func( @@ -1822,15 +1823,23 @@ func TestRunToolWithEventContexts_ToolPermissionPolicyDenySkipsExecution( return &tool.AfterToolResult{}, nil }) invocation := &agent.Invocation{ - RunOptions: agent.NewRunOptions(agent.WithToolPermissionPolicyFunc( - func(_ context.Context, req *tool.PermissionRequest) (tool.PermissionDecision, error) { - policyCalled = true - require.Equal(t, toolName, req.ToolName) - require.Equal(t, toolCallID, req.ToolCallID) - require.JSONEq(t, rewrittenArgs, string(req.Arguments)) - return tool.DenyPermission(denyReason), nil - }, - )), + RunOptions: agent.NewRunOptions( + agent.WithMandatoryToolPermissionPolicyFunc( + func(_ context.Context, req *tool.PermissionRequest) (tool.PermissionDecision, error) { + mandatoryCalled = true + require.Equal(t, toolName, req.ToolName) + require.Equal(t, toolCallID, req.ToolCallID) + require.JSONEq(t, rewrittenArgs, string(req.Arguments)) + return tool.DenyPermission(denyReason), nil + }, + ), + agent.WithToolPermissionPolicyFunc( + func(context.Context, *tool.PermissionRequest) (tool.PermissionDecision, error) { + ordinaryCalled = true + return tool.AllowPermission(), nil + }, + ), + ), } ctx := agent.NewInvocationContext(context.Background(), invocation) tl := &captureTool{name: toolName, result: map[string]any{"ok": true}} @@ -1853,7 +1862,8 @@ func TestRunToolWithEventContexts_ToolPermissionPolicyDenySkipsExecution( ) require.NoError(t, err) require.True(t, beforeCalled) - require.True(t, policyCalled) + require.True(t, mandatoryCalled) + require.False(t, ordinaryCalled) require.False(t, afterCalled) require.False(t, tl.called) require.JSONEq(t, rewrittenArgs, string(modifiedArgs)) @@ -1864,6 +1874,165 @@ func TestRunToolWithEventContexts_ToolPermissionPolicyDenySkipsExecution( require.Equal(t, denyReason, permissionResult.Reason) } +func TestRunToolWithEventContexts_MandatoryToolPermissionPolicySurvivesCallbackInvocationReplacement( + t *testing.T, +) { + const ( + toolName = "delete_file" + toolCallID = "call-deny" + denyReason = "tenant policy" + originalArgs = `{"path":"unsafe"}` + rewrittenArgs = `{"path":"safe"}` + ) + + var ( + mandatoryCalled bool + ordinaryCalled bool + afterCalled bool + ) + callbackInvocation := agent.NewInvocation( + agent.WithInvocationID("callback-invocation"), + agent.WithInvocationRunOptions(agent.NewRunOptions( + agent.WithToolPermissionPolicyFunc( + func(context.Context, *tool.PermissionRequest) (tool.PermissionDecision, error) { + ordinaryCalled = true + return tool.AllowPermission(), nil + }, + ), + )), + ) + callbacks := tool.NewCallbacks() + callbacks.RegisterBeforeTool(func( + _ context.Context, + _ *tool.BeforeToolArgs, + ) (*tool.BeforeToolResult, error) { + return &tool.BeforeToolResult{ + Context: agent.NewInvocationContext( + context.Background(), + callbackInvocation, + ), + ModifiedArguments: []byte(rewrittenArgs), + }, nil + }) + callbacks.RegisterAfterTool(func( + _ context.Context, + _ *tool.AfterToolArgs, + ) (*tool.AfterToolResult, error) { + afterCalled = true + return &tool.AfterToolResult{}, nil + }) + originalInvocation := agent.NewInvocation( + agent.WithInvocationRunOptions(agent.NewRunOptions( + agent.WithMandatoryToolPermissionPolicyFunc( + func( + _ context.Context, + req *tool.PermissionRequest, + ) (tool.PermissionDecision, error) { + mandatoryCalled = true + require.Equal(t, toolName, req.ToolName) + require.Equal(t, toolCallID, req.ToolCallID) + require.JSONEq(t, rewrittenArgs, string(req.Arguments)) + return tool.DenyPermission(denyReason), nil + }, + ), + )), + ) + ctx := agent.NewInvocationContext(context.Background(), originalInvocation) + tl := &captureTool{name: toolName, result: map[string]any{"ok": true}} + toolCall := model.ToolCall{ + ID: toolCallID, + Function: model.FunctionDefinitionParam{ + Name: toolName, + Arguments: []byte(originalArgs), + }, + } + + _, startInvocation, _, _, result, modifiedArgs, err := + runToolWithEventContexts( + ctx, + toolCall, + callbacks, + tl, + State{}, + nil, + 0, + ) + require.NoError(t, err) + require.Same(t, callbackInvocation, startInvocation) + require.True(t, mandatoryCalled) + require.False(t, ordinaryCalled) + require.False(t, afterCalled) + require.False(t, tl.called) + require.JSONEq(t, rewrittenArgs, string(modifiedArgs)) + permissionResult, ok := result.(*tool.PermissionResult) + require.True(t, ok) + require.Equal(t, tool.PermissionResultStatusDenied, permissionResult.Status) + require.Equal(t, toolName, permissionResult.Tool) + require.Equal(t, denyReason, permissionResult.Reason) +} + +func TestRunToolWithEventContexts_MandatoryToolFilterDenySkipsCallbacksAndExecution( + t *testing.T, +) { + const toolName = "hidden_tool" + var ( + beforeCalled bool + permissionCalled bool + ) + callbacks := tool.NewCallbacks() + callbacks.RegisterBeforeTool(func( + _ context.Context, + _ *tool.BeforeToolArgs, + ) (*tool.BeforeToolResult, error) { + beforeCalled = true + return &tool.BeforeToolResult{}, nil + }) + invocation := agent.NewInvocation( + agent.WithInvocationRunOptions(agent.NewRunOptions( + agent.WithMandatoryToolFilter( + func(_ context.Context, candidate tool.Tool) bool { + return candidate.Declaration().Name != toolName + }, + ), + agent.WithToolPermissionPolicyFunc( + func(context.Context, *tool.PermissionRequest) (tool.PermissionDecision, error) { + permissionCalled = true + return tool.AllowPermission(), nil + }, + ), + )), + ) + ctx := agent.NewInvocationContext(context.Background(), invocation) + tl := &captureTool{name: toolName, result: map[string]any{"ok": true}} + toolCall := model.ToolCall{ + ID: "call-hidden", + Function: model.FunctionDefinitionParam{ + Name: toolName, + Arguments: []byte(`{"value":"blocked"}`), + }, + } + + _, _, _, _, result, modifiedArgs, err := runToolWithEventContexts( + ctx, + toolCall, + callbacks, + tl, + State{}, + nil, + 0, + ) + require.NoError(t, err) + require.False(t, beforeCalled) + require.False(t, permissionCalled) + require.False(t, tl.called) + require.Equal(t, toolCall.Function.Arguments, modifiedArgs) + permissionResult, ok := result.(*tool.PermissionResult) + require.True(t, ok) + require.Equal(t, tool.PermissionResultStatusDenied, permissionResult.Status) + require.Equal(t, toolName, permissionResult.Tool) + require.Contains(t, permissionResult.Reason, "mandatory tool filter") +} + func TestNewToolsNodeFunc_ToolCallbacksPrecedence(t *testing.T) { // Test that node-configured callbacks take precedence over state callbacks. var nodeCallbackUsed, stateCallbackUsed bool diff --git a/graph/surface_runtime_test.go b/graph/surface_runtime_test.go index 465a69f4e0..1bbad62666 100644 --- a/graph/surface_runtime_test.go +++ b/graph/surface_runtime_test.go @@ -131,6 +131,91 @@ func TestLLMNode_SurfacePatch_AppendsTools(t *testing.T) { require.Contains(t, m.lastReq.Tools, "frontend_tool") } +func TestLLMNode_MandatoryToolFilterHidesRequestTools(t *testing.T) { + m := &captureModel{} + sg := NewStateGraph(MessagesStateSchema()) + sg.AddLLMNode( + "llm", + m, + "static instruction", + map[string]tool.Tool{ + "allowed_tool": &echoTool{name: "allowed_tool"}, + "hidden_tool": &echoTool{name: "hidden_tool"}, + }, + ) + inv := agent.NewInvocation( + agent.WithInvocationTraceNodeID("graph"), + agent.WithInvocationRunOptions(agent.NewRunOptions( + agent.WithMandatoryToolFilter( + tool.NewIncludeToolNamesFilter("allowed_tool"), + ), + )), + ) + ctx := agent.NewInvocationContext(context.Background(), inv) + node := sg.graph.nodes["llm"] + exec := &ExecutionContext{InvocationID: inv.InvocationID, Invocation: inv} + + _, err := node.Function(ctx, State{ + StateKeyExecContext: exec, + StateKeyCurrentNodeID: "llm", + StateKeyUserInput: "actual user", + }) + require.NoError(t, err) + + require.NotNil(t, m.lastReq) + require.Contains(t, m.lastReq.Tools, "allowed_tool") + require.NotContains(t, m.lastReq.Tools, "hidden_tool") +} + +func TestRunModelStream_ReappliesMandatoryToolFilterAfterBeforeModelCallbacks( + t *testing.T, +) { + m := &captureModel{} + allowed := &echoTool{name: "allowed_tool"} + hidden := &echoTool{name: "hidden_tool"} + callbacks := model.NewCallbacks().RegisterBeforeModel( + func( + _ context.Context, + req *model.Request, + ) (*model.Response, error) { + req.Tools["hidden_tool"] = hidden + return nil, nil + }, + ) + inv := agent.NewInvocation( + agent.WithInvocationRunOptions(agent.NewRunOptions( + agent.WithMandatoryToolFilter( + tool.NewIncludeToolNamesFilter("allowed_tool"), + ), + )), + ) + req := &model.Request{ + Messages: []model.Message{model.NewUserMessage("actual user")}, + Tools: map[string]tool.Tool{ + "allowed_tool": allowed, + }, + } + + _, stream, err := runModelStream( + agent.NewInvocationContext(context.Background(), inv), + inv, + callbacks, + m, + req, + nil, + ) + require.NoError(t, err) + require.NotNil(t, stream.Ch) + for range stream.Ch { + } + + require.NotNil(t, m.lastReq) + require.Contains(t, m.lastReq.Tools, "allowed_tool") + require.NotContains(t, m.lastReq.Tools, "hidden_tool") + require.Contains(t, req.Tools, "allowed_tool") + require.NotContains(t, req.Tools, "hidden_tool") +} + func TestToolsNode_SurfacePatch_OverridesExplicitTools(t *testing.T) { sg := NewStateGraph(MessagesStateSchema()) sg.AddToolsNode("tools", map[string]tool.Tool{ diff --git a/internal/flow/llmflow/llmflow.go b/internal/flow/llmflow/llmflow.go index 05c6a73c6c..c1f825ac0f 100644 --- a/internal/flow/llmflow/llmflow.go +++ b/internal/flow/llmflow/llmflow.go @@ -1367,6 +1367,10 @@ func (f *Flow) preprocess( eventChan chan<- *event.Event, ) *contextCompactionRebuildPlan { var rebuildPlan *contextCompactionRebuildPlan + var mandatoryToolFilter tool.FilterFunc + if invocation != nil { + mandatoryToolFilter = invocation.RunOptions.MandatoryToolFilter + } ctx, span, started := startLatencySpan( ctx, invocation, @@ -1425,6 +1429,7 @@ func (f *Flow) preprocess( } finishLatencySpan(stageSpan, stageStarted, nil) } + applyMandatoryRequestToolFilter(ctx, mandatoryToolFilter, llmRequest) // Sanitize invalid tool calls in history to avoid poisoning future requests. llmRequest.Messages = toolcall.SanitizeMessagesWithTools(ctx, llmRequest.Messages, llmRequest.Tools) return rebuildPlan @@ -1632,6 +1637,10 @@ func (f *Flow) rebuildRequestForContextCompaction( if rebuilt.Tools == nil { rebuilt.Tools = make(map[string]tool.Tool) } + var mandatoryToolFilter tool.FilterFunc + if invocation != nil { + mandatoryToolFilter = invocation.RunOptions.MandatoryToolFilter + } rebuildPlan.contentProcessor.ProcessRequest(ctx, invocation, rebuilt, nil) for _, tailProcessor := range rebuildPlan.tailProcessors { tailProcessor.RebuildRequestForContextCompaction( @@ -1640,6 +1649,7 @@ func (f *Flow) rebuildRequestForContextCompaction( rebuilt, ) } + applyMandatoryRequestToolFilter(ctx, mandatoryToolFilter, rebuilt) rebuilt.Messages = toolcall.SanitizeMessagesWithTools( ctx, rebuilt.Messages, @@ -1988,7 +1998,16 @@ func (f *Flow) getFilteredTools( hasUserToolTracking, invocation.RunOptions, ) - if f.toolActivationApplier != nil { + var activationApplied bool + allTools, userToolNames, externalToolNames, activationApplied = + toolsurface.ApplyInvocationToolActivation( + ctx, + invocation, + allTools, + userToolNames, + externalToolNames, + ) + if !activationApplied && f.toolActivationApplier != nil { allTools = append([]tool.Tool(nil), allTools...) if userToolNames != nil { userToolNames = copyToolNames(userToolNames) @@ -2006,6 +2025,15 @@ func (f *Flow) getFilteredTools( ) hasUserToolTracking = userToolNames != nil } + allTools, userToolNames, externalToolNames = + toolsurface.ApplyMandatoryToolFilter( + ctx, + allTools, + userToolNames, + externalToolNames, + invocation.RunOptions, + ) + hasUserToolTracking = userToolNames != nil // If no filter is specified, return all tools for this invocation. if invocation.RunOptions.ToolFilter == nil { @@ -2181,6 +2209,10 @@ func (f *Flow) callLLM( llmRequest *model.Request, callModel model.Model, ) (context.Context, model.Seq[*model.Response], error) { + var mandatoryToolFilter tool.FilterFunc + if invocation != nil { + mandatoryToolFilter = invocation.RunOptions.MandatoryToolFilter + } ctx, span, started := startLatencySpan( ctx, invocation, @@ -2216,6 +2248,7 @@ func (f *Flow) callLLM( if err != nil { return ctx, nil, err } + applyMandatoryRequestToolFilter(ctx, mandatoryToolFilter, llmRequest) if customResp != nil { return ctx, func(yield func(*model.Response) bool) { yield(customResp) @@ -2229,6 +2262,22 @@ func (f *Flow) callLLM( return ctx, seq, nil } +func applyMandatoryRequestToolFilter( + ctx context.Context, + mandatoryFilter tool.FilterFunc, + req *model.Request, +) { + if mandatoryFilter == nil || req == nil || len(req.Tools) == 0 { + return + } + for name, candidate := range req.Tools { + if toolName(candidate) == "" || + !mandatoryFilter(ctx, itool.ResolveDeclaration(candidate)) { + delete(req.Tools, name) + } + } +} + func (f *Flow) runBeforeModelCallbacks( ctx context.Context, invocation *agent.Invocation, diff --git a/internal/flow/llmflow/llmflow_test.go b/internal/flow/llmflow/llmflow_test.go index a7424ad98e..8b8c46285b 100644 --- a/internal/flow/llmflow/llmflow_test.go +++ b/internal/flow/llmflow/llmflow_test.go @@ -327,6 +327,36 @@ func TestPreprocess_AddsAgentToolsWhenPresent(t *testing.T) { require.Contains(t, req.Tools, "t1") } +func TestPreprocess_ReappliesMandatoryToolFilterAfterRequestProcessors( + t *testing.T, +) { + allowed := &mockTool{name: "allowed"} + hidden := &mockTool{name: "hidden"} + f := New( + []flow.RequestProcessor{&injectToolsRequestProcessor{ + tools: map[string]tool.Tool{"hidden": hidden}, + }}, + nil, + Options{}, + ) + req := &model.Request{Tools: map[string]tool.Tool{}} + inv := agent.NewInvocation( + agent.WithInvocationAgent(&minimalAgent{ + tools: []tool.Tool{allowed, hidden}, + }), + agent.WithInvocationRunOptions(agent.NewRunOptions( + agent.WithMandatoryToolFilter( + tool.NewIncludeToolNamesFilter("allowed"), + ), + )), + ) + + f.preprocess(context.Background(), inv, req, make(chan *event.Event, 1)) + + require.Contains(t, req.Tools, "allowed") + require.NotContains(t, req.Tools, "hidden") +} + func TestPreprocess_DowngradesOrphanToolCallBeforeModel(t *testing.T) { modelStub := &mockModel{ responses: []*model.Response{ @@ -2044,6 +2074,24 @@ func (p *seedMessagesRequestProcessor) ProcessRequest( req.Messages = append(req.Messages, cloneMessagesForTest(p.messages)...) } +type injectToolsRequestProcessor struct { + tools map[string]tool.Tool +} + +func (p *injectToolsRequestProcessor) ProcessRequest( + _ context.Context, + _ *agent.Invocation, + req *model.Request, + _ chan<- *event.Event, +) { + if req.Tools == nil { + req.Tools = make(map[string]tool.Tool) + } + for name, candidate := range p.tools { + req.Tools[name] = candidate + } +} + const flowRunPanicTestMsg = "boom" type panicRequestProcessor struct{} @@ -3134,6 +3182,49 @@ func TestFlow_CallLLM_PluginBeforeModelCanShortCircuit(t *testing.T) { require.False(t, m.called) } +func TestFlow_CallLLM_ReappliesMandatoryToolFilterAfterBeforeModelCallbacks( + t *testing.T, +) { + allowed := &mockTool{name: "allowed"} + hidden := &mockTool{name: "hidden"} + callbacks := model.NewCallbacks().RegisterBeforeModel( + func( + _ context.Context, + req *model.Request, + ) (*model.Response, error) { + req.Tools["hidden"] = hidden + return nil, nil + }, + ) + f := New(nil, nil, Options{ModelCallbacks: callbacks}) + selectedModel := &namedFlowModel{name: "selected"} + inv := agent.NewInvocation( + agent.WithInvocationModel(selectedModel), + agent.WithInvocationRunOptions(agent.NewRunOptions( + agent.WithMandatoryToolFilter( + tool.NewIncludeToolNamesFilter("allowed"), + ), + )), + ) + req := &model.Request{ + Messages: []model.Message{model.NewUserMessage("check tools")}, + Tools: map[string]tool.Tool{"allowed": allowed}, + } + + _, seq, err := f.callLLM( + context.Background(), + inv, + req, + selectedModel, + ) + require.NoError(t, err) + seq(func(*model.Response) bool { return true }) + + require.True(t, selectedModel.Called()) + require.Contains(t, req.Tools, "allowed") + require.NotContains(t, req.Tools, "hidden") +} + type testCtxKey struct{} func TestFlow_CallLLM_PluginBeforeModelError(t *testing.T) { diff --git a/internal/flow/llmflow/tool_filter_test.go b/internal/flow/llmflow/tool_filter_test.go index d2e0900eb4..00a85762f4 100644 --- a/internal/flow/llmflow/tool_filter_test.go +++ b/internal/flow/llmflow/tool_filter_test.go @@ -550,6 +550,102 @@ func TestGetFilteredTools_AppendsRunOptionTools(t *testing.T) { require.Empty(t, traceableNames) } +func TestGetFilteredTools_MandatoryFilterBlocksAdditionalTools( + t *testing.T, +) { + f := New(nil, nil, Options{}) + frameworkTool := &mockTool{name: "framework_tool"} + additionalTool := &mockTool{name: "additional_tool"} + mockAgent := &mockAgentWithInvocationToolSurface{ + name: "test-agent", + allTools: []tool.Tool{frameworkTool}, + userToolNames: map[string]bool{}, + } + inv := agent.NewInvocation( + agent.WithInvocationAgent(mockAgent), + agent.WithInvocationRunOptions(agent.NewRunOptions( + agent.WithAdditionalTools([]tool.Tool{additionalTool}), + agent.WithMandatoryToolFilter( + tool.NewIncludeToolNamesFilter("framework_tool"), + ), + )), + ) + + filtered := f.getFilteredTools(context.Background(), inv) + + require.Equal(t, []tool.Tool{frameworkTool}, filtered) + hasUserTools, ok := InvocationHasFilteredUserTools(inv) + require.True(t, ok) + require.False(t, hasUserTools) +} + +func TestGetFilteredTools_MandatoryFilterBlocksExternalTools( + t *testing.T, +) { + f := New(nil, nil, Options{}) + frameworkTool := &mockTool{name: "framework_tool"} + externalTool := &mockTool{name: "external_tool"} + mockAgent := &mockAgentWithInvocationToolSurface{ + name: "test-agent", + allTools: []tool.Tool{frameworkTool}, + userToolNames: map[string]bool{}, + } + inv := agent.NewInvocation( + agent.WithInvocationAgent(mockAgent), + agent.WithInvocationRunOptions(agent.NewRunOptions( + agent.WithExternalTools([]tool.Tool{externalTool}), + agent.WithMandatoryToolFilter( + tool.NewIncludeToolNamesFilter("framework_tool"), + ), + )), + ) + + filtered := f.getFilteredTools(context.Background(), inv) + + require.Equal(t, []tool.Tool{frameworkTool}, filtered) + require.Empty(t, inv.RunOptions.ExternalToolNames) +} + +func TestGetFilteredTools_MandatoryFilterBlocksActivatedTools( + t *testing.T, +) { + frameworkTool := &mockTool{name: "framework_tool"} + activatedTool := &mockTool{name: "activated_tool"} + f := New(nil, nil, Options{ + ToolActivationApplier: func( + _ context.Context, + _ *agent.Invocation, + tools []tool.Tool, + userNames map[string]bool, + externalNames map[string]bool, + ) ([]tool.Tool, map[string]bool, map[string]bool) { + tools = append(tools, activatedTool) + userNames[activatedTool.Declaration().Name] = true + return tools, userNames, externalNames + }, + }) + mockAgent := &mockAgentWithInvocationToolSurface{ + name: "test-agent", + allTools: []tool.Tool{frameworkTool}, + userToolNames: map[string]bool{}, + } + inv := agent.NewInvocation( + agent.WithInvocationAgent(mockAgent), + agent.WithInvocationRunOptions(agent.NewRunOptions( + agent.WithMandatoryToolFilter( + tool.NewIncludeToolNamesFilter("framework_tool"), + ), + )), + ) + + filtered := f.getFilteredTools(context.Background(), inv) + + require.Equal(t, []tool.Tool{frameworkTool}, filtered) + hasUserTools, ok := InvocationHasFilteredUserTools(inv) + require.True(t, ok) + require.False(t, hasUserTools) +} + func TestGetFilteredTools_FiltersRunOptionToolsWithFilterProvider( t *testing.T, ) { diff --git a/internal/flow/processor/functioncall.go b/internal/flow/processor/functioncall.go index f17df244ae..12c92ee200 100644 --- a/internal/flow/processor/functioncall.go +++ b/internal/flow/processor/functioncall.go @@ -2429,6 +2429,22 @@ func (p *FunctionCallResponseProcessor) executeToolWithCallbacks( } rememberExecutingToolArgs(ctx, toolCall.Function.Arguments) toolDeclaration := tl.Declaration() + visibilityResult, err := checkMandatoryToolVisibility( + ctx, + invocation, + toolCall, + tl, + toolDeclaration, + ) + if err != nil { + return ctx, nil, toolCall.Function.Arguments, false, false, err + } + if visibilityResult != nil { + ctx = withSkippedToolStateDelta(ctx) + ctx = withSkippedToolSkipSummarization(ctx) + return ctx, *visibilityResult, toolCall.Function.Arguments, false, + false, nil + } ctx, toolCall, customResult, err := p.runBeforeToolPluginCallbacks( ctx, invocation, @@ -2542,6 +2558,42 @@ func (p *FunctionCallResponseProcessor) executeToolWithCallbacks( suppressDefaultToolMessage, skipSummarization || localSkip, toolErr } +func checkMandatoryToolVisibility( + ctx context.Context, + invocation *agent.Invocation, + toolCall model.ToolCall, + tl tool.Tool, + decl *tool.Declaration, +) (*tool.PermissionResult, error) { + if invocation == nil || invocation.RunOptions.MandatoryToolFilter == nil { + return nil, nil + } + if invocation.RunOptions.MandatoryToolFilter( + ctx, + itool.ResolveDeclaration(tl), + ) { + return nil, nil + } + req := &tool.PermissionRequest{ + Tool: tl, + ToolName: toolCall.Function.Name, + ToolCallID: toolCall.ID, + Declaration: decl, + Arguments: toolCall.Function.Arguments, + Metadata: tool.MetadataOf(itool.ResolveSemantic(tl)), + } + return normalizeToolPermissionResult( + req, + tool.DenyPermission( + fmt.Sprintf( + "tool %q is hidden by mandatory tool filter", + req.ToolName, + ), + ), + nil, + ) +} + func (p *FunctionCallResponseProcessor) checkToolPermission( ctx context.Context, invocation *agent.Invocation, @@ -2565,10 +2617,10 @@ func (p *FunctionCallResponseProcessor) checkToolPermission( return result, err } } - if invocation == nil || invocation.RunOptions.ToolPermissionPolicy == nil { + if invocation == nil { return nil, nil } - decision, err := invocation.RunOptions.ToolPermissionPolicy.CheckToolPermission(ctx, req) + decision, err := invocation.RunOptions.CheckToolPermission(ctx, req) return normalizeToolPermissionResult(req, decision, err) } diff --git a/internal/flow/processor/functioncall_test.go b/internal/flow/processor/functioncall_test.go index d020f081f9..0e1ca08a88 100644 --- a/internal/flow/processor/functioncall_test.go +++ b/internal/flow/processor/functioncall_test.go @@ -124,6 +124,7 @@ type permissionMockTool struct { *mockCallableTool metadata tool.ToolMetadata decision tool.PermissionDecision + permissionCalled bool stateDelta map[string][]byte stateDeltaCalled bool skipSummarize bool @@ -149,6 +150,7 @@ func (m *permissionMockTool) CheckPermission( _ context.Context, _ *tool.PermissionRequest, ) (tool.PermissionDecision, error) { + m.permissionCalled = true return m.decision, nil } @@ -9006,6 +9008,78 @@ func TestExecuteToolWithCallbacks_ToolPermissionPolicyDenySkipsExecution( require.JSONEq(t, permissionJSON, string(mustJSON(res))) } +func TestExecuteToolWithCallbacks_MandatoryToolFilterDenySkipsCallbacksCheckerAndExecution( + t *testing.T, +) { + const ( + toolName = "hidden_tool" + permissionJSON = `{"status":"denied","tool":"hidden_tool","reason":"tool \"hidden_tool\" is hidden by mandatory tool filter"}` + ) + var ( + calledTool bool + calledCallback bool + ordinaryPolicy bool + ) + callbacks := tool.NewCallbacks() + callbacks.RegisterBeforeTool(func( + _ context.Context, + _ *tool.BeforeToolArgs, + ) (*tool.BeforeToolResult, error) { + calledCallback = true + return &tool.BeforeToolResult{}, nil + }) + tl := &permissionMockTool{ + mockCallableTool: &mockCallableTool{ + declaration: &tool.Declaration{Name: toolName}, + callFn: func(context.Context, []byte) (any, error) { + calledTool = true + return map[string]any{"ok": true}, nil + }, + }, + decision: tool.AllowPermission(), + } + inv := agent.NewInvocation( + agent.WithInvocationRunOptions(agent.NewRunOptions( + agent.WithMandatoryToolFilter( + tool.NewIncludeToolNamesFilter("allowed_tool"), + ), + agent.WithToolPermissionPolicyFunc( + func( + context.Context, + *tool.PermissionRequest, + ) (tool.PermissionDecision, error) { + ordinaryPolicy = true + return tool.AllowPermission(), nil + }, + ), + )), + ) + + _, result, modifiedArgs, _, _, err := + NewFunctionCallResponseProcessor(false, callbacks). + executeToolWithCallbacks( + context.Background(), + inv, + model.ToolCall{ + ID: "call-hidden", + Function: model.FunctionDefinitionParam{ + Name: toolName, + Arguments: []byte(`{"value":"blocked"}`), + }, + }, + tl, + nil, + ) + + require.NoError(t, err) + require.False(t, calledCallback) + require.False(t, tl.permissionCalled) + require.False(t, ordinaryPolicy) + require.False(t, calledTool) + require.JSONEq(t, `{"value":"blocked"}`, string(modifiedArgs)) + require.JSONEq(t, permissionJSON, string(mustJSON(result))) +} + func TestExecuteToolCall_ToolPermissionResultSkipsToolResultMessagesCallback( t *testing.T, ) { @@ -9071,6 +9145,67 @@ func TestExecuteToolCall_ToolPermissionResultSkipsToolResultMessagesCallback( require.JSONEq(t, permissionJSON, choices[0].Message.Content) } +func TestExecuteToolWithCallbacks_MandatoryPermissionDenyCannotBeOverridden( + t *testing.T, +) { + const ( + toolName = "shell" + denyReason = "tenant policy denied shell" + permissionJSON = `{"status":"denied","tool":"shell","reason":"tenant policy denied shell"}` + ) + var ( + calledTool bool + ordinaryCalled bool + ) + tl := &mockCallableTool{ + declaration: &tool.Declaration{Name: toolName}, + callFn: func(context.Context, []byte) (any, error) { + calledTool = true + return map[string]any{"ok": true}, nil + }, + } + inv := &agent.Invocation{ + RunOptions: agent.NewRunOptions( + agent.WithMandatoryToolPermissionPolicyFunc( + func( + context.Context, + *tool.PermissionRequest, + ) (tool.PermissionDecision, error) { + return tool.DenyPermission(denyReason), nil + }, + ), + agent.WithToolPermissionPolicyFunc( + func( + context.Context, + *tool.PermissionRequest, + ) (tool.PermissionDecision, error) { + ordinaryCalled = true + return tool.AllowPermission(), nil + }, + ), + ), + } + + _, res, _, _, _, err := NewFunctionCallResponseProcessor(false, nil). + executeToolWithCallbacks( + context.Background(), + inv, + model.ToolCall{ + ID: "call-shell", + Function: model.FunctionDefinitionParam{ + Name: toolName, + Arguments: []byte(`{}`), + }, + }, + tl, + nil, + ) + require.NoError(t, err) + require.False(t, calledTool) + require.False(t, ordinaryCalled) + require.JSONEq(t, permissionJSON, string(mustJSON(res))) +} + func TestExecuteSingleToolCallSequential_ToolPermissionResultSkipsStateDelta( t *testing.T, ) { diff --git a/internal/toolsurface/toolsurface.go b/internal/toolsurface/toolsurface.go index cb5e815001..89530072b5 100644 --- a/internal/toolsurface/toolsurface.go +++ b/internal/toolsurface/toolsurface.go @@ -9,7 +9,8 @@ // Package toolsurface resolves the effective tool surface for an invocation: // the base surface exposed by the agent plus the run-scoped tools, with the -// run-scoped tool filter applied. It is the single source of truth shared by +// mandatory and ordinary run-scoped tool filters applied. It is the single +// source of truth shared by // the LLM flow (which uses it to build the model request) and by helpers such // as the dynamic AgentTool (which derives a child capability surface from a // parent invocation). Keeping the logic here avoids both behavioral drift and @@ -124,6 +125,37 @@ func AppendRunOptionTools( return allTools, userToolNames, hasUserToolTracking, externalNames } +// ApplyInvocationToolActivation applies the agent's invocation-scoped +// activation layer, if supported. The inputs are copied before invoking the +// provider so activation cannot mutate the configured/base surface. +func ApplyInvocationToolActivation( + ctx context.Context, + invocation *agent.Invocation, + allTools []tool.Tool, + userToolNames map[string]bool, + externalToolNames map[string]bool, +) ([]tool.Tool, map[string]bool, map[string]bool, bool) { + if invocation == nil || invocation.Agent == nil { + return allTools, userToolNames, externalToolNames, false + } + provider, ok := invocation.Agent.(agent.InvocationToolActivationProvider) + if !ok { + return allTools, userToolNames, externalToolNames, false + } + allTools = append([]tool.Tool(nil), allTools...) + userToolNames = copyToolNames(userToolNames) + externalToolNames = copyToolNames(externalToolNames) + allTools, userToolNames, externalToolNames = + provider.ApplyInvocationToolActivation( + ctx, + invocation, + allTools, + userToolNames, + externalToolNames, + ) + return allTools, userToolNames, externalToolNames, true +} + // ApplyToolFilter applies the run-scoped ToolFilter to allTools, always keeping // framework tools and keeping user tools only when the filter passes. The // result is sorted by name for stable prompt-cache behavior. It assumes @@ -166,6 +198,40 @@ func ApplyToolFilter( return filtered } +// ApplyMandatoryToolFilter applies the non-negotiable run filter to the +// complete tool surface and removes hidden names from the user/external +// classification maps. Unlike ApplyToolFilter, framework tools are not exempt. +func ApplyMandatoryToolFilter( + ctx context.Context, + allTools []tool.Tool, + userToolNames map[string]bool, + externalToolNames map[string]bool, + opts agent.RunOptions, +) ([]tool.Tool, map[string]bool, map[string]bool) { + if opts.MandatoryToolFilter == nil { + return allTools, userToolNames, externalToolNames + } + filtered := make([]tool.Tool, 0, len(allTools)) + visibleNames := make(map[string]bool, len(allTools)) + for _, candidate := range allTools { + name := toolName(candidate) + if name == "" { + continue + } + if !opts.MandatoryToolFilter( + ctx, + itool.ResolveDeclaration(candidate), + ) { + continue + } + filtered = append(filtered, candidate) + visibleNames[name] = true + } + return filtered, + visibleToolNames(userToolNames, visibleNames), + visibleToolNames(externalToolNames, visibleNames) +} + // Effective returns the effective tool surface for the invocation: the base // surface (InvocationToolSurface / FilterTools / Tools) plus // RunOptions.AdditionalTools and ExternalTools, with the run-scoped @@ -213,6 +279,23 @@ func EffectiveWithExternal( hasUserToolTracking, invocation.RunOptions, ) + allTools, userToolNames, externalNames, _ = + ApplyInvocationToolActivation( + ctx, + invocation, + allTools, + userToolNames, + externalNames, + ) + hasUserToolTracking = userToolNames != nil + allTools, userToolNames, externalNames = ApplyMandatoryToolFilter( + ctx, + allTools, + userToolNames, + externalNames, + invocation.RunOptions, + ) + hasUserToolTracking = userToolNames != nil if invocation.RunOptions.ToolFilter == nil { return allTools, userToolNames, externalNames } @@ -260,6 +343,9 @@ func collectToolNames(tools []tool.Tool) map[string]bool { } func copyToolNames(src map[string]bool) map[string]bool { + if src == nil { + return nil + } dst := make(map[string]bool, len(src)) for name, ok := range src { dst[name] = ok @@ -267,6 +353,22 @@ func copyToolNames(src map[string]bool) map[string]bool { return dst } +func visibleToolNames( + names map[string]bool, + visibleNames map[string]bool, +) map[string]bool { + if names == nil { + return nil + } + visible := make(map[string]bool, len(names)) + for name, enabled := range names { + if enabled && visibleNames[name] { + visible[name] = true + } + } + return visible +} + func toolName(tl tool.Tool) string { if tl == nil { return "" diff --git a/internal/toolsurface/toolsurface_test.go b/internal/toolsurface/toolsurface_test.go index 19e46fa335..7874767522 100644 --- a/internal/toolsurface/toolsurface_test.go +++ b/internal/toolsurface/toolsurface_test.go @@ -123,6 +123,75 @@ type stubSurfaceAgent struct { userTools []tool.Tool } +type stubActivationSurfaceAgent struct { + *stubSurfaceAgent + seen []string +} + +type stubUntrackedActivationAgent struct { + tools []tool.Tool + sawNilUsers bool +} + +func (s *stubActivationSurfaceAgent) ApplyInvocationToolActivation( + _ context.Context, + _ *agent.Invocation, + tools []tool.Tool, + userToolNames map[string]bool, + externalToolNames map[string]bool, +) ([]tool.Tool, map[string]bool, map[string]bool) { + s.seen = make([]string, 0, len(tools)) + for _, candidate := range tools { + s.seen = append(s.seen, candidate.Declaration().Name) + } + filtered := make([]tool.Tool, 0, len(tools)) + for _, candidate := range tools { + if candidate.Declaration().Name == "disabled" { + delete(userToolNames, "disabled") + delete(externalToolNames, "disabled") + continue + } + filtered = append(filtered, candidate) + } + return filtered, userToolNames, externalToolNames +} + +func (s *stubUntrackedActivationAgent) ApplyInvocationToolActivation( + _ context.Context, + _ *agent.Invocation, + tools []tool.Tool, + userToolNames map[string]bool, + externalToolNames map[string]bool, +) ([]tool.Tool, map[string]bool, map[string]bool) { + s.sawNilUsers = userToolNames == nil + return tools, userToolNames, externalToolNames +} + +func (s *stubUntrackedActivationAgent) Run( + context.Context, + *agent.Invocation, +) (<-chan *event.Event, error) { + ch := make(chan *event.Event) + close(ch) + return ch, nil +} + +func (s *stubUntrackedActivationAgent) Tools() []tool.Tool { + return s.tools +} + +func (s *stubUntrackedActivationAgent) Info() agent.Info { + return agent.Info{Name: "untracked-activation-agent"} +} + +func (s *stubUntrackedActivationAgent) SubAgents() []agent.Agent { + return nil +} + +func (s *stubUntrackedActivationAgent) FindSubAgent(string) agent.Agent { + return nil +} + func (s *stubSurfaceAgent) Run( context.Context, *agent.Invocation, @@ -236,6 +305,103 @@ func TestEffectiveWithExternal_AppendsAndClassifiesRunOptionTools(t *testing.T) require.Equal(t, map[string]bool{"external": true}, externalNames) } +func TestEffectiveWithExternal_AppliesMandatoryFilterAfterRunOptionTools( + t *testing.T, +) { + agt := &stubSurfaceAgent{ + tools: []tool.Tool{surfaceTool("base")}, + userTools: []tool.Tool{surfaceTool("base")}, + } + inv := agent.NewInvocation( + agent.WithInvocationAgent(agt), + agent.WithInvocationRunOptions(agent.NewRunOptions( + agent.WithAdditionalTools([]tool.Tool{ + surfaceTool("added"), + }), + agent.WithExternalTools([]tool.Tool{ + surfaceTool("external"), + }), + agent.WithMandatoryToolFilter( + tool.NewIncludeToolNamesFilter("base"), + ), + )), + ) + + tools, userToolNames, externalNames := EffectiveWithExternal( + context.Background(), + inv, + ) + + requireToolNames(t, tools, []string{"base"}) + require.Equal(t, map[string]bool{"base": true}, userToolNames) + require.Empty(t, externalNames) +} + +func TestEffectiveWithExternal_AppliesInvocationActivationAfterRunOptionTools( + t *testing.T, +) { + base := &stubSurfaceAgent{ + tools: []tool.Tool{ + surfaceTool("base"), + surfaceTool("disabled"), + }, + userTools: []tool.Tool{ + surfaceTool("base"), + surfaceTool("disabled"), + }, + } + agt := &stubActivationSurfaceAgent{stubSurfaceAgent: base} + inv := agent.NewInvocation( + agent.WithInvocationAgent(agt), + agent.WithInvocationRunOptions(agent.NewRunOptions( + agent.WithAdditionalTools([]tool.Tool{surfaceTool("added")}), + )), + ) + + tools, userToolNames, externalNames := EffectiveWithExternal( + context.Background(), + inv, + ) + + require.ElementsMatch(t, []string{"base", "disabled", "added"}, agt.seen) + requireToolNames(t, tools, []string{"base", "added"}) + require.Equal(t, map[string]bool{ + "base": true, + "added": true, + }, userToolNames) + require.Empty(t, externalNames) + requireToolNames(t, base.tools, []string{"base", "disabled"}) +} + +func TestEffectiveWithExternal_ActivationPreservesMissingUserToolTracking( + t *testing.T, +) { + agt := &stubUntrackedActivationAgent{ + tools: []tool.Tool{ + surfaceTool("keep"), + surfaceTool("drop"), + }, + } + inv := agent.NewInvocation( + agent.WithInvocationAgent(agt), + agent.WithInvocationRunOptions(agent.NewRunOptions( + agent.WithToolFilter( + tool.NewIncludeToolNamesFilter("keep"), + ), + )), + ) + + tools, userToolNames, externalNames := EffectiveWithExternal( + context.Background(), + inv, + ) + + require.True(t, agt.sawNilUsers) + require.Nil(t, userToolNames) + require.Nil(t, externalNames) + requireToolNames(t, tools, []string{"keep"}) +} + func TestApplyDeclarations_OverridesDeclarationAndPreservesCall(t *testing.T) { base := &stubCallableSurfaceTool{ stubSurfaceTool: stubSurfaceTool{decl: &tool.Declaration{ diff --git a/platform/gateway/errors.go b/platform/gateway/errors.go index 9db62b54db..bedcf57cc3 100644 --- a/platform/gateway/errors.go +++ b/platform/gateway/errors.go @@ -17,6 +17,8 @@ var ( ErrRuntimeInactive = errors.New("gateway runtime inactive") // ErrRuntimeMismatch indicates that a runtime's tenant, app, binding, or inbound identifiers do not match. ErrRuntimeMismatch = errors.New("gateway runtime identifiers mismatch") + // ErrToolPermissionPolicyRequired indicates that a governed app runtime has no enforcement policy. + ErrToolPermissionPolicyRequired = errors.New("gateway tool permission policy is required") // ErrBindingAccessDenied indicates that a binding policy rejects the inbound sender or conversation. ErrBindingAccessDenied = errors.New("gateway binding access denied") // ErrBindingMentionRequired indicates that a group/thread message did not mention the agent. diff --git a/platform/gateway/registry.go b/platform/gateway/registry.go index d370596827..1e3a4c6827 100644 --- a/platform/gateway/registry.go +++ b/platform/gateway/registry.go @@ -10,10 +10,12 @@ package gateway import ( "context" + "strings" "sync" "trpc.group/trpc-go/trpc-agent-go/platform" "trpc.group/trpc-go/trpc-agent-go/runner" + "trpc.group/trpc-go/trpc-agent-go/tool" ) // Runtime contains the platform configuration and runner for one active binding. @@ -23,6 +25,10 @@ type Runtime struct { Binding platform.ChannelBinding Runner runner.Runner Audit platform.AuditSink + // ToolFilter narrows user-visible tools for this runtime. + ToolFilter tool.FilterFunc + // ToolPermissionPolicy enforces tool-call authorization before execution. + ToolPermissionPolicy tool.PermissionPolicy } // Validate checks that the runtime can process inbound messages. @@ -39,6 +45,10 @@ func (r Runtime) Validate() error { if r.Runner == nil { return ErrRuntimeNotFound } + if strings.TrimSpace(r.App.ToolPolicyID) != "" && + isNilInterfaceValue(r.ToolPermissionPolicy) { + return ErrToolPermissionPolicyRequired + } if r.App.TenantID != r.Tenant.TenantID || r.Binding.TenantID != r.Tenant.TenantID || r.Binding.AppID != r.App.AppID { diff --git a/platform/gateway/registry_test.go b/platform/gateway/registry_test.go index 078fe00f58..632324751f 100644 --- a/platform/gateway/registry_test.go +++ b/platform/gateway/registry_test.go @@ -16,6 +16,7 @@ import ( "github.com/stretchr/testify/require" "trpc.group/trpc-go/trpc-agent-go/platform" + "trpc.group/trpc-go/trpc-agent-go/tool" ) func TestInMemoryRegistryAvoidsTenantAppDelimiterCollision(t *testing.T) { @@ -39,6 +40,39 @@ func TestInMemoryRegistryAvoidsTenantAppDelimiterCollision(t *testing.T) { assert.Same(t, secondRunner, gotSecond.Runner) } +func TestInMemoryRegistryRejectsGovernedRuntimeWithoutPermissionPolicy( + t *testing.T, +) { + var typedNil *nilablePermissionPolicy + tests := map[string]tool.PermissionPolicy{ + "nil": nil, + "typed nil": typedNil, + } + for name, policy := range tests { + t.Run(name, func(t *testing.T) { + registry := NewInMemoryRegistry() + runtime := validRuntime( + "tenant-a", + &recordingRunner{response: "unused"}, + ) + runtime.App.ToolPolicyID = "policy-a" + runtime.ToolPermissionPolicy = policy + + err := registry.Register(runtime) + require.ErrorIs(t, err, ErrToolPermissionPolicyRequired) + }) + } +} + +type nilablePermissionPolicy struct{} + +func (*nilablePermissionPolicy) CheckToolPermission( + context.Context, + *tool.PermissionRequest, +) (tool.PermissionDecision, error) { + return tool.AllowPermission(), nil +} + func inboundForRegistryRuntime(runtime Runtime) platform.InboundMessage { return platform.InboundMessage{ TenantID: runtime.Tenant.TenantID, diff --git a/platform/gateway/service.go b/platform/gateway/service.go index ef12dd2fb6..112440b73e 100644 --- a/platform/gateway/service.go +++ b/platform/gateway/service.go @@ -424,14 +424,31 @@ func (s *Service) runGatewayRunner( if input.FencingToken > 0 { runnerSpan.SetAttributes(attribute.Int64("storage.fencing_token", input.FencingToken)) } + runOptions := []agent.RunOption{ + agent.WithRequestID(input.RequestID), + agent.WithLatencyDiagnostics(true), + agent.WithLatencyDiagnosticsEvents(false), + } + if runtime.ToolFilter != nil { + runOptions = append( + runOptions, + agent.WithMandatoryToolFilter(runtime.ToolFilter), + ) + } + if !isNilInterfaceValue(runtime.ToolPermissionPolicy) { + runOptions = append( + runOptions, + agent.WithMandatoryToolPermissionPolicy( + runtime.ToolPermissionPolicy, + ), + ) + } ch, err := runtime.Runner.Run( runnerCtx, input.InternalUserID, input.SessionID, model.NewUserMessage(input.Text), - agent.WithRequestID(input.RequestID), - agent.WithLatencyDiagnostics(true), - agent.WithLatencyDiagnosticsEvents(false), + runOptions..., ) if err != nil { s.writeAuditTo(auditCtx, auditSink, auditFromMessage(msg, input.SessionID, input.InternalUserID, "runner_error", err.Error(), input.Start, err)) @@ -758,24 +775,24 @@ func (s *Service) writeAuditTo( auditSink platform.AuditSink, record platform.AuditRecord, ) { - if isNilAuditSink(auditSink) { + if isNilInterfaceValue(auditSink) { return } _ = auditSink.WriteAudit(ctx, record) } func (s *Service) auditSinkForRuntime(runtime Runtime) platform.AuditSink { - if !isNilAuditSink(runtime.Audit) { + if !isNilInterfaceValue(runtime.Audit) { return runtime.Audit } return s.auditSink } -func isNilAuditSink(auditSink platform.AuditSink) bool { - if auditSink == nil { +func isNilInterfaceValue(value any) bool { + if value == nil { return true } - reflected := reflect.ValueOf(auditSink) + reflected := reflect.ValueOf(value) switch reflected.Kind() { case reflect.Chan, reflect.Func, diff --git a/platform/toolpolicy/filter.go b/platform/toolpolicy/filter.go new file mode 100644 index 0000000000..53831e48b3 --- /dev/null +++ b/platform/toolpolicy/filter.go @@ -0,0 +1,57 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package toolpolicy + +import ( + "context" + "strings" + + "trpc.group/trpc-go/trpc-agent-go/tool" +) + +// ToolFilter returns the run-scoped visibility filter for this policy. +// +// A non-empty whitelist is an allow boundary. Tool and platform denylists +// always override the whitelist. Nil is returned when the policy does not +// constrain tool names. +func (p *Policy) ToolFilter() tool.FilterFunc { + if p == nil { + return nil + } + whitelist := nameSet(normalizedList(p.policy.ToolWhitelist)) + denylist := nameSet(policyDenylist(p.policy)) + if len(whitelist) == 0 && len(denylist) == 0 { + return nil + } + return func(_ context.Context, candidate tool.Tool) bool { + if candidate == nil || candidate.Declaration() == nil { + return false + } + name := strings.TrimSpace(candidate.Declaration().Name) + if name == "" { + return false + } + if _, denied := denylist[name]; denied { + return false + } + if len(whitelist) == 0 { + return true + } + _, allowed := whitelist[name] + return allowed + } +} + +func nameSet(names []string) map[string]struct{} { + set := make(map[string]struct{}, len(names)) + for _, name := range names { + set[name] = struct{}{} + } + return set +} diff --git a/platform/toolpolicy/filter_test.go b/platform/toolpolicy/filter_test.go new file mode 100644 index 0000000000..b869520b91 --- /dev/null +++ b/platform/toolpolicy/filter_test.go @@ -0,0 +1,80 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package toolpolicy + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "trpc.group/trpc-go/trpc-agent-go/platform" + "trpc.group/trpc-go/trpc-agent-go/tool" +) + +func TestPolicyToolFilterMatchesNameGovernance(t *testing.T) { + policy, err := New(platform.ToolPolicy{ + TenantID: "tenant-a", + AppID: "app-a", + PolicyID: "policy-a", + ToolWhitelist: []string{"read_file", "workspace_write"}, + ToolDenylist: []string{"workspace_write"}, + PlatformDenylist: []string{"shell"}, + }) + require.NoError(t, err) + filter := policy.ToolFilter() + require.NotNil(t, filter) + + assert.True(t, filter(context.Background(), namedTool("read_file"))) + assert.False(t, filter(context.Background(), namedTool("workspace_write"))) + assert.False(t, filter(context.Background(), namedTool("shell"))) + assert.False(t, filter(context.Background(), namedTool("unknown"))) + assert.False(t, filter(context.Background(), nil)) +} + +func TestPolicyToolFilterAllowsNonDeniedToolsWithoutWhitelist(t *testing.T) { + policy, err := New(platform.ToolPolicy{ + TenantID: "tenant-a", + AppID: "app-a", + PolicyID: "policy-a", + ToolDenylist: []string{"blocked"}, + }) + require.NoError(t, err) + filter := policy.ToolFilter() + require.NotNil(t, filter) + + assert.True(t, filter(context.Background(), namedTool("allowed"))) + assert.False(t, filter(context.Background(), namedTool("blocked"))) +} + +func TestPolicyToolFilterReturnsNilWithoutNameConstraints(t *testing.T) { + policy, err := New(platform.ToolPolicy{ + TenantID: "tenant-a", + AppID: "app-a", + PolicyID: "policy-a", + }) + require.NoError(t, err) + + assert.Nil(t, policy.ToolFilter()) +} + +type filterTestTool struct { + declaration *tool.Declaration +} + +func namedTool(name string) tool.Tool { + return &filterTestTool{ + declaration: &tool.Declaration{Name: name}, + } +} + +func (t *filterTestTool) Declaration() *tool.Declaration { + return t.declaration +} diff --git a/platform/toolpolicy/policy.go b/platform/toolpolicy/policy.go index 89771b21fe..58e334398e 100644 --- a/platform/toolpolicy/policy.go +++ b/platform/toolpolicy/policy.go @@ -103,6 +103,7 @@ func New(policy platform.ToolPolicy, opts ...Option) (*Policy, error) { if err := validateRuntimeIdentity(policy); err != nil { return nil, err } + policy = cloneToolPolicy(policy) redactor, err := platform.NewRedactor(policy.ArgumentRedactionRules...) if err != nil { return nil, fmt.Errorf("newing platform tool policy: redaction rules: %w", err) @@ -121,6 +122,18 @@ func New(policy platform.ToolPolicy, opts ...Option) (*Policy, error) { return p, nil } +func cloneToolPolicy(policy platform.ToolPolicy) platform.ToolPolicy { + policy.ToolWhitelist = append([]string(nil), policy.ToolWhitelist...) + policy.ToolDenylist = append([]string(nil), policy.ToolDenylist...) + policy.ArgumentRedactionRules = append( + []string(nil), + policy.ArgumentRedactionRules..., + ) + policy.PlatformDenylist = append([]string(nil), policy.PlatformDenylist...) + policy.HighRiskTools = append([]string(nil), policy.HighRiskTools...) + return policy +} + // Name implements plugin.Plugin when Policy is registered with a plugin manager. func (p *Policy) Name() string { if p == nil || p.name == "" { diff --git a/platform/toolpolicy/policy_test.go b/platform/toolpolicy/policy_test.go index 5e4818f57d..41762ac246 100644 --- a/platform/toolpolicy/policy_test.go +++ b/platform/toolpolicy/policy_test.go @@ -15,6 +15,8 @@ import ( "testing" "time" + "github.com/stretchr/testify/require" + "trpc.group/trpc-go/trpc-agent-go/platform" "trpc.group/trpc-go/trpc-agent-go/plugin" "trpc.group/trpc-go/trpc-agent-go/plugin/guardrail/approval" @@ -592,6 +594,34 @@ func TestApprovalOptionsWithReviewerAllowsWhitelistedHighRiskAsk(t *testing.T) { } } +func TestNewCopiesMutablePolicySlices(t *testing.T) { + source := defaultPolicy(platform.ToolPolicy{ + ToolWhitelist: []string{"read_file"}, + ToolDenylist: []string{"shell"}, + ArgumentRedactionRules: []string{"secret"}, + PlatformDenylist: []string{"admin"}, + HighRiskTools: []string{"workspace_write"}, + DangerousToolAction: platform.DangerousToolActionDeny, + NetworkPolicyJSON: `{"mode":"deny"}`, + FilesystemPolicyJSON: `{"mode":"deny"}`, + ToolBudgetRemainingJSON: `{"calls":1}`, + }) + p, err := New(source) + require.NoError(t, err) + + source.ToolWhitelist[0] = "mutated" + source.ToolDenylist[0] = "mutated" + source.ArgumentRedactionRules[0] = "mutated" + source.PlatformDenylist[0] = "mutated" + source.HighRiskTools[0] = "mutated" + + require.Equal(t, []string{"read_file"}, p.policy.ToolWhitelist) + require.Equal(t, []string{"shell"}, p.policy.ToolDenylist) + require.Equal(t, []string{"secret"}, p.policy.ArgumentRedactionRules) + require.Equal(t, []string{"admin"}, p.policy.PlatformDenylist) + require.Equal(t, []string{"workspace_write"}, p.policy.HighRiskTools) +} + func newPolicy(t *testing.T, policy platform.ToolPolicy, opts ...Option) *Policy { t.Helper() policy = defaultPolicy(policy) diff --git a/platform/worker/builder.go b/platform/worker/builder.go index 2576f4efe5..74b30bf135 100644 --- a/platform/worker/builder.go +++ b/platform/worker/builder.go @@ -23,6 +23,7 @@ import ( "trpc.group/trpc-go/trpc-agent-go/platform/storagerouter" "trpc.group/trpc-go/trpc-agent-go/runner" "trpc.group/trpc-go/trpc-agent-go/session" + "trpc.group/trpc-go/trpc-agent-go/tool" ) // AgentDependencies contains tenant-scoped services available while building @@ -37,6 +38,12 @@ type AgentDependencies struct { Artifact artifact.Service Knowledge knowledge.Knowledge Audit platform.AuditSink + // ToolPolicy is the tenant app policy used to build the agent tool surface. + ToolPolicy platform.ToolPolicy + // ToolFilter narrows tools visible to the model for this runtime. + ToolFilter tool.FilterFunc + // ToolPermissionPolicy enforces tool-call authorization before execution. + ToolPermissionPolicy tool.PermissionPolicy } // AgentFactory builds an agent for one tenant app runtime. @@ -57,14 +64,26 @@ func (f AgentFactoryFunc) BuildAgent( // RuntimeBuilder assembles gateway runtimes from tenant storage profiles. type RuntimeBuilder struct { - router storagerouter.Router - factory AgentFactory + router storagerouter.Router + factory AgentFactory + toolPolicyProvider ToolPolicyProvider +} + +// RuntimeBuilderOption configures RuntimeBuilder. +type RuntimeBuilderOption func(*RuntimeBuilder) + +// WithToolPolicyProvider resolves configured app tool policies. +func WithToolPolicyProvider(provider ToolPolicyProvider) RuntimeBuilderOption { + return func(builder *RuntimeBuilder) { + builder.toolPolicyProvider = provider + } } // NewRuntimeBuilder creates a runtime builder. func NewRuntimeBuilder( router storagerouter.Router, factory AgentFactory, + opts ...RuntimeBuilderOption, ) (*RuntimeBuilder, error) { if isNilDependency(router) { return nil, ErrStorageRouterRequired @@ -72,10 +91,16 @@ func NewRuntimeBuilder( if isNilDependency(factory) { return nil, ErrAgentFactoryRequired } - return &RuntimeBuilder{ + builder := &RuntimeBuilder{ router: router, factory: factory, - }, nil + } + for _, opt := range opts { + if opt != nil { + opt(builder) + } + } + return builder, nil } // Build resolves tenant-scoped storage services, builds the configured agent, @@ -92,6 +117,10 @@ func (b *RuntimeBuilder) Build( if err := validateRuntimeConfig(tenant, app, binding); err != nil { return gateway.Runtime{}, err } + resolvedPolicy, err := b.resolveToolPolicyConfig(ctx, tenant, app) + if err != nil { + return gateway.Runtime{}, err + } storage, err := b.router.Adapter(ctx, tenant.TenantID, app.StorageProfileID) if err != nil { @@ -117,17 +146,28 @@ func (b *RuntimeBuilder) Build( if err != nil { return gateway.Runtime{}, fmt.Errorf("resolve audit sink: %w", err) } + permissionPolicy, err := compileToolPolicy(resolvedPolicy, auditSink) + if err != nil { + return gateway.Runtime{}, err + } + var toolFilter tool.FilterFunc + if permissionPolicy != nil { + toolFilter = permissionPolicy.ToolFilter() + } dependencies := AgentDependencies{ - Tenant: tenant, - App: app, - Binding: binding, - Storage: storage, - Session: sessionService, - Memory: memoryService, - Artifact: artifactService, - Knowledge: knowledgeService, - Audit: auditSink, + Tenant: tenant, + App: app, + Binding: binding, + Storage: storage, + Session: sessionService, + Memory: memoryService, + Artifact: artifactService, + Knowledge: knowledgeService, + Audit: auditSink, + ToolPolicy: resolvedPolicy, + ToolFilter: toolFilter, + ToolPermissionPolicy: permissionPolicy, } ag, err := b.factory.BuildAgent(ctx, dependencies) if err != nil { @@ -151,7 +191,9 @@ func (b *RuntimeBuilder) Build( runner.WithMemoryService(memoryService), runner.WithArtifactService(artifactService), ), - Audit: auditSink, + Audit: auditSink, + ToolFilter: toolFilter, + ToolPermissionPolicy: permissionPolicy, } if err := runtime.Validate(); err != nil { _ = runtime.Runner.Close() diff --git a/platform/worker/errors.go b/platform/worker/errors.go index b650e54c64..ca2316c59e 100644 --- a/platform/worker/errors.go +++ b/platform/worker/errors.go @@ -27,4 +27,8 @@ var ( ErrAgentRequired = errors.New("worker agent is required") // ErrAgentNameMismatch indicates that the built agent does not match app config. ErrAgentNameMismatch = errors.New("worker agent name does not match app config") + // ErrToolPolicyProviderRequired indicates that an app policy cannot be resolved. + ErrToolPolicyProviderRequired = errors.New("worker tool policy provider is required") + // ErrToolPolicyIdentityMismatch indicates that a resolved policy belongs elsewhere. + ErrToolPolicyIdentityMismatch = errors.New("worker tool policy identity mismatch") ) diff --git a/platform/worker/governance.go b/platform/worker/governance.go new file mode 100644 index 0000000000..1833163cf2 --- /dev/null +++ b/platform/worker/governance.go @@ -0,0 +1,104 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package worker + +import ( + "context" + "fmt" + "strings" + + "trpc.group/trpc-go/trpc-agent-go/platform" + "trpc.group/trpc-go/trpc-agent-go/platform/toolpolicy" +) + +// ToolPolicyProvider resolves a configured tenant app tool policy. +type ToolPolicyProvider interface { + ResolveToolPolicy( + ctx context.Context, + tenantID string, + appID string, + policyID string, + ) (platform.ToolPolicy, error) +} + +// ToolPolicyProviderFunc adapts a function into a ToolPolicyProvider. +type ToolPolicyProviderFunc func( + context.Context, + string, + string, + string, +) (platform.ToolPolicy, error) + +// ResolveToolPolicy implements ToolPolicyProvider. +func (f ToolPolicyProviderFunc) ResolveToolPolicy( + ctx context.Context, + tenantID string, + appID string, + policyID string, +) (platform.ToolPolicy, error) { + return f(ctx, tenantID, appID, policyID) +} + +func (b *RuntimeBuilder) resolveToolPolicyConfig( + ctx context.Context, + tenant platform.Tenant, + app platform.AgentApp, +) (platform.ToolPolicy, error) { + policyID := strings.TrimSpace(app.ToolPolicyID) + if policyID == "" { + return platform.ToolPolicy{}, nil + } + if isNilDependency(b.toolPolicyProvider) { + return platform.ToolPolicy{}, ErrToolPolicyProviderRequired + } + policy, err := b.toolPolicyProvider.ResolveToolPolicy( + ctx, + tenant.TenantID, + app.AppID, + policyID, + ) + if err != nil { + return platform.ToolPolicy{}, fmt.Errorf("resolve tool policy: %w", err) + } + if policy.TenantID != tenant.TenantID || + policy.AppID != app.AppID || + policy.PolicyID != policyID { + return platform.ToolPolicy{}, ErrToolPolicyIdentityMismatch + } + return cloneToolPolicy(policy), nil +} + +func compileToolPolicy( + policy platform.ToolPolicy, + auditSink platform.AuditSink, +) (*toolpolicy.Policy, error) { + if strings.TrimSpace(policy.PolicyID) == "" { + return nil, nil + } + compiled, err := toolpolicy.New( + policy, + toolpolicy.WithAuditSink(auditSink), + ) + if err != nil { + return nil, fmt.Errorf("compile tool policy: %w", err) + } + return compiled, nil +} + +func cloneToolPolicy(policy platform.ToolPolicy) platform.ToolPolicy { + policy.ToolWhitelist = append([]string(nil), policy.ToolWhitelist...) + policy.ToolDenylist = append([]string(nil), policy.ToolDenylist...) + policy.ArgumentRedactionRules = append( + []string(nil), + policy.ArgumentRedactionRules..., + ) + policy.PlatformDenylist = append([]string(nil), policy.PlatformDenylist...) + policy.HighRiskTools = append([]string(nil), policy.HighRiskTools...) + return policy +} diff --git a/platform/worker/governance_test.go b/platform/worker/governance_test.go new file mode 100644 index 0000000000..8718cf36cc --- /dev/null +++ b/platform/worker/governance_test.go @@ -0,0 +1,420 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package worker + +import ( + "context" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "trpc.group/trpc-go/trpc-agent-go/agent" + "trpc.group/trpc-go/trpc-agent-go/event" + memoryinmemory "trpc.group/trpc-go/trpc-agent-go/memory/inmemory" + "trpc.group/trpc-go/trpc-agent-go/model" + "trpc.group/trpc-go/trpc-agent-go/platform" + "trpc.group/trpc-go/trpc-agent-go/platform/artifactstore" + "trpc.group/trpc-go/trpc-agent-go/platform/gateway" + "trpc.group/trpc-go/trpc-agent-go/platform/storagerouter" + sessioninmemory "trpc.group/trpc-go/trpc-agent-go/session/inmemory" + "trpc.group/trpc-go/trpc-agent-go/tool" +) + +func TestRuntimeBuilderAppliesToolGovernanceEndToEnd(t *testing.T) { + ctx := context.Background() + router, auditSink := governanceTestRouter(t) + tenant, app, binding := governanceRuntimeConfig() + app.ToolPolicyID = "policy-a" + policy := platform.ToolPolicy{ + TenantID: tenant.TenantID, + AppID: app.AppID, + PolicyID: app.ToolPolicyID, + ToolWhitelist: []string{"read_file", "shell"}, + ToolDenylist: []string{"shell"}, + HighRiskTools: []string{"shell"}, + DangerousToolAction: platform.DangerousToolActionDeny, + } + providerCalled := false + provider := ToolPolicyProviderFunc(func( + _ context.Context, + tenantID string, + appID string, + policyID string, + ) (platform.ToolPolicy, error) { + providerCalled = true + assert.Equal(t, tenant.TenantID, tenantID) + assert.Equal(t, app.AppID, appID) + assert.Equal(t, app.ToolPolicyID, policyID) + return policy, nil + }) + var captured AgentDependencies + builder, err := NewRuntimeBuilder( + router, + AgentFactoryFunc(func( + _ context.Context, + dependencies AgentDependencies, + ) (agent.Agent, error) { + captured = dependencies + return newGovernanceProbeAgent(app.AgentName), nil + }), + WithToolPolicyProvider(provider), + ) + require.NoError(t, err) + runtime, err := builder.Build(ctx, tenant, app, binding) + require.NoError(t, err) + t.Cleanup(func() { + require.NoError(t, runtime.Runner.Close()) + }) + + assert.True(t, providerCalled) + assert.Equal(t, policy, captured.ToolPolicy) + require.NotNil(t, captured.ToolFilter) + require.NotNil(t, captured.ToolPermissionPolicy) + require.NotNil(t, runtime.ToolFilter) + require.NotNil(t, runtime.ToolPermissionPolicy) + + registry := gateway.NewInMemoryRegistry() + require.NoError(t, registry.Register(runtime)) + service := gateway.NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + gateway.NewInMemoryOutboundStore(), + ) + result, err := service.HandleInbound(ctx, governanceInbound(tenant, app, binding)) + require.NoError(t, err) + assert.Equal(t, "visible=read_file permission=deny", result.Outbound.Content) + + records := auditSink.Records() + require.Len(t, records, 2) + assert.Equal(t, "shell", records[0].ToolName) + assert.Equal(t, string(tool.PermissionActionDeny), records[0].Decision) + assert.Equal(t, tenant.TenantID, records[0].TenantID) + assert.Equal(t, app.AppID, records[0].AppID) + assert.Equal(t, "completed", records[1].Decision) +} + +func TestRuntimeBuilderRejectsMissingToolPolicyProviderBeforeStorage(t *testing.T) { + tenant, app, binding := governanceRuntimeConfig() + app.ToolPolicyID = "policy-a" + router := &countingRouter{Router: storagerouter.NewInMemoryRouter()} + builder, err := NewRuntimeBuilder( + router, + AgentFactoryFunc(func( + context.Context, + AgentDependencies, + ) (agent.Agent, error) { + return newGovernanceProbeAgent(app.AgentName), nil + }), + ) + require.NoError(t, err) + + _, err = builder.Build(context.Background(), tenant, app, binding) + assert.ErrorIs(t, err, ErrToolPolicyProviderRequired) + assert.Zero(t, router.adapterCalls) +} + +func TestRuntimeBuilderRejectsMismatchedToolPolicyBeforeStorage(t *testing.T) { + tenant, app, binding := governanceRuntimeConfig() + app.ToolPolicyID = "policy-a" + router := &countingRouter{Router: storagerouter.NewInMemoryRouter()} + builder, err := NewRuntimeBuilder( + router, + AgentFactoryFunc(func( + context.Context, + AgentDependencies, + ) (agent.Agent, error) { + return newGovernanceProbeAgent(app.AgentName), nil + }), + WithToolPolicyProvider(ToolPolicyProviderFunc(func( + context.Context, + string, + string, + string, + ) (platform.ToolPolicy, error) { + return platform.ToolPolicy{ + TenantID: "tenant-b", + AppID: app.AppID, + PolicyID: app.ToolPolicyID, + }, nil + })), + ) + require.NoError(t, err) + + _, err = builder.Build(context.Background(), tenant, app, binding) + assert.ErrorIs(t, err, ErrToolPolicyIdentityMismatch) + assert.Zero(t, router.adapterCalls) +} + +func TestRuntimeBuilderIsolatesResolvedAndCompiledToolPolicySlices( + t *testing.T, +) { + ctx := context.Background() + router, _ := governanceTestRouter(t) + tenant, app, binding := governanceRuntimeConfig() + app.ToolPolicyID = "policy-a" + source := platform.ToolPolicy{ + TenantID: tenant.TenantID, + AppID: app.AppID, + PolicyID: app.ToolPolicyID, + ToolWhitelist: []string{"read_file"}, + ToolDenylist: []string{"shell"}, + ArgumentRedactionRules: []string{"secret"}, + PlatformDenylist: []string{"admin"}, + HighRiskTools: []string{"shell"}, + DangerousToolAction: platform.DangerousToolActionDeny, + } + builder, err := NewRuntimeBuilder( + router, + AgentFactoryFunc(func( + _ context.Context, + dependencies AgentDependencies, + ) (agent.Agent, error) { + source.ToolWhitelist[0] = "provider_mutation" + source.ToolDenylist[0] = "provider_mutation" + source.ArgumentRedactionRules[0] = "provider_mutation" + source.PlatformDenylist[0] = "provider_mutation" + source.HighRiskTools[0] = "provider_mutation" + require.Equal(t, "read_file", dependencies.ToolPolicy.ToolWhitelist[0]) + require.Equal(t, "shell", dependencies.ToolPolicy.ToolDenylist[0]) + require.Equal(t, "secret", dependencies.ToolPolicy.ArgumentRedactionRules[0]) + require.Equal(t, "admin", dependencies.ToolPolicy.PlatformDenylist[0]) + require.Equal(t, "shell", dependencies.ToolPolicy.HighRiskTools[0]) + + dependencies.ToolPolicy.ToolWhitelist[0] = "factory_mutation" + dependencies.ToolPolicy.ToolDenylist[0] = "factory_mutation" + return newGovernanceProbeAgent(app.AgentName), nil + }), + WithToolPolicyProvider(ToolPolicyProviderFunc(func( + context.Context, + string, + string, + string, + ) (platform.ToolPolicy, error) { + return source, nil + })), + ) + require.NoError(t, err) + + runtime, err := builder.Build(ctx, tenant, app, binding) + require.NoError(t, err) + t.Cleanup(func() { + require.NoError(t, runtime.Runner.Close()) + }) + require.True(t, runtime.ToolFilter( + ctx, + &governanceProbeTool{name: "read_file"}, + )) + require.False(t, runtime.ToolFilter( + ctx, + &governanceProbeTool{name: "shell"}, + )) + decision, err := runtime.ToolPermissionPolicy.CheckToolPermission( + ctx, + &tool.PermissionRequest{ToolName: "shell"}, + ) + require.NoError(t, err) + require.Equal(t, tool.PermissionActionDeny, decision.Action) +} + +func governanceTestRouter( + t *testing.T, +) (*storagerouter.InMemoryRouter, *platform.InMemoryAuditSink) { + t.Helper() + tenantID := "tenant-a" + profileID := "profile-a" + backendID := "backend-a" + namespace := "tenant/" + tenantID + sessionService := sessioninmemory.NewSessionService() + t.Cleanup(func() { + require.NoError(t, sessionService.Close()) + }) + memoryService := memoryinmemory.NewMemoryService() + t.Cleanup(func() { + require.NoError(t, memoryService.Close()) + }) + artifactService, err := artifactstore.New(artifactstore.ServiceConfig{ + TenantID: tenantID, + Namespace: namespace, + MetadataStore: artifactstore.NewInMemoryMetadataStore(), + ObjectStore: artifactstore.NewInMemoryObjectStore(), + MaxAttempts: 2, + }) + require.NoError(t, err) + auditSink := platform.NewInMemoryAuditSink() + router := storagerouter.NewInMemoryRouter() + require.NoError(t, router.RegisterBackend(storagerouter.BackendSet{ + TenantID: tenantID, + BackendID: backendID, + Session: sessionService, + Summary: sessionService, + Memory: memoryService, + Artifact: artifactService, + Knowledge: &stubKnowledge{}, + Audit: auditSink, + })) + require.NoError(t, router.RegisterProfile(platform.StorageProfile{ + TenantID: tenantID, + ProfileID: profileID, + SessionBackend: backendID, + MemoryBackend: backendID, + SummaryBackend: backendID, + ArtifactBackend: backendID, + KnowledgeBackend: backendID, + AuditBackend: backendID, + DSNRef: "secret://storage/" + tenantID, + Namespace: namespace, + })) + return router, auditSink +} + +func governanceRuntimeConfig() ( + platform.Tenant, + platform.AgentApp, + platform.ChannelBinding, +) { + tenant := platform.Tenant{ + TenantID: "tenant-a", + Status: platform.TenantStatusActive, + } + app := platform.AgentApp{ + TenantID: tenant.TenantID, + AppID: "app-a", + AppName: "app", + AgentName: "governance-probe", + StorageProfileID: "profile-a", + Status: platform.AppStatusActive, + } + binding := platform.ChannelBinding{ + TenantID: tenant.TenantID, + AppID: app.AppID, + BindingID: "binding-a", + Channel: "wecom", + AccountID: "account-a", + WebhookPath: "/callback", + TokenRef: "secret://token", + SecretRef: "secret://secret", + Status: platform.BindingStatusActive, + } + return tenant, app, binding +} + +func governanceInbound( + tenant platform.Tenant, + app platform.AgentApp, + binding platform.ChannelBinding, +) platform.InboundMessage { + return platform.InboundMessage{ + TenantID: tenant.TenantID, + AppID: app.AppID, + BindingID: binding.BindingID, + Channel: binding.Channel, + ChannelAccountID: binding.AccountID, + PlatformMessageID: "governance-message", + ExternalUserID: "external-user", + ConversationType: platform.ConversationTypeDM, + MessageType: platform.MessageTypeText, + ContentParts: []platform.ContentPart{ + {Type: platform.ContentPartTypeText, Text: "check governance"}, + }, + ReceivedAt: time.Unix(100, 0), + } +} + +type governanceProbeAgent struct { + name string + tools []tool.Tool +} + +func newGovernanceProbeAgent(name string) *governanceProbeAgent { + return &governanceProbeAgent{ + name: name, + tools: []tool.Tool{ + &governanceProbeTool{name: "read_file"}, + &governanceProbeTool{name: "shell"}, + }, + } +} + +func (a *governanceProbeAgent) Run( + ctx context.Context, + invocation *agent.Invocation, +) (<-chan *event.Event, error) { + visible := tool.FilterTools( + ctx, + a.tools, + invocation.RunOptions.MandatoryToolFilter, + ) + visibleNames := make([]string, 0, len(visible)) + for _, candidate := range visible { + visibleNames = append(visibleNames, candidate.Declaration().Name) + } + shell := a.tools[1] + decision, err := invocation.RunOptions.CheckToolPermission( + ctx, + &tool.PermissionRequest{ + Tool: shell, + ToolName: shell.Declaration().Name, + ToolCallID: "call-shell", + Declaration: shell.Declaration(), + Arguments: []byte(`{"command":"restricted"}`), + }, + ) + if err != nil { + return nil, err + } + out := make(chan *event.Event, 1) + out <- event.NewResponseEvent( + invocation.InvocationID, + a.name, + &model.Response{ + ID: "governance-probe-response", + Object: model.ObjectTypeChatCompletion, + Done: true, + Choices: []model.Choice{ + { + Index: 0, + Message: model.Message{ + Role: model.RoleAssistant, + Content: "visible=" + strings.Join(visibleNames, ",") + + " permission=" + string(decision.Action), + }, + }, + }, + }, + ) + close(out) + return out, nil +} + +func (a *governanceProbeAgent) Tools() []tool.Tool { + return a.tools +} + +func (a *governanceProbeAgent) Info() agent.Info { + return agent.Info{Name: a.name} +} + +func (a *governanceProbeAgent) SubAgents() []agent.Agent { + return nil +} + +func (a *governanceProbeAgent) FindSubAgent(string) agent.Agent { + return nil +} + +type governanceProbeTool struct { + name string +} + +func (t *governanceProbeTool) Declaration() *tool.Declaration { + return &tool.Declaration{Name: t.name} +} diff --git a/tool/agent/agent_tool.go b/tool/agent/agent_tool.go index d19c4b3e7a..061e188d62 100644 --- a/tool/agent/agent_tool.go +++ b/tool/agent/agent_tool.go @@ -1802,7 +1802,23 @@ func (at *Tool) fallbackRunnerRunOptions(ctx context.Context) []agent.RunOption if !ok || parentInv == nil { return nil } - opts := make([]agent.RunOption, 0, 3) + opts := make([]agent.RunOption, 0, 5) + if parentInv.RunOptions.MandatoryToolFilter != nil { + opts = append( + opts, + agent.WithMandatoryToolFilter( + parentInv.RunOptions.MandatoryToolFilter, + ), + ) + } + if parentInv.RunOptions.MandatoryToolPermissionPolicy != nil { + opts = append( + opts, + agent.WithMandatoryToolPermissionPolicy( + parentInv.RunOptions.MandatoryToolPermissionPolicy, + ), + ) + } if agent.IsGraphCompletionEventDisabled(parentInv) { opts = append(opts, agent.WithDisableGraphCompletionEvent(true)) } diff --git a/tool/agent/agent_tool_test.go b/tool/agent/agent_tool_test.go index 203fb458d1..3dbad62a19 100644 --- a/tool/agent/agent_tool_test.go +++ b/tool/agent/agent_tool_test.go @@ -4508,6 +4508,24 @@ func TestTool_FallbackRunnerRunOptions_PreserveOnlyCompatibilityControls(t *test agent.WithInvocationRunOptions(agent.NewRunOptions( agent.WithStreamMode(agent.StreamModeUpdates), agent.WithGraphEmitFinalModelResponses(true), + agent.WithToolFilter(func(context.Context, tool.Tool) bool { + return true + }), + agent.WithToolPermissionPolicyFunc( + func(context.Context, *tool.PermissionRequest) (tool.PermissionDecision, error) { + return tool.AllowPermission(), nil + }, + ), + agent.WithMandatoryToolFilter( + func(context.Context, tool.Tool) bool { + return true + }, + ), + agent.WithMandatoryToolPermissionPolicyFunc( + func(context.Context, *tool.PermissionRequest) (tool.PermissionDecision, error) { + return tool.AllowPermission(), nil + }, + ), agent.WithDisableGraphCompletionEvent(true), agent.WithDisableGraphExecutorEvents(true), agent.WithEventChannelBufferSize(7), @@ -4519,6 +4537,10 @@ func TestTool_FallbackRunnerRunOptions_PreserveOnlyCompatibilityControls(t *test require.False(t, runOptions.StreamModeEnabled) require.False(t, runOptions.GraphEmitFinalModelResponses) + require.Nil(t, runOptions.ToolFilter) + require.Nil(t, runOptions.ToolPermissionPolicy) + require.NotNil(t, runOptions.MandatoryToolFilter) + require.NotNil(t, runOptions.MandatoryToolPermissionPolicy) require.True(t, agent.IsGraphCompletionEventDisabled(child)) require.True(t, agent.IsGraphExecutorEventsDisabled(child)) require.Equal(t, 7, agent.GetEventChannelBufferSize(child)) diff --git a/tool/agent/dynamic_tool.go b/tool/agent/dynamic_tool.go index 0e84c23625..b433d0828b 100644 --- a/tool/agent/dynamic_tool.go +++ b/tool/agent/dynamic_tool.go @@ -1058,10 +1058,11 @@ func (at *Tool) dynamicChildInvocationOptions( // - prompt: Instruction/GlobalInstruction outrank the template prompt — a // model-provided instruction still applies because it travels via the // surface patch, which is resolved before RunOptions; -// - execution: CodeExecutor, plus the execution-policy filters +// - execution: CodeExecutor, plus the ordinary execution-policy filters // (ToolExecutionFilter defers tool calls, ToolPermissionPolicy gates them) // which have no natural external-continuation channel for a synchronous -// sub-agent. Inheriting these requires an explicit future Option. +// sub-agent. MandatoryToolFilter and MandatoryToolPermissionPolicy are +// intentionally preserved as non-negotiable parent governance. func (at *Tool) sanitizeChildRunOptions( runOpts *agent.RunOptions, enforceTemplateBoundary bool, diff --git a/tool/agent/dynamic_tool_test.go b/tool/agent/dynamic_tool_test.go index 3e8dc7a8f7..76cdbc6a7a 100644 --- a/tool/agent/dynamic_tool_test.go +++ b/tool/agent/dynamic_tool_test.go @@ -18,6 +18,7 @@ import ( "sort" "strings" "sync" + "sync/atomic" "testing" "time" @@ -972,6 +973,79 @@ func (m *dynRecordingModel) snapshot() [][]string { return out } +type dynToolCallModel struct { + name string + toolName string + calls atomic.Int32 +} + +func (m *dynToolCallModel) GenerateContent( + _ context.Context, + _ *model.Request, +) (<-chan *model.Response, error) { + call := m.calls.Add(1) + response := &model.Response{Done: true} + if call == 1 { + response.Choices = []model.Choice{{ + Message: model.Message{ + Role: model.RoleAssistant, + ToolCalls: []model.ToolCall{{ + ID: "call-restricted", + Type: "function", + Function: model.FunctionDefinitionParam{ + Name: m.toolName, + Arguments: []byte(`{}`), + }, + }}, + }, + }} + } else { + response.Choices = []model.Choice{{ + Message: model.NewAssistantMessage("child-done"), + }} + } + ch := make(chan *model.Response, 1) + ch <- response + close(ch) + return ch, nil +} + +func (m *dynToolCallModel) Info() model.Info { + return model.Info{Name: m.name} +} + +type dynActivationSurfaceAgent struct { + agent.Agent + disabled string +} + +func (a *dynActivationSurfaceAgent) InvocationToolSurface( + ctx context.Context, + inv *agent.Invocation, +) ([]tool.Tool, map[string]bool) { + provider := a.Agent.(agent.InvocationToolSurfaceProvider) + return provider.InvocationToolSurface(ctx, inv) +} + +func (a *dynActivationSurfaceAgent) ApplyInvocationToolActivation( + _ context.Context, + _ *agent.Invocation, + tools []tool.Tool, + userToolNames map[string]bool, + externalToolNames map[string]bool, +) ([]tool.Tool, map[string]bool, map[string]bool) { + filtered := make([]tool.Tool, 0, len(tools)) + for _, candidate := range tools { + if candidate.Declaration().Name == a.disabled { + delete(userToolNames, a.disabled) + delete(externalToolNames, a.disabled) + continue + } + filtered = append(filtered, candidate) + } + return filtered, userToolNames, externalToolNames +} + type dynBlockingModel struct{} func (m *dynBlockingModel) GenerateContent( @@ -1060,6 +1134,42 @@ func TestNewDynamicTool_Integration_DefaultAllTools(t *testing.T) { require.Equal(t, []string{"tool_a", "tool_b"}, seen[0]) } +func TestNewDynamicTool_Integration_AppliesParentInvocationActivation( + t *testing.T, +) { + parentModel := &dynRecordingModel{name: "parent", response: "unused"} + parentBase := llmagent.New( + "main", + llmagent.WithModel(parentModel), + llmagent.WithTools([]tool.Tool{ + newDynTestTool("tool_a"), + newDynTestTool("tool_b"), + }), + ) + parent := &dynActivationSurfaceAgent{ + Agent: parentBase, + disabled: "tool_b", + } + templateModel := &dynRecordingModel{name: "template", response: "done"} + template := llmagent.New("subagent", llmagent.WithModel(templateModel)) + at := NewDynamicTool(WithTemplateAgent(template)) + inv := agent.NewInvocation( + agent.WithInvocationAgent(parent), + agent.WithInvocationSession(session.NewSession("app", "user", "session")), + agent.WithInvocationEventFilterKey("main"), + ) + ctx := agent.NewInvocationContext(context.Background(), inv) + + got, err := at.Call(ctx, []byte(`{"request":"use available tools"}`)) + require.NoError(t, err) + require.Equal(t, "done", got) + require.Empty(t, parentModel.snapshot()) + seen := templateModel.snapshot() + require.Len(t, seen, 1) + require.Equal(t, []string{"tool_a"}, seen[0], + "dynamic child must use the parent's activated tool surface") +} + func TestNewDynamicTool_Integration_UnavailableReasonReturnedToParent(t *testing.T) { recModel := &dynRecordingModel{name: "rec", response: "child-done"} main := llmagent.New("main", llmagent.WithModel(recModel)) @@ -1172,6 +1282,179 @@ func TestNewDynamicTool_Integration_WithTemplateAgent(t *testing.T) { "tools selected from the parent surface must be injected into the template") } +func TestNewDynamicTool_Integration_TemplatePreservesMandatoryPermissionPolicy( + t *testing.T, +) { + const restrictedToolName = "restricted_tool" + var executions atomic.Int32 + restrictedTool := function.NewFunctionTool( + func(_ context.Context, _ struct{}) (string, error) { + executions.Add(1) + return "executed", nil + }, + function.WithName(restrictedToolName), + function.WithDescription("must remain tenant-governed"), + ) + parentModel := &dynRecordingModel{ + name: "parent", + response: "parent-should-not-run", + } + templateModel := &dynToolCallModel{ + name: "template", + toolName: restrictedToolName, + } + main := llmagent.New( + "main", + llmagent.WithModel(parentModel), + llmagent.WithTools([]tool.Tool{restrictedTool}), + ) + subTemplate := llmagent.New( + "subagent", + llmagent.WithModel(templateModel), + ) + at := NewDynamicTool(WithTemplateAgent(subTemplate)) + + sess := session.NewSession("app", "user", "session") + parent := agent.NewInvocation( + agent.WithInvocationAgent(main), + agent.WithInvocationSession(sess), + agent.WithInvocationEventFilterKey("main"), + agent.WithInvocationRunOptions(agent.NewRunOptions( + agent.WithMandatoryToolPermissionPolicyFunc( + func( + _ context.Context, + req *tool.PermissionRequest, + ) (tool.PermissionDecision, error) { + if req.ToolName == restrictedToolName { + return tool.DenyPermission("tenant policy"), nil + } + return tool.AllowPermission(), nil + }, + ), + )), + ) + ctx := agent.NewInvocationContext(context.Background(), parent) + + got, err := at.Call( + ctx, + []byte(`{"request":"run restricted tool","tools":["restricted_tool"]}`), + ) + require.NoError(t, err) + require.Equal(t, "child-done", got) + require.Equal(t, int32(0), executions.Load(), + "template child must not clear tenant mandatory permission policy") + require.Equal(t, int32(2), templateModel.calls.Load(), + "permission denial should be returned to the child model") + require.Empty(t, parentModel.snapshot(), + "template model must remain the child execution boundary") +} + +func TestTool_Integration_FallbackRunnerPreservesMandatoryPermissionPolicy( + t *testing.T, +) { + const restrictedToolName = "restricted_tool" + var executions atomic.Int32 + restrictedTool := function.NewFunctionTool( + func(_ context.Context, _ struct{}) (string, error) { + executions.Add(1) + return "executed", nil + }, + function.WithName(restrictedToolName), + function.WithDescription("must remain parent-governed"), + ) + childModel := &dynToolCallModel{ + name: "child", + toolName: restrictedToolName, + } + child := llmagent.New( + "child", + llmagent.WithModel(childModel), + llmagent.WithTools([]tool.Tool{restrictedTool}), + ) + at := NewTool(child) + parent := agent.NewInvocation( + agent.WithInvocationRunOptions(agent.NewRunOptions( + agent.WithMandatoryToolPermissionPolicyFunc( + func( + _ context.Context, + req *tool.PermissionRequest, + ) (tool.PermissionDecision, error) { + if req.ToolName == restrictedToolName { + return tool.DenyPermission("parent policy"), nil + } + return tool.AllowPermission(), nil + }, + ), + )), + ) + ctx := agent.NewInvocationContext(context.Background(), parent) + + got, err := at.Call(ctx, []byte(`{"request":"run restricted tool"}`)) + require.NoError(t, err) + require.Equal(t, "child-done", got) + require.Equal(t, int32(0), executions.Load()) + require.Equal(t, int32(2), childModel.calls.Load()) +} + +func TestTool_Integration_FallbackRunnerPreservesMandatoryToolFilter( + t *testing.T, +) { + for _, stream := range []bool{false, true} { + name := "sync" + if stream { + name = "stream" + } + t.Run(name, func(t *testing.T) { + childModel := &dynRecordingModel{name: "child", response: "done"} + child := llmagent.New( + "child", + llmagent.WithModel(childModel), + llmagent.WithTools([]tool.Tool{ + newDynTestTool("allowed_tool"), + newDynTestTool("hidden_tool"), + }), + ) + var opts []Option + if stream { + opts = append(opts, WithStreamInner(true)) + } + at := NewTool(child, opts...) + parent := agent.NewInvocation( + agent.WithInvocationRunOptions(agent.NewRunOptions( + agent.WithMandatoryToolFilter( + tool.NewIncludeToolNamesFilter("allowed_tool"), + ), + )), + ) + ctx := agent.NewInvocationContext(context.Background(), parent) + + if stream { + reader, err := at.StreamableCall( + ctx, + []byte(`{"request":"list tools"}`), + ) + require.NoError(t, err) + defer reader.Close() + for { + _, recvErr := reader.Recv() + if recvErr == io.EOF { + break + } + require.NoError(t, recvErr) + } + } else { + got, err := at.Call(ctx, []byte(`{"request":"list tools"}`)) + require.NoError(t, err) + require.Equal(t, "done", got) + } + + seen := childModel.snapshot() + require.Len(t, seen, 1) + require.Equal(t, []string{"allowed_tool"}, seen[0]) + }) + } +} + // TestNewDynamicTool_Integration_ExcludesSelf ensures the dynamic tool never // leaks itself into the child surface, preventing runaway recursion. func TestNewDynamicTool_Integration_ExcludesSelf(t *testing.T) { @@ -1814,12 +2097,13 @@ func TestChildCodeExecutor_TemplatePassesNonNilInvocation(t *testing.T) { func fullyPopulatedChildRunOptions() agent.RunOptions { return agent.RunOptions{ - AdditionalTools: stubTools("extra"), - ExternalTools: stubTools("ext"), - ExternalToolNames: map[string]bool{"ext": true}, - ToolFilter: func(context.Context, tool.Tool) bool { return true }, - Model: &dynRecordingModel{name: "m"}, - ModelName: "mname", + AdditionalTools: stubTools("extra"), + ExternalTools: stubTools("ext"), + ExternalToolNames: map[string]bool{"ext": true}, + ToolFilter: func(context.Context, tool.Tool) bool { return true }, + MandatoryToolFilter: func(context.Context, tool.Tool) bool { return true }, + Model: &dynRecordingModel{name: "m"}, + ModelName: "mname", ModelSelector: func(context.Context, *agent.Invocation) (model.Model, error) { return nil, nil }, @@ -1832,6 +2116,11 @@ func fullyPopulatedChildRunOptions() agent.RunOptions { return tool.AllowPermission(), nil }, ), + MandatoryToolPermissionPolicy: tool.PermissionPolicyFunc( + func(context.Context, *tool.PermissionRequest) (tool.PermissionDecision, error) { + return tool.DenyPermission("tenant policy"), nil + }, + ), } } @@ -1859,6 +2148,8 @@ func TestSanitizeChildRunOptions_TemplateBoundaryClearsAll(t *testing.T) { require.Nil(t, runOpts.CodeExecutor) require.Nil(t, runOpts.ToolExecutionFilter) require.Nil(t, runOpts.ToolPermissionPolicy) + require.NotNil(t, runOpts.MandatoryToolFilter) + require.NotNil(t, runOpts.MandatoryToolPermissionPolicy) } // TestSanitizeChildRunOptions_NoTemplateKeepsBoundaryFields verifies that @@ -1884,6 +2175,8 @@ func TestSanitizeChildRunOptions_NoTemplateKeepsBoundaryFields(t *testing.T) { require.NotNil(t, runOpts.CodeExecutor) require.NotNil(t, runOpts.ToolExecutionFilter) require.NotNil(t, runOpts.ToolPermissionPolicy) + require.NotNil(t, runOpts.MandatoryToolFilter) + require.NotNil(t, runOpts.MandatoryToolPermissionPolicy) } // TestNewDynamicTool_Integration_CapabilityToolsBypassParentFilter is this diff --git a/tool/dynamicworkflow/tool.go b/tool/dynamicworkflow/tool.go index 9d3beaeccc..806ce6962a 100644 --- a/tool/dynamicworkflow/tool.go +++ b/tool/dynamicworkflow/tool.go @@ -25,6 +25,7 @@ import ( "trpc.group/trpc-go/trpc-agent-go/internal/state/eventstream" "trpc.group/trpc-go/trpc-agent-go/internal/state/flush" "trpc.group/trpc-go/trpc-agent-go/internal/state/livesession" + itool "trpc.group/trpc-go/trpc-agent-go/internal/tool" "trpc.group/trpc-go/trpc-agent-go/model" "trpc.group/trpc-go/trpc-agent-go/tool" ) @@ -230,7 +231,13 @@ func (g *workflowGateway) callTool(ctx context.Context, call Call) (json.RawMess if err := json.Unmarshal(call.Args, &args); err != nil || args == nil { return nil, fmt.Errorf("dynamicworkflow: tool %q requires a JSON object argument", call.Name) } - permissionResult, err := g.checkToolPermission(ctx, call, candidate) + permissionResult, err := g.checkMandatoryToolVisibility(ctx, call, candidate) + if err != nil { + return nil, fmt.Errorf("dynamicworkflow: check visibility for tool %q: %w", call.Name, err) + } + if permissionResult == nil { + permissionResult, err = g.checkToolPermission(ctx, call, candidate) + } if err != nil { return nil, fmt.Errorf("dynamicworkflow: check permission for tool %q: %w", call.Name, err) } @@ -252,19 +259,39 @@ func (g *workflowGateway) callTool(ctx context.Context, call Call) (json.RawMess return raw, nil } -func (g *workflowGateway) checkToolPermission( +func (g *workflowGateway) checkMandatoryToolVisibility( ctx context.Context, call Call, candidate tool.CallableTool, ) (*tool.PermissionResult, error) { - req := &tool.PermissionRequest{ - Tool: candidate, - ToolName: call.Name, - ToolCallID: call.ID, - Declaration: candidate.Declaration(), - Arguments: append([]byte(nil), call.Args...), - Metadata: tool.MetadataOf(candidate), + if g == nil || g.parent == nil || g.parent.RunOptions.MandatoryToolFilter == nil { + return nil, nil + } + if g.parent.RunOptions.MandatoryToolFilter( + ctx, + itool.ResolveDeclaration(candidate), + ) { + return nil, nil } + req := workflowToolPermissionRequest(call, candidate) + return normalizeWorkflowToolPermissionResult( + req, + tool.DenyPermission( + fmt.Sprintf( + "tool %q is hidden by mandatory tool filter", + req.ToolName, + ), + ), + nil, + ) +} + +func (g *workflowGateway) checkToolPermission( + ctx context.Context, + call Call, + candidate tool.CallableTool, +) (*tool.PermissionResult, error) { + req := workflowToolPermissionRequest(call, candidate) if checker, ok := candidate.(tool.PermissionChecker); ok { decision, err := checker.CheckPermission(ctx, req) result, err := normalizeWorkflowToolPermissionResult(req, decision, err) @@ -272,13 +299,27 @@ func (g *workflowGateway) checkToolPermission( return result, err } } - if g == nil || g.parent == nil || g.parent.RunOptions.ToolPermissionPolicy == nil { + if g == nil || g.parent == nil { return nil, nil } - decision, err := g.parent.RunOptions.ToolPermissionPolicy.CheckToolPermission(ctx, req) + decision, err := g.parent.RunOptions.CheckToolPermission(ctx, req) return normalizeWorkflowToolPermissionResult(req, decision, err) } +func workflowToolPermissionRequest( + call Call, + candidate tool.CallableTool, +) *tool.PermissionRequest { + return &tool.PermissionRequest{ + Tool: candidate, + ToolName: call.Name, + ToolCallID: call.ID, + Declaration: candidate.Declaration(), + Arguments: append([]byte(nil), call.Args...), + Metadata: tool.MetadataOf(candidate), + } +} + func normalizeWorkflowToolPermissionResult( req *tool.PermissionRequest, decision tool.PermissionDecision, diff --git a/tool/dynamicworkflow/tool_test.go b/tool/dynamicworkflow/tool_test.go index fd3c9028c4..3b9e2def6f 100644 --- a/tool/dynamicworkflow/tool_test.go +++ b/tool/dynamicworkflow/tool_test.go @@ -503,6 +503,57 @@ func TestWorkflowCoordinatesExplicitAgentAndTool(t *testing.T) { func TestWorkflowCallToolHonorsPermissionBoundaries(t *testing.T) { reviewer := &testAgent{name: "reviewer"} + t.Run("mandatory filter deny skips checker policy and execution", func(t *testing.T) { + sensitive := &permissionTestTool{ + name: "sensitive", + decision: tool.AllowPermission(), + } + workflow, err := NewTool(scriptedRuntime{run: func(ctx context.Context, handler CallHandler) (Result, error) { + raw, err := handler.HandleWorkflowCall(ctx, Call{ + ID: "tool-1", Kind: CallKindTool, Name: "sensitive", Args: json.RawMessage(`{"id":"42"}`), + }) + return Result{Value: raw}, err + }}, []agent.Agent{reviewer}, WithCodeCallableTools(sensitive)) + require.NoError(t, err) + + filterCalled := false + policyCalled := false + parent := agent.NewInvocation( + agent.WithInvocationSession(&session.Session{ID: "session-1", AppName: "app", UserID: "user"}), + agent.WithInvocationRunOptions(agent.NewRunOptions( + agent.WithMandatoryToolFilter( + func(_ context.Context, candidate tool.Tool) bool { + filterCalled = true + require.Equal(t, "sensitive", candidate.Declaration().Name) + return false + }, + ), + agent.WithToolPermissionPolicyFunc( + func(context.Context, *tool.PermissionRequest) (tool.PermissionDecision, error) { + policyCalled = true + return tool.AllowPermission(), nil + }, + ), + )), + ) + + value, err := workflow.Call( + agent.NewInvocationContext(context.Background(), parent), + []byte(`{"code":"return None"}`), + ) + require.NoError(t, err) + result := value.(Result) + require.JSONEq( + t, + `{"status":"denied","tool":"sensitive","reason":"tool \"sensitive\" is hidden by mandatory tool filter"}`, + string(result.Value), + ) + require.True(t, filterCalled) + require.False(t, sensitive.checkerCalled) + require.False(t, policyCalled) + require.False(t, sensitive.called) + }) + t.Run("tool checker deny skips execution and run policy", func(t *testing.T) { sensitive := &permissionTestTool{ name: "sensitive", @@ -563,6 +614,55 @@ func TestWorkflowCallToolHonorsPermissionBoundaries(t *testing.T) { require.JSONEq(t, `{"status":"approval_required","tool":"sensitive","reason":"needs approval"}`, string(result.Value)) require.False(t, sensitive.called) }) + + t.Run("mandatory deny cannot be overridden by run policy", func(t *testing.T) { + sensitive := &permissionTestTool{name: "sensitive", decision: tool.AllowPermission()} + workflow, err := NewTool(scriptedRuntime{run: func(ctx context.Context, handler CallHandler) (Result, error) { + raw, err := handler.HandleWorkflowCall(ctx, Call{ + ID: "tool-1", Kind: CallKindTool, Name: "sensitive", Args: json.RawMessage(`{"id":"42"}`), + }) + return Result{Value: raw}, err + }}, []agent.Agent{reviewer}, WithCodeCallableTools(sensitive)) + require.NoError(t, err) + + ordinaryCalled := false + parent := agent.NewInvocation( + agent.WithInvocationSession(&session.Session{ID: "session-1", AppName: "app", UserID: "user"}), + agent.WithInvocationRunOptions(agent.NewRunOptions( + agent.WithMandatoryToolPermissionPolicyFunc( + func( + context.Context, + *tool.PermissionRequest, + ) (tool.PermissionDecision, error) { + return tool.DenyPermission("tenant policy"), nil + }, + ), + agent.WithToolPermissionPolicyFunc( + func( + context.Context, + *tool.PermissionRequest, + ) (tool.PermissionDecision, error) { + ordinaryCalled = true + return tool.AllowPermission(), nil + }, + ), + )), + ) + + value, err := workflow.Call( + agent.NewInvocationContext(context.Background(), parent), + []byte(`{"code":"return None"}`), + ) + require.NoError(t, err) + result := value.(Result) + require.JSONEq( + t, + `{"status":"denied","tool":"sensitive","reason":"tenant policy"}`, + string(result.Value), + ) + require.False(t, sensitive.called) + require.False(t, ordinaryCalled) + }) } func TestWorkflowChildAgentToolsHonorParentPermissionPolicy(t *testing.T) { @@ -1891,10 +1991,11 @@ func (a *schemaTestAgent) SubAgents() []agent.Agent { return nil } func (a *schemaTestAgent) FindSubAgent(string) agent.Agent { return nil } type permissionTestTool struct { - name string - decision tool.PermissionDecision - err error - called bool + name string + decision tool.PermissionDecision + err error + called bool + checkerCalled bool } func (t *permissionTestTool) Declaration() *tool.Declaration { @@ -1910,6 +2011,7 @@ func (t *permissionTestTool) CheckPermission( context.Context, *tool.PermissionRequest, ) (tool.PermissionDecision, error) { + t.checkerCalled = true return t.decision, t.err } From 541dd0962484ae760fe787c624fe21c7ac4cab81 Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 15:21:47 +0800 Subject: [PATCH 58/95] fix lint failures in phase2 tool surface stack --- graph/state_graph.go | 88 +++++++++++++++++++++-------- internal/flow/llmflow/llmflow.go | 3 +- internal/toolsurface/toolsurface.go | 1 - 3 files changed, 67 insertions(+), 25 deletions(-) diff --git a/graph/state_graph.go b/graph/state_graph.go index e5450ba16e..c371d07b56 100644 --- a/graph/state_graph.go +++ b/graph/state_graph.go @@ -4493,14 +4493,16 @@ func runToolWithEventContexts( ) completeInvocation = invocationFromContextOrFallback(ctx, completeInvocation) if err != nil { - if customResult != nil { - return startCtx, startInvocation, ctx, completeInvocation, customResult, toolCall.Function.Arguments, err - } - var interruptErr *InterruptError - if errors.As(err, &interruptErr) { - return startCtx, startInvocation, ctx, completeInvocation, result, toolCall.Function.Arguments, err - } - return startCtx, startInvocation, ctx, completeInvocation, nil, toolCall.Function.Arguments, err + return toolCallbackErrorResult( + startCtx, + startInvocation, + ctx, + completeInvocation, + result, + customResult, + toolCall.Function.Arguments, + err, + ) } if customResult != nil { return startCtx, startInvocation, ctx, completeInvocation, customResult, toolCall.Function.Arguments, nil @@ -4516,30 +4518,72 @@ func runToolWithEventContexts( ) completeInvocation = invocationFromContextOrFallback(ctx, completeInvocation) if err != nil { - if customResult != nil { - return startCtx, startInvocation, ctx, completeInvocation, customResult, toolCall.Function.Arguments, err - } - var interruptErr *InterruptError - if errors.As(err, &interruptErr) { - return startCtx, startInvocation, ctx, completeInvocation, result, toolCall.Function.Arguments, err - } - return startCtx, startInvocation, ctx, completeInvocation, nil, toolCall.Function.Arguments, err + return toolCallbackErrorResult( + startCtx, + startInvocation, + ctx, + completeInvocation, + result, + customResult, + toolCall.Function.Arguments, + err, + ) } if customResult != nil { return startCtx, startInvocation, ctx, completeInvocation, customResult, toolCall.Function.Arguments, nil } if toolErr != nil { - var interruptErr *InterruptError - if errors.As(toolErr, &interruptErr) { - return startCtx, startInvocation, ctx, completeInvocation, result, toolCall.Function.Arguments, toolErr - } - return startCtx, startInvocation, ctx, completeInvocation, nil, toolCall.Function.Arguments, - fmt.Errorf("tool %s call failed: %w", toolCall.Function.Name, toolErr) + return toolRunErrorResult( + startCtx, + startInvocation, + ctx, + completeInvocation, + result, + toolCall, + toolErr, + ) } return startCtx, startInvocation, ctx, completeInvocation, result, toolCall.Function.Arguments, nil } +func toolCallbackErrorResult( + startCtx context.Context, + startInvocation *agent.Invocation, + completeCtx context.Context, + completeInvocation *agent.Invocation, + result any, + customResult any, + modifiedArgs []byte, + err error, +) (context.Context, *agent.Invocation, context.Context, *agent.Invocation, any, []byte, error) { + if customResult != nil { + return startCtx, startInvocation, completeCtx, completeInvocation, customResult, modifiedArgs, err + } + var interruptErr *InterruptError + if errors.As(err, &interruptErr) { + return startCtx, startInvocation, completeCtx, completeInvocation, result, modifiedArgs, err + } + return startCtx, startInvocation, completeCtx, completeInvocation, nil, modifiedArgs, err +} + +func toolRunErrorResult( + startCtx context.Context, + startInvocation *agent.Invocation, + completeCtx context.Context, + completeInvocation *agent.Invocation, + result any, + toolCall model.ToolCall, + err error, +) (context.Context, *agent.Invocation, context.Context, *agent.Invocation, any, []byte, error) { + var interruptErr *InterruptError + if errors.As(err, &interruptErr) { + return startCtx, startInvocation, completeCtx, completeInvocation, result, toolCall.Function.Arguments, err + } + return startCtx, startInvocation, completeCtx, completeInvocation, nil, toolCall.Function.Arguments, + fmt.Errorf("tool %s call failed: %w", toolCall.Function.Name, err) +} + func agentToolGraphRuntimeContext( invocation *agent.Invocation, state State, diff --git a/internal/flow/llmflow/llmflow.go b/internal/flow/llmflow/llmflow.go index c1f825ac0f..ac5bb92155 100644 --- a/internal/flow/llmflow/llmflow.go +++ b/internal/flow/llmflow/llmflow.go @@ -1991,7 +1991,7 @@ func (f *Flow) getFilteredTools( hasUserToolTracking, userToolNames, ) - allTools, userToolNames, hasUserToolTracking, externalToolNames := + allTools, userToolNames, _, externalToolNames := toolsurface.AppendRunOptionTools( allTools, userToolNames, @@ -2023,7 +2023,6 @@ func (f *Flow) getFilteredTools( userToolNames, externalToolNames, ) - hasUserToolTracking = userToolNames != nil } allTools, userToolNames, externalToolNames = toolsurface.ApplyMandatoryToolFilter( diff --git a/internal/toolsurface/toolsurface.go b/internal/toolsurface/toolsurface.go index 89530072b5..fdab7ebd64 100644 --- a/internal/toolsurface/toolsurface.go +++ b/internal/toolsurface/toolsurface.go @@ -287,7 +287,6 @@ func EffectiveWithExternal( userToolNames, externalNames, ) - hasUserToolTracking = userToolNames != nil allTools, userToolNames, externalNames = ApplyMandatoryToolFilter( ctx, allTools, From a8099e4224f7afaebb7a68621c8077fc54f811d9 Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 15:55:28 +0800 Subject: [PATCH 59/95] fix remaining user tool tracking ineffassign --- internal/flow/llmflow/llmflow.go | 8 ++++---- internal/toolsurface/toolsurface.go | 3 +-- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/internal/flow/llmflow/llmflow.go b/internal/flow/llmflow/llmflow.go index ac5bb92155..0561134a37 100644 --- a/internal/flow/llmflow/llmflow.go +++ b/internal/flow/llmflow/llmflow.go @@ -2032,7 +2032,7 @@ func (f *Flow) getFilteredTools( externalToolNames, invocation.RunOptions, ) - hasUserToolTracking = userToolNames != nil + filteredHasUserToolTracking := userToolNames != nil // If no filter is specified, return all tools for this invocation. if invocation.RunOptions.ToolFilter == nil { @@ -2041,7 +2041,7 @@ func (f *Flow) getFilteredTools( toolsnapshot.Set( invocation, allTools, - len(trackedUserToolNames(allTools, hasUserToolTracking, userToolNames)) > 0, + len(trackedUserToolNames(allTools, filteredHasUserToolTracking, userToolNames)) > 0, filteredTraceableToolNames(allTools, traceableUserToolNames), ) return allTools @@ -2054,7 +2054,7 @@ func (f *Flow) getFilteredTools( ctx, allTools, userToolNames, - hasUserToolTracking, + filteredHasUserToolTracking, invocation.RunOptions, ) @@ -2062,7 +2062,7 @@ func (f *Flow) getFilteredTools( toolsnapshot.Set( invocation, filtered, - len(trackedUserToolNames(filtered, hasUserToolTracking, userToolNames)) > 0, + len(trackedUserToolNames(filtered, filteredHasUserToolTracking, userToolNames)) > 0, filteredTraceableToolNames(filtered, traceableUserToolNames), ) diff --git a/internal/toolsurface/toolsurface.go b/internal/toolsurface/toolsurface.go index fdab7ebd64..10c59fe4fc 100644 --- a/internal/toolsurface/toolsurface.go +++ b/internal/toolsurface/toolsurface.go @@ -294,7 +294,6 @@ func EffectiveWithExternal( externalNames, invocation.RunOptions, ) - hasUserToolTracking = userToolNames != nil if invocation.RunOptions.ToolFilter == nil { return allTools, userToolNames, externalNames } @@ -302,7 +301,7 @@ func EffectiveWithExternal( ctx, allTools, userToolNames, - hasUserToolTracking, + userToolNames != nil, invocation.RunOptions, ), userToolNames, externalNames } From 4cd83295dae42074f50a9598a385676707e257fc Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 16:16:40 +0800 Subject: [PATCH 60/95] preserve runtime builder API compatibility --- platform/worker/builder.go | 8 ++++++++ platform/worker/governance_test.go | 8 ++++---- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/platform/worker/builder.go b/platform/worker/builder.go index 74b30bf135..daf23d4a2a 100644 --- a/platform/worker/builder.go +++ b/platform/worker/builder.go @@ -83,6 +83,14 @@ func WithToolPolicyProvider(provider ToolPolicyProvider) RuntimeBuilderOption { func NewRuntimeBuilder( router storagerouter.Router, factory AgentFactory, +) (*RuntimeBuilder, error) { + return NewRuntimeBuilderWithOptions(router, factory) +} + +// NewRuntimeBuilderWithOptions creates a runtime builder with options. +func NewRuntimeBuilderWithOptions( + router storagerouter.Router, + factory AgentFactory, opts ...RuntimeBuilderOption, ) (*RuntimeBuilder, error) { if isNilDependency(router) { diff --git a/platform/worker/governance_test.go b/platform/worker/governance_test.go index 8718cf36cc..71114d874b 100644 --- a/platform/worker/governance_test.go +++ b/platform/worker/governance_test.go @@ -57,7 +57,7 @@ func TestRuntimeBuilderAppliesToolGovernanceEndToEnd(t *testing.T) { return policy, nil }) var captured AgentDependencies - builder, err := NewRuntimeBuilder( + builder, err := NewRuntimeBuilderWithOptions( router, AgentFactoryFunc(func( _ context.Context, @@ -106,7 +106,7 @@ func TestRuntimeBuilderRejectsMissingToolPolicyProviderBeforeStorage(t *testing. tenant, app, binding := governanceRuntimeConfig() app.ToolPolicyID = "policy-a" router := &countingRouter{Router: storagerouter.NewInMemoryRouter()} - builder, err := NewRuntimeBuilder( + builder, err := NewRuntimeBuilderWithOptions( router, AgentFactoryFunc(func( context.Context, @@ -126,7 +126,7 @@ func TestRuntimeBuilderRejectsMismatchedToolPolicyBeforeStorage(t *testing.T) { tenant, app, binding := governanceRuntimeConfig() app.ToolPolicyID = "policy-a" router := &countingRouter{Router: storagerouter.NewInMemoryRouter()} - builder, err := NewRuntimeBuilder( + builder, err := NewRuntimeBuilderWithOptions( router, AgentFactoryFunc(func( context.Context, @@ -172,7 +172,7 @@ func TestRuntimeBuilderIsolatesResolvedAndCompiledToolPolicySlices( HighRiskTools: []string{"shell"}, DangerousToolAction: platform.DangerousToolActionDeny, } - builder, err := NewRuntimeBuilder( + builder, err := NewRuntimeBuilderWithOptions( router, AgentFactoryFunc(func( _ context.Context, From f1f4548948dc248fc753126ca302cebfaf5df392 Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 17:01:24 +0800 Subject: [PATCH 61/95] fix base lint issues in tool execution --- graph/state_graph.go | 166 +++++++++++++++++++--------- internal/toolsurface/toolsurface.go | 44 ++++---- 2 files changed, 135 insertions(+), 75 deletions(-) diff --git a/graph/state_graph.go b/graph/state_graph.go index c371d07b56..feccb2a1cc 100644 --- a/graph/state_graph.go +++ b/graph/state_graph.go @@ -4400,73 +4400,23 @@ func runToolWithEventContexts( retryPolicy *tool.RetryPolicy, toolCallIndex int, ) (context.Context, *agent.Invocation, context.Context, *agent.Invocation, any, []byte, error) { - ctx = context.WithValue(ctx, tool.ContextKeyToolCallID{}, toolCall.ID) - if invocation, ok := agent.InvocationFromContext(ctx); ok && jsonrepair.IsToolCallArgumentsJSONRepairEnabled(invocation) { - jsonrepair.RepairToolCallArgumentsInPlace(ctx, &toolCall) - } decl := t.Declaration() - startInvocation := invocationFromContextOrFallback(ctx, nil) - var mandatoryToolPermissionPolicy tool.PermissionPolicy - if startInvocation != nil { - mandatoryToolPermissionPolicy = - startInvocation.RunOptions.MandatoryToolPermissionPolicy - } - visibilityResult, err := checkMandatoryToolVisibility( - ctx, - startInvocation, - toolCall, - t, - decl, - ) - if err != nil { - return ctx, startInvocation, ctx, startInvocation, nil, toolCall.Function.Arguments, err - } - if visibilityResult != nil { - return ctx, startInvocation, ctx, startInvocation, visibilityResult, toolCall.Function.Arguments, nil - } - - ctx, toolCall, customResult, err := runBeforeToolPluginCallbacks( + prepared, customResult, err := prepareToolCall( ctx, toolCall, - decl, - state, - ) - startInvocation = invocationFromContextOrFallback(ctx, startInvocation) - if err != nil { - return ctx, startInvocation, ctx, startInvocation, customResult, toolCall.Function.Arguments, err - } - if customResult != nil { - return ctx, startInvocation, ctx, startInvocation, customResult, toolCall.Function.Arguments, nil - } - - ctx, toolCall, customResult, err = runBeforeToolCallbacks( - ctx, - toolCall, - decl, toolCallbacks, + t, state, ) - startInvocation = invocationFromContextOrFallback(ctx, startInvocation) + ctx = prepared.ctx + toolCall = prepared.toolCall + startInvocation := prepared.startInvocation if err != nil { return ctx, startInvocation, ctx, startInvocation, customResult, toolCall.Function.Arguments, err } if customResult != nil { return ctx, startInvocation, ctx, startInvocation, customResult, toolCall.Function.Arguments, nil } - permissionResult, err := checkToolPermission( - ctx, - mandatoryToolPermissionPolicy, - startInvocation, - toolCall, - t, - decl, - ) - if err != nil { - return ctx, startInvocation, ctx, startInvocation, nil, toolCall.Function.Arguments, err - } - if permissionResult != nil { - return ctx, startInvocation, ctx, startInvocation, permissionResult, toolCall.Function.Arguments, nil - } startCtx := ctx callableTool, err := ensureCallableTool(t, toolCall.Function.Name) @@ -4547,6 +4497,112 @@ func runToolWithEventContexts( return startCtx, startInvocation, ctx, completeInvocation, result, toolCall.Function.Arguments, nil } +type preparedToolCall struct { + ctx context.Context + toolCall model.ToolCall + startInvocation *agent.Invocation + mandatoryToolPermissionPolicy tool.PermissionPolicy +} + +func prepareToolCall( + ctx context.Context, + toolCall model.ToolCall, + toolCallbacks *tool.Callbacks, + t tool.Tool, + state State, +) (preparedToolCall, any, error) { + ctx = context.WithValue(ctx, tool.ContextKeyToolCallID{}, toolCall.ID) + if invocation, ok := agent.InvocationFromContext(ctx); ok && jsonrepair.IsToolCallArgumentsJSONRepairEnabled(invocation) { + jsonrepair.RepairToolCallArgumentsInPlace(ctx, &toolCall) + } + decl := t.Declaration() + startInvocation := invocationFromContextOrFallback(ctx, nil) + prepared := preparedToolCall{ + ctx: ctx, + toolCall: toolCall, + startInvocation: startInvocation, + mandatoryToolPermissionPolicy: mandatoryToolPermissionPolicy(startInvocation), + } + customResult, err := runPreToolChecks(&prepared, toolCallbacks, t, decl, state) + return prepared, customResult, err +} + +func runPreToolChecks( + prepared *preparedToolCall, + toolCallbacks *tool.Callbacks, + t tool.Tool, + decl *tool.Declaration, + state State, +) (any, error) { + visibilityResult, err := checkMandatoryToolVisibility( + prepared.ctx, + prepared.startInvocation, + prepared.toolCall, + t, + decl, + ) + if err != nil || visibilityResult != nil { + return visibilityResult, err + } + + customResult, err := runPreToolCallbacks(prepared, toolCallbacks, decl, state) + if err != nil || customResult != nil { + return customResult, err + } + + permissionResult, err := checkToolPermission( + prepared.ctx, + prepared.mandatoryToolPermissionPolicy, + prepared.startInvocation, + prepared.toolCall, + t, + decl, + ) + if err != nil || permissionResult != nil { + return permissionResult, err + } + return nil, nil +} + +func runPreToolCallbacks( + prepared *preparedToolCall, + toolCallbacks *tool.Callbacks, + decl *tool.Declaration, + state State, +) (any, error) { + ctx, toolCall, customResult, err := runBeforeToolPluginCallbacks( + prepared.ctx, + prepared.toolCall, + decl, + state, + ) + prepared.ctx = ctx + prepared.toolCall = toolCall + prepared.startInvocation = invocationFromContextOrFallback(ctx, prepared.startInvocation) + if err != nil || customResult != nil { + return customResult, err + } + + ctx, toolCall, customResult, err = runBeforeToolCallbacks( + prepared.ctx, + prepared.toolCall, + decl, + toolCallbacks, + state, + ) + prepared.ctx = ctx + prepared.toolCall = toolCall + prepared.startInvocation = invocationFromContextOrFallback(ctx, prepared.startInvocation) + return customResult, err +} + +func mandatoryToolPermissionPolicy(invocation *agent.Invocation) tool.PermissionPolicy { + if invocation == nil { + return nil + } + return invocation.RunOptions.MandatoryToolPermissionPolicy +} + func toolCallbackErrorResult( startCtx context.Context, startInvocation *agent.Invocation, diff --git a/internal/toolsurface/toolsurface.go b/internal/toolsurface/toolsurface.go index 10c59fe4fc..9afb419048 100644 --- a/internal/toolsurface/toolsurface.go +++ b/internal/toolsurface/toolsurface.go @@ -49,33 +49,37 @@ func ResolveBase( ctx context.Context, invocation *agent.Invocation, ) ([]tool.Tool, map[string]bool, bool) { - var allTools []tool.Tool - var userToolNames map[string]bool - hasUserToolTracking := false if provider, ok := invocation.Agent.(agent.InvocationToolSurfaceProvider); ok { - allTools, userToolNames = provider.InvocationToolSurface(ctx, invocation) - hasUserToolTracking = userToolNames != nil - } else if provider, ok := invocation.Agent.(ToolFilterProvider); ok { - allTools = provider.FilterTools(ctx) - } else { - allTools = invocation.Agent.Tools() + allTools, userToolNames := provider.InvocationToolSurface(ctx, invocation) + if userToolNames != nil { + return allTools, userToolNames, true + } + return withTrackedUserTools(invocation, allTools) + } + if provider, ok := invocation.Agent.(ToolFilterProvider); ok { + return withTrackedUserTools(invocation, provider.FilterTools(ctx)) } + return withTrackedUserTools(invocation, invocation.Agent.Tools()) +} +func withTrackedUserTools( + invocation *agent.Invocation, + allTools []tool.Tool, +) ([]tool.Tool, map[string]bool, bool) { // User tools are those explicitly registered via WithTools and // WithToolSets. Framework tools (Knowledge, SubAgents) are never filtered. - if !hasUserToolTracking { - if provider, ok := invocation.Agent.(UserToolsProvider); ok { - userTools := provider.UserTools() - hasUserToolTracking = true - userToolNames = make(map[string]bool, len(userTools)) - for _, t := range userTools { - if name := toolName(t); name != "" { - userToolNames[name] = true - } - } + provider, ok := invocation.Agent.(UserToolsProvider) + if !ok { + return allTools, nil, false + } + userTools := provider.UserTools() + userToolNames := make(map[string]bool, len(userTools)) + for _, t := range userTools { + if name := toolName(t); name != "" { + userToolNames[name] = true } } - return allTools, userToolNames, hasUserToolTracking + return allTools, userToolNames, true } // AppendRunOptionTools appends RunOptions.AdditionalTools and ExternalTools to From dc0aa47b84370968271d43b3953131af0c30c09e Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 21:01:18 +0800 Subject: [PATCH 62/95] fix toolsurface lint ineffassign --- internal/toolsurface/toolsurface.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/internal/toolsurface/toolsurface.go b/internal/toolsurface/toolsurface.go index 9afb419048..518c759400 100644 --- a/internal/toolsurface/toolsurface.go +++ b/internal/toolsurface/toolsurface.go @@ -276,7 +276,7 @@ func EffectiveWithExternal( ctx = context.Background() } allTools, userToolNames, hasUserToolTracking := ResolveBase(ctx, invocation) - allTools, userToolNames, hasUserToolTracking, externalNames := + allTools, userToolNames, _, externalNames := AppendRunOptionTools( allTools, userToolNames, From 133a545412322a6eb0b5f94be3caf92996d24511 Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 08:20:28 +0800 Subject: [PATCH 63/95] feat(gateway): add budget preflight gate --- platform/gateway/errors.go | 2 + platform/gateway/service.go | 179 +++++++++++++++++++++++++++++++ platform/gateway/service_test.go | 80 ++++++++++++++ 3 files changed, 261 insertions(+) diff --git a/platform/gateway/errors.go b/platform/gateway/errors.go index bedcf57cc3..5fbacdf10b 100644 --- a/platform/gateway/errors.go +++ b/platform/gateway/errors.go @@ -27,6 +27,8 @@ var ( ErrUnsupportedMessageType = errors.New("gateway only supports text messages") // ErrEmptyText indicates that a text message does not contain usable text. ErrEmptyText = errors.New("gateway text content is required") + // ErrBudgetExceeded indicates that a runtime budget gate rejected the request. + ErrBudgetExceeded = errors.New("gateway budget exceeded") // ErrRunnerResponseEmpty indicates that the runner completed without assistant text. ErrRunnerResponseEmpty = errors.New("gateway runner response is empty") ) diff --git a/platform/gateway/service.go b/platform/gateway/service.go index 112440b73e..78ed972c59 100644 --- a/platform/gateway/service.go +++ b/platform/gateway/service.go @@ -37,12 +37,48 @@ type Service struct { leaseStore SessionLeaseStore auditSink platform.AuditSink messageEventSink platform.MessageEventSink + budgetEstimator BudgetEstimator now func() time.Time } // Option configures a Service. type Option func(*Service) +// BudgetEstimateRequest contains safe request metadata for gateway budget checks. +type BudgetEstimateRequest struct { + Runtime Runtime + Message platform.InboundMessage + Text string + SessionID string + RequestID string + InternalUserID string +} + +// BudgetEstimator estimates maximum pre-run token and cost usage for one request. +type BudgetEstimator interface { + EstimateBudget( + ctx context.Context, + request BudgetEstimateRequest, + ) (platform.UsageEstimate, error) +} + +// BudgetEstimatorFunc adapts a function into a BudgetEstimator. +type BudgetEstimatorFunc func( + ctx context.Context, + request BudgetEstimateRequest, +) (platform.UsageEstimate, error) + +// EstimateBudget implements BudgetEstimator. +func (f BudgetEstimatorFunc) EstimateBudget( + ctx context.Context, + request BudgetEstimateRequest, +) (platform.UsageEstimate, error) { + if f == nil { + return platform.UsageEstimate{}, nil + } + return f(ctx, request) +} + // WithAuditSink sets the audit sink used by the service. func WithAuditSink(sink platform.AuditSink) Option { return func(s *Service) { @@ -73,6 +109,13 @@ func WithSessionLeaseStore(store SessionLeaseStore) Option { } } +// WithBudgetEstimator enables pre-run tenant budget checks. +func WithBudgetEstimator(estimator BudgetEstimator) Option { + return func(s *Service) { + s.budgetEstimator = estimator + } +} + // NewService creates a gateway service. func NewService( registry Registry, @@ -147,6 +190,21 @@ func (s *Service) HandleInbound( internalUserID := platform.InternalUserID(msg.TenantID, msg.Channel, msg.ExternalUserID) setInboundTraceAttributes(callbackSpan, msg, sessionID, requestID, internalUserID) setInboundTraceAttributes(routeSpan, msg, sessionID, requestID, internalUserID) + if err := s.checkBudget( + routeCtx, + ctx, + routeSpan, + runtime, + auditSink, + msg, + text, + sessionID, + requestID, + internalUserID, + start, + ); err != nil { + return Result{}, err + } key := platform.IdempotencyKey( msg.TenantID, msg.Channel, @@ -240,6 +298,111 @@ func (s *Service) lookupRuntime( return runtime, nil } +func (s *Service) checkBudget( + routeCtx context.Context, + auditCtx context.Context, + routeSpan oteltrace.Span, + runtime Runtime, + auditSink platform.AuditSink, + msg platform.InboundMessage, + text string, + sessionID string, + requestID string, + internalUserID string, + start time.Time, +) error { + if s.budgetEstimator == nil { + return nil + } + budgetCtx, budgetSpan := telemetrytrace.Tracer.Start(routeCtx, "gateway.budget") + defer budgetSpan.End() + setInboundTraceAttributes(budgetSpan, msg, sessionID, requestID, internalUserID) + quota, err := platform.ParseTenantQuota(runtime.Tenant) + if err != nil { + s.writeRejectAuditTo(auditCtx, auditSink, msg, start, err) + recordSpanError(routeSpan, err) + recordSpanError(budgetSpan, err) + return err + } + estimate, err := s.budgetEstimator.EstimateBudget( + budgetCtx, + BudgetEstimateRequest{ + Runtime: runtime, + Message: msg, + Text: text, + SessionID: sessionID, + RequestID: requestID, + InternalUserID: internalUserID, + }, + ) + if err != nil { + s.writeRejectAuditTo(auditCtx, auditSink, msg, start, err) + recordSpanError(routeSpan, err) + recordSpanError(budgetSpan, err) + return err + } + decision, err := quota.Check(estimate) + if err != nil { + s.writeRejectAuditTo(auditCtx, auditSink, msg, start, err) + recordSpanError(routeSpan, err) + recordSpanError(budgetSpan, err) + return err + } + budgetSpan.SetAttributes( + attribute.String("decision", "allow"), + attribute.Int("estimated_total_tokens", estimatedTotalTokens(estimate)), + ) + if decision.Allowed { + return nil + } + budgetSpan.SetAttributes(attribute.String("decision", "deny")) + s.writeBudgetDeniedAudit( + auditCtx, + auditSink, + runtime, + requestID, + decision, + estimate, + quota, + start, + ) + err = fmt.Errorf("%w: %s", ErrBudgetExceeded, decision.Reason) + recordSpanError(routeSpan, err) + recordSpanError(budgetSpan, err) + return err +} + +func (s *Service) writeBudgetDeniedAudit( + ctx context.Context, + auditSink platform.AuditSink, + runtime Runtime, + requestID string, + decision platform.BudgetDecision, + estimate platform.UsageEstimate, + quota platform.TenantQuota, + start time.Time, +) { + record, err := platform.NewBudgetDecisionAuditRecord(platform.BudgetDecisionAuditInput{ + TenantID: runtime.Tenant.TenantID, + AppID: runtime.App.AppID, + RequestID: requestID, + TraceID: requestID, + Decision: decision, + Estimate: estimate, + Quota: quota, + Outcome: platform.BudgetDecisionOutcomeDeny, + CreatedAt: start, + }) + if err != nil { + s.writeRejectAuditTo(ctx, auditSink, platform.InboundMessage{ + TenantID: runtime.Tenant.TenantID, + AppID: runtime.App.AppID, + }, start, err) + return + } + s.writeAuditTo(ctx, auditSink, record) +} + func validateRuntimeForMessage(runtime Runtime, msg platform.InboundMessage) error { if err := runtime.Validate(); err != nil { return err @@ -931,5 +1094,21 @@ var traceErrorTypes = []struct { {ErrBindingMentionRequired, "binding_mention_required"}, {ErrUnsupportedMessageType, "unsupported_message_type"}, {ErrEmptyText, "empty_text"}, + {ErrBudgetExceeded, "budget_exceeded"}, {ErrRunnerResponseEmpty, "runner_response_empty"}, } + +func estimatedTotalTokens(estimate platform.UsageEstimate) int { + if estimate.PromptTokens > maxInt()-estimate.CompletionTokens { + return estimate.TotalTokens + } + sum := estimate.PromptTokens + estimate.CompletionTokens + if sum > estimate.TotalTokens { + return sum + } + return estimate.TotalTokens +} + +func maxInt() int { + return int(^uint(0) >> 1) +} diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go index 985dd06453..8fc26a9e78 100644 --- a/platform/gateway/service_test.go +++ b/platform/gateway/service_test.go @@ -858,6 +858,86 @@ func TestServiceHandleInboundRejectsMissingRequiredMention(t *testing.T) { assert.Equal(t, ErrBindingMentionRequired.Error(), audit.Records()[0].DecisionReason) } +func TestServiceHandleInboundRejectsBudgetExceededBeforeIdempotency(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "unused"} + runtime := validRuntime("tenant-a", r) + runtime.Tenant.QuotaJSON = `{"max_total_tokens":10}` + require.NoError(t, registry.Register(runtime)) + audit := platform.NewInMemoryAuditSink() + store := platform.NewInMemoryIdempotencyStore() + var estimateRequest BudgetEstimateRequest + svc := NewService( + registry, + store, + NewInMemoryOutboundStore(), + WithAuditSink(audit), + WithBudgetEstimator(BudgetEstimatorFunc(func( + ctx context.Context, + request BudgetEstimateRequest, + ) (platform.UsageEstimate, error) { + estimateRequest = request + return platform.UsageEstimate{PromptTokens: 8, CompletionTokens: 5}, nil + })), + ) + msg := inbound("tenant-a", "msg-1", "user-1", "hello") + msg.TraceContext = map[string]string{"request_id": "req-budget"} + + _, err := svc.HandleInbound(ctx, msg) + + require.ErrorIs(t, err, ErrBudgetExceeded) + assert.Empty(t, r.calls) + _, ok, getErr := store.Get(ctx, platform.IdempotencyKey("tenant-a", "wecom", "acct", "msg-1")) + require.NoError(t, getErr) + assert.False(t, ok) + require.Len(t, audit.Records(), 1) + record := audit.Records()[0] + assert.Equal(t, "budget:tenant", record.ToolName) + assert.Equal(t, string(platform.BudgetDecisionOutcomeDeny), record.Decision) + assert.Equal(t, "total_tokens_exceeded", record.DecisionReason) + assert.Equal(t, "req-budget", record.RequestID) + assert.Equal(t, "req-budget", record.TraceID) + assert.Contains(t, record.TokenUsageJSON, "prompt_tokens:8") + assert.Contains(t, record.TokenUsageJSON, "completion_tokens:5") + assert.Contains(t, record.TokenUsageJSON, "total_tokens:13") + assert.Equal(t, runtime.Tenant.TenantID, estimateRequest.Runtime.Tenant.TenantID) + assert.Equal(t, "hello", estimateRequest.Text) + assert.Equal(t, "req-budget", estimateRequest.RequestID) + assert.NotEmpty(t, estimateRequest.SessionID) + assert.NotEmpty(t, estimateRequest.InternalUserID) +} + +func TestServiceHandleInboundAllowsWithinBudget(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "within budget"} + runtime := validRuntime("tenant-a", r) + runtime.Tenant.QuotaJSON = `{"max_total_tokens":20}` + require.NoError(t, registry.Register(runtime)) + audit := platform.NewInMemoryAuditSink() + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + WithAuditSink(audit), + WithBudgetEstimator(BudgetEstimatorFunc(func( + ctx context.Context, + request BudgetEstimateRequest, + ) (platform.UsageEstimate, error) { + return platform.UsageEstimate{PromptTokens: 8, CompletionTokens: 5}, nil + })), + ) + + result, err := svc.HandleInbound(ctx, inbound("tenant-a", "msg-1", "user-1", "hello")) + + require.NoError(t, err) + assert.Equal(t, "within budget", result.Outbound.Content) + require.Len(t, r.calls, 1) + require.Len(t, audit.Records(), 1) + assert.Equal(t, "completed", audit.Records()[0].Decision) +} + func TestServiceHandleInboundAllowsAuthorizedGroupMention(t *testing.T) { ctx := context.Background() registry := NewInMemoryRegistry() From 3b80b0e1184988d311eedbb1a2edf3634ffc222f Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 08:27:05 +0800 Subject: [PATCH 64/95] feat(gateway): record runner usage --- platform/gateway/service.go | 82 +++++++++++++++++++++++++++----- platform/gateway/service_test.go | 63 +++++++++++++++++++++++- 2 files changed, 130 insertions(+), 15 deletions(-) diff --git a/platform/gateway/service.go b/platform/gateway/service.go index 78ed972c59..fd7dc6f8a5 100644 --- a/platform/gateway/service.go +++ b/platform/gateway/service.go @@ -37,6 +37,7 @@ type Service struct { leaseStore SessionLeaseStore auditSink platform.AuditSink messageEventSink platform.MessageEventSink + usageSink platform.UsageSink budgetEstimator BudgetEstimator now func() time.Time } @@ -93,6 +94,13 @@ func WithMessageEventSink(sink platform.MessageEventSink) Option { } } +// WithUsageSink sets the usage sink used for post-run accounting records. +func WithUsageSink(sink platform.UsageSink) Option { + return func(s *Service) { + s.usageSink = sink + } +} + // WithNow sets the clock used by the service. func WithNow(now func() time.Time) Option { return func(s *Service) { @@ -261,6 +269,11 @@ type inboundRunInput struct { Start time.Time } +type runnerOutput struct { + Content string + Usage *model.Usage +} + func (s *Service) lookupRuntime( routeCtx context.Context, auditCtx context.Context, @@ -550,7 +563,7 @@ func (s *Service) runAndReply( msg platform.InboundMessage, input inboundRunInput, ) (Result, error) { - content, err := s.runGatewayRunner( + output, err := s.runGatewayRunner( routeCtx, auditCtx, runtime, @@ -561,15 +574,20 @@ func (s *Service) runAndReply( if err != nil { return Result{}, err } - return s.writeReply( + result, err := s.writeReply( routeCtx, auditCtx, runtime, auditSink, msg, input, - content, + output.Content, ) + if err != nil { + return Result{}, err + } + s.writeUsageRecord(auditCtx, runtime, msg, input, output.Usage) + return result, nil } func (s *Service) runGatewayRunner( @@ -579,7 +597,7 @@ func (s *Service) runGatewayRunner( auditSink platform.AuditSink, msg platform.InboundMessage, input inboundRunInput, -) (string, error) { +) (runnerOutput, error) { runnerCtx, runnerSpan := telemetrytrace.Tracer.Start(routeCtx, "runner.run") defer runnerSpan.End() runnerCtx = platform.ContextWithStorageFencingToken(runnerCtx, input.FencingToken) @@ -616,15 +634,41 @@ func (s *Service) runGatewayRunner( if err != nil { s.writeAuditTo(auditCtx, auditSink, auditFromMessage(msg, input.SessionID, input.InternalUserID, "runner_error", err.Error(), input.Start, err)) recordSpanError(runnerSpan, err) - return "", err + return runnerOutput{}, err } - content, err := collectAssistantText(auditCtx, ch) + output, err := collectAssistantOutput(auditCtx, ch) if err != nil { s.writeAuditTo(auditCtx, auditSink, auditFromMessage(msg, input.SessionID, input.InternalUserID, "runner_error", err.Error(), input.Start, err)) recordSpanError(runnerSpan, err) - return "", err + return runnerOutput{}, err + } + return output, nil +} + +func (s *Service) writeUsageRecord( + ctx context.Context, + runtime Runtime, + msg platform.InboundMessage, + input inboundRunInput, + usage *model.Usage, +) { + if isNilInterfaceValue(s.usageSink) || usage == nil { + return } - return content, nil + record := platform.UsageRecord{ + TenantID: runtime.Tenant.TenantID, + AppID: runtime.App.AppID, + UserIDHash: platform.UserIDHash(msg.TenantID, msg.Channel, msg.ExternalUserID), + SessionID: input.SessionID, + RequestID: input.RequestID, + ModelName: runtime.App.ModelProfileID, + PromptTokens: usage.PromptTokens, + CompletionTokens: usage.CompletionTokens, + CachedTokens: usage.PromptTokensDetails.CachedTokens, + TraceID: input.RequestID, + CreatedAt: s.now(), + } + _ = s.usageSink.WriteUsage(ctx, record) } func (s *Service) writeReply( @@ -823,13 +867,22 @@ func inboundText(msg platform.InboundMessage) (string, error) { } func collectAssistantText(ctx context.Context, ch <-chan *event.Event) (string, error) { + output, err := collectAssistantOutput(ctx, ch) + if err != nil { + return "", err + } + return output.Content, nil +} + +func collectAssistantOutput(ctx context.Context, ch <-chan *event.Event) (runnerOutput, error) { var parts []string var final string + var usage *model.Usage for { var evt *event.Event select { case <-ctx.Done(): - return "", ctx.Err() + return runnerOutput{}, ctx.Err() case next, ok := <-ch: if !ok { goto done @@ -839,8 +892,11 @@ func collectAssistantText(ctx context.Context, ch <-chan *event.Event) (string, if evt == nil || evt.Response == nil { continue } + if evt.Response.Usage != nil { + usage = evt.Response.Usage + } if evt.IsTerminalError() { - return "", evt.Response.Error + return runnerOutput{}, evt.Response.Error } if evt.IsRunnerCompletion() { break @@ -864,13 +920,13 @@ func collectAssistantText(ctx context.Context, ch <-chan *event.Event) (string, } done: if strings.TrimSpace(final) != "" { - return strings.TrimSpace(final), nil + return runnerOutput{Content: strings.TrimSpace(final), Usage: usage}, nil } content := strings.TrimSpace(strings.Join(parts, "")) if content == "" { - return "", ErrRunnerResponseEmpty + return runnerOutput{}, ErrRunnerResponseEmpty } - return content, nil + return runnerOutput{Content: content, Usage: usage}, nil } func requestIDFor(msg platform.InboundMessage) string { diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go index 8fc26a9e78..568304e21f 100644 --- a/platform/gateway/service_test.go +++ b/platform/gateway/service_test.go @@ -1060,6 +1060,56 @@ func TestServiceHandleInboundUsesRequestIDAndStreamsText(t *testing.T) { assert.Equal(t, result.Outbound.TraceID, events[1].TraceID) } +func TestServiceHandleInboundWritesUsageRecord(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{ + response: "usage reply", + usage: &model.Usage{ + PromptTokens: 11, + CompletionTokens: 7, + TotalTokens: 18, + PromptTokensDetails: model.PromptTokensDetails{ + CachedTokens: 3, + }, + }, + } + runtime := validRuntime("tenant-a", r) + runtime.App.ModelProfileID = "profile-gpt" + require.NoError(t, registry.Register(runtime)) + usageSink := platform.NewInMemoryUsageSink() + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + WithUsageSink(usageSink), + ) + msg := inbound("tenant-a", "msg-1", "external-user-raw", "hello") + msg.TraceContext = map[string]string{"request_id": "req-usage"} + + result, err := svc.HandleInbound(ctx, msg) + + require.NoError(t, err) + assert.Equal(t, "usage reply", result.Outbound.Content) + records := usageSink.Records() + require.Len(t, records, 1) + record := records[0] + assert.Equal(t, "tenant-a", record.TenantID) + assert.Equal(t, "app", record.AppID) + assert.Equal(t, platform.UserIDHash("tenant-a", "wecom", "external-user-raw"), record.UserIDHash) + assert.Equal(t, result.SessionID, record.SessionID) + assert.Equal(t, "req-usage", record.RequestID) + assert.Equal(t, "profile-gpt", record.ModelName) + assert.Equal(t, 11, record.PromptTokens) + assert.Equal(t, 7, record.CompletionTokens) + assert.Equal(t, 3, record.CachedTokens) + assert.Equal(t, "req-usage", record.TraceID) + assert.False(t, record.CreatedAt.IsZero()) + assert.Zero(t, record.ModelCost) + assert.Zero(t, record.ToolCost) + assert.Zero(t, record.TotalCost) +} + func TestServiceHandleInboundEmitsTraceSkeleton(t *testing.T) { recorder := useGatewaySpanRecorder(t) ctx := context.Background() @@ -1431,6 +1481,7 @@ type runnerCall struct { type recordingRunner struct { response string chunks []string + usage *model.Usage runErr error calls []runnerCall } @@ -1480,11 +1531,19 @@ func (r *recordingRunner) Run( defer close(out) if len(r.chunks) > 0 { for i, chunk := range r.chunks { - out <- chunkEvent(chunk, i != len(r.chunks)-1) + evt := chunkEvent(chunk, i != len(r.chunks)-1) + if i == len(r.chunks)-1 && r.usage != nil { + evt.Response.Usage = r.usage + } + out <- evt } return } - out <- responseEvent(r.response, true) + evt := responseEvent(r.response, true) + if r.usage != nil { + evt.Response.Usage = r.usage + } + out <- evt }() return out, nil } From 76a147595f85857dca46f209195b3ddbb86a43cd Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 08:37:16 +0800 Subject: [PATCH 65/95] feat(gateway): apply model cost policy --- platform/gateway/registry.go | 28 ++++++++-- platform/gateway/service.go | 17 +++++- platform/gateway/service_test.go | 16 +++++- platform/model_cost_policy.go | 89 ++++++++++++++++++++++++++++++ platform/model_cost_policy_test.go | 76 +++++++++++++++++++++++++ platform/types_test.go | 11 ++++ platform/validation.go | 3 + 7 files changed, 231 insertions(+), 9 deletions(-) create mode 100644 platform/model_cost_policy.go create mode 100644 platform/model_cost_policy_test.go diff --git a/platform/gateway/registry.go b/platform/gateway/registry.go index 1e3a4c6827..cc80a21ece 100644 --- a/platform/gateway/registry.go +++ b/platform/gateway/registry.go @@ -20,11 +20,12 @@ import ( // Runtime contains the platform configuration and runner for one active binding. type Runtime struct { - Tenant platform.Tenant - App platform.AgentApp - Binding platform.ChannelBinding - Runner runner.Runner - Audit platform.AuditSink + Tenant platform.Tenant + App platform.AgentApp + Binding platform.ChannelBinding + ModelProfile platform.ModelProfile + Runner runner.Runner + Audit platform.AuditSink // ToolFilter narrows user-visible tools for this runtime. ToolFilter tool.FilterFunc // ToolPermissionPolicy enforces tool-call authorization before execution. @@ -42,6 +43,9 @@ func (r Runtime) Validate() error { if err := r.Binding.Validate(); err != nil { return err } + if err := validateRuntimeModelProfile(r); err != nil { + return err + } if r.Runner == nil { return ErrRuntimeNotFound } @@ -66,6 +70,20 @@ func (r Runtime) Validate() error { return nil } +func validateRuntimeModelProfile(r Runtime) error { + if r.ModelProfile == (platform.ModelProfile{}) { + return nil + } + if err := r.ModelProfile.Validate(); err != nil { + return err + } + if r.ModelProfile.TenantID != r.Tenant.TenantID || + r.ModelProfile.ProfileID != r.App.ModelProfileID { + return ErrRuntimeMismatch + } + return nil +} + func (r Runtime) matchesInbound(msg platform.InboundMessage) bool { return r.Tenant.TenantID == msg.TenantID && r.App.TenantID == msg.TenantID && diff --git a/platform/gateway/service.go b/platform/gateway/service.go index fd7dc6f8a5..4b1a3ca031 100644 --- a/platform/gateway/service.go +++ b/platform/gateway/service.go @@ -655,22 +655,37 @@ func (s *Service) writeUsageRecord( if isNilInterfaceValue(s.usageSink) || usage == nil { return } + modelCost, err := platform.ModelUsageCostForProfile(runtime.ModelProfile, usage) + if err != nil { + return + } record := platform.UsageRecord{ TenantID: runtime.Tenant.TenantID, AppID: runtime.App.AppID, UserIDHash: platform.UserIDHash(msg.TenantID, msg.Channel, msg.ExternalUserID), SessionID: input.SessionID, RequestID: input.RequestID, - ModelName: runtime.App.ModelProfileID, + ModelName: usageModelName(runtime), PromptTokens: usage.PromptTokens, CompletionTokens: usage.CompletionTokens, CachedTokens: usage.PromptTokensDetails.CachedTokens, + ModelUnitPrice: modelCost.UnitPrice, + ModelCost: modelCost.Cost, + TotalCost: modelCost.Cost, TraceID: input.RequestID, CreatedAt: s.now(), } _ = s.usageSink.WriteUsage(ctx, record) } +func usageModelName(runtime Runtime) string { + modelName := strings.TrimSpace(runtime.ModelProfile.Model) + if modelName != "" { + return modelName + } + return runtime.App.ModelProfileID +} + func (s *Service) writeReply( routeCtx context.Context, auditCtx context.Context, diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go index 568304e21f..51296ac0c6 100644 --- a/platform/gateway/service_test.go +++ b/platform/gateway/service_test.go @@ -1076,6 +1076,15 @@ func TestServiceHandleInboundWritesUsageRecord(t *testing.T) { } runtime := validRuntime("tenant-a", r) runtime.App.ModelProfileID = "profile-gpt" + runtime.ModelProfile = platform.ModelProfile{ + TenantID: "tenant-a", + ProfileID: "profile-gpt", + Model: "gpt-test", + CostPolicyJSON: `{ + "input_token_price_per_token":0.000001, + "output_token_price_per_token":0.000002 + }`, + } require.NoError(t, registry.Register(runtime)) usageSink := platform.NewInMemoryUsageSink() svc := NewService( @@ -1099,15 +1108,16 @@ func TestServiceHandleInboundWritesUsageRecord(t *testing.T) { assert.Equal(t, platform.UserIDHash("tenant-a", "wecom", "external-user-raw"), record.UserIDHash) assert.Equal(t, result.SessionID, record.SessionID) assert.Equal(t, "req-usage", record.RequestID) - assert.Equal(t, "profile-gpt", record.ModelName) + assert.Equal(t, "gpt-test", record.ModelName) assert.Equal(t, 11, record.PromptTokens) assert.Equal(t, 7, record.CompletionTokens) assert.Equal(t, 3, record.CachedTokens) assert.Equal(t, "req-usage", record.TraceID) assert.False(t, record.CreatedAt.IsZero()) - assert.Zero(t, record.ModelCost) + assert.InDelta(t, 0.000025/18, record.ModelUnitPrice, 0.000000000001) + assert.InDelta(t, 0.000025, record.ModelCost, 0.000000000001) assert.Zero(t, record.ToolCost) - assert.Zero(t, record.TotalCost) + assert.InDelta(t, 0.000025, record.TotalCost, 0.000000000001) } func TestServiceHandleInboundEmitsTraceSkeleton(t *testing.T) { diff --git a/platform/model_cost_policy.go b/platform/model_cost_policy.go new file mode 100644 index 0000000000..718f6384e5 --- /dev/null +++ b/platform/model_cost_policy.go @@ -0,0 +1,89 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "encoding/json" + "fmt" + "strings" + + "trpc.group/trpc-go/trpc-agent-go/model" +) + +// ModelCostPolicy captures per-token model pricing from ModelProfile.CostPolicyJSON. +type ModelCostPolicy struct { + InputTokenPricePerToken float64 `json:"input_token_price_per_token,omitempty"` + OutputTokenPricePerToken float64 `json:"output_token_price_per_token,omitempty"` +} + +// ModelUsageCost is the calculated cost for one model usage payload. +type ModelUsageCost struct { + UnitPrice float64 + Cost float64 +} + +// ParseModelCostPolicy parses ModelProfile.CostPolicyJSON. Empty policy means zero-cost accounting. +func ParseModelCostPolicy(profile ModelProfile) (ModelCostPolicy, error) { + costPolicyJSON := strings.TrimSpace(profile.CostPolicyJSON) + if costPolicyJSON == "" { + return ModelCostPolicy{}, nil + } + var policy ModelCostPolicy + if err := json.Unmarshal([]byte(costPolicyJSON), &policy); err != nil { + return ModelCostPolicy{}, fmt.Errorf("parsing model cost_policy_json: %w", err) + } + if err := policy.Validate(); err != nil { + return ModelCostPolicy{}, err + } + return policy, nil +} + +// Validate checks model pricing assumptions are safe to use for accounting. +func (p ModelCostPolicy) Validate() error { + if !isFiniteNonNegative(p.InputTokenPricePerToken) { + return fmt.Errorf("input_token_price_per_token must be finite and non-negative") + } + if !isFiniteNonNegative(p.OutputTokenPricePerToken) { + return fmt.Errorf("output_token_price_per_token must be finite and non-negative") + } + return nil +} + +// Cost calculates model cost for one usage payload. +func (p ModelCostPolicy) Cost(usage *model.Usage) (ModelUsageCost, error) { + if err := p.Validate(); err != nil { + return ModelUsageCost{}, err + } + if usage == nil { + return ModelUsageCost{}, nil + } + if usage.PromptTokens < 0 || usage.CompletionTokens < 0 { + return ModelUsageCost{}, fmt.Errorf("model usage token values must be non-negative") + } + cost := (float64(usage.PromptTokens) * p.InputTokenPricePerToken) + + (float64(usage.CompletionTokens) * p.OutputTokenPricePerToken) + totalTokens := usage.PromptTokens + usage.CompletionTokens + unitPrice := 0.0 + if totalTokens > 0 { + unitPrice = cost / float64(totalTokens) + } + return ModelUsageCost{ + UnitPrice: unitPrice, + Cost: cost, + }, nil +} + +// ModelUsageCostForProfile calculates model usage cost from a model profile. +func ModelUsageCostForProfile(profile ModelProfile, usage *model.Usage) (ModelUsageCost, error) { + policy, err := ParseModelCostPolicy(profile) + if err != nil { + return ModelUsageCost{}, err + } + return policy.Cost(usage) +} diff --git a/platform/model_cost_policy_test.go b/platform/model_cost_policy_test.go new file mode 100644 index 0000000000..60006377c0 --- /dev/null +++ b/platform/model_cost_policy_test.go @@ -0,0 +1,76 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "math" + "strings" + "testing" + + "trpc.group/trpc-go/trpc-agent-go/model" +) + +func TestModelCostPolicyCalculatesUsageCost(t *testing.T) { + policy, err := ParseModelCostPolicy(ModelProfile{ + TenantID: "tenant", + ProfileID: "profile", + CostPolicyJSON: `{"input_token_price_per_token":0.000001,"output_token_price_per_token":0.000002}`, + }) + if err != nil { + t.Fatalf("ParseModelCostPolicy: %v", err) + } + + cost, err := policy.Cost(&model.Usage{ + PromptTokens: 100, + CompletionTokens: 50, + }) + if err != nil { + t.Fatalf("Cost: %v", err) + } + + assertFloat(t, "Cost", cost.Cost, 0.0002) + assertFloat(t, "UnitPrice", cost.UnitPrice, 0.0002/150) +} + +func TestParseModelCostPolicyAllowsEmptyPolicy(t *testing.T) { + policy, err := ParseModelCostPolicy(ModelProfile{}) + if err != nil { + t.Fatalf("ParseModelCostPolicy: %v", err) + } + + cost, err := policy.Cost(&model.Usage{PromptTokens: 10, CompletionTokens: 5}) + if err != nil { + t.Fatalf("Cost: %v", err) + } + + assertFloat(t, "Cost", cost.Cost, 0) + assertFloat(t, "UnitPrice", cost.UnitPrice, 0) +} + +func TestModelCostPolicyRejectsInvalidInputs(t *testing.T) { + _, err := ParseModelCostPolicy(ModelProfile{CostPolicyJSON: `{"input_token_price_per_token":`}) + if err == nil || !strings.Contains(err.Error(), "cost_policy_json") { + t.Fatalf("expected parse error, got %v", err) + } + + _, err = ParseModelCostPolicy(ModelProfile{CostPolicyJSON: `{"input_token_price_per_token":-0.01}`}) + if err == nil || !strings.Contains(err.Error(), "input_token_price_per_token") { + t.Fatalf("expected negative input price error, got %v", err) + } + + _, err = ModelCostPolicy{OutputTokenPricePerToken: math.Inf(1)}.Cost(&model.Usage{}) + if err == nil || !strings.Contains(err.Error(), "output_token_price_per_token") { + t.Fatalf("expected infinite output price error, got %v", err) + } + + _, err = ModelCostPolicy{}.Cost(&model.Usage{PromptTokens: -1}) + if err == nil || !strings.Contains(err.Error(), "token values") { + t.Fatalf("expected negative usage error, got %v", err) + } +} diff --git a/platform/types_test.go b/platform/types_test.go index 4125bb8c77..0f5b686d68 100644 --- a/platform/types_test.go +++ b/platform/types_test.go @@ -762,6 +762,17 @@ func TestModelProfileRejectsInlineSecrets(t *testing.T) { } } +func TestModelProfileRejectsInvalidCostPolicy(t *testing.T) { + profile := ModelProfile{ + TenantID: "tenant", + ProfileID: "model", + CostPolicyJSON: `{"output_token_price_per_token":-0.01}`, + } + if err := profile.Validate(); err == nil || !strings.Contains(err.Error(), "output_token_price_per_token") { + t.Fatalf("expected cost policy validation error, got %v", err) + } +} + func TestIdempotencyUpdateUnknownKeyFails(t *testing.T) { store := NewInMemoryIdempotencyStore() _, err := store.Complete(context.Background(), "missing", "result") diff --git a/platform/validation.go b/platform/validation.go index ed9610639a..5fb4fb7995 100644 --- a/platform/validation.go +++ b/platform/validation.go @@ -185,6 +185,9 @@ func (p ModelProfile) Validate() error { if err := validateSecretReference("api_key_ref", p.APIKeyRef); err != nil { return err } + if _, err := ParseModelCostPolicy(p); err != nil { + return err + } return nil } From d6cf2ab14df7f0ba36bbd75a1015b864d16d5c33 Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 08:51:17 +0800 Subject: [PATCH 66/95] feat(platform): add budget usage snapshots --- platform/budget.go | 175 ++++++++++++++++++++++++++++++++ platform/budget_test.go | 217 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 392 insertions(+) diff --git a/platform/budget.go b/platform/budget.go index a3c4b6fdd1..61e1e631c2 100644 --- a/platform/budget.go +++ b/platform/budget.go @@ -37,6 +37,25 @@ type BudgetDecision struct { Reason string } +// BudgetUsageSnapshot captures accumulated usage against one quota boundary. +type BudgetUsageSnapshot struct { + TenantID string + AppID string + PromptTokensUsed int + CompletionTokensUsed int + TotalTokensUsed int + CostUsed float64 + MaxPromptTokens int + MaxCompletionTokens int + MaxTotalTokens int + MaxCost float64 + PromptTokensRemaining int + CompletionTokensRemaining int + TotalTokensRemaining int + CostRemaining float64 + Decision BudgetDecision +} + // ParseTenantQuota parses Tenant.QuotaJSON. Empty quota means no budget limits. func ParseTenantQuota(tenant Tenant) (TenantQuota, error) { if err := tenant.Validate(); err != nil { @@ -65,6 +84,77 @@ func CheckTenantBudget(tenant Tenant, estimate UsageEstimate) (BudgetDecision, e return quota.Check(estimate) } +// CheckUsageSummaryBudget checks accumulated usage against the tenant quota. +func CheckUsageSummaryBudget(tenant Tenant, summary UsageSummary) (BudgetUsageSnapshot, error) { + quota, err := ParseTenantQuota(tenant) + if err != nil { + return BudgetUsageSnapshot{}, err + } + if strings.TrimSpace(summary.TenantID) != strings.TrimSpace(tenant.TenantID) { + return BudgetUsageSnapshot{}, fmt.Errorf("usage summary tenant_id mismatch") + } + estimate := UsageEstimate{ + PromptTokens: summary.PromptTokens, + CompletionTokens: summary.CompletionTokens, + TotalTokens: summary.TotalTokens, + Cost: summary.TotalCost, + } + decision, err := quota.Check(estimate) + if err != nil { + return BudgetUsageSnapshot{}, err + } + return NewBudgetUsageSnapshot(summary, quota, decision) +} + +// NewBudgetUsageSnapshot builds an accumulated usage snapshot for dashboards and budget counters. +func NewBudgetUsageSnapshot( + summary UsageSummary, + quota TenantQuota, + decision BudgetDecision, +) (BudgetUsageSnapshot, error) { + if err := quota.Validate(); err != nil { + return BudgetUsageSnapshot{}, err + } + estimate := UsageEstimate{ + PromptTokens: summary.PromptTokens, + CompletionTokens: summary.CompletionTokens, + TotalTokens: summary.TotalTokens, + Cost: summary.TotalCost, + } + expected, err := quota.Check(estimate) + if err != nil { + return BudgetUsageSnapshot{}, err + } + effectiveTotalTokens, err := estimate.effectiveTotalTokens() + if err != nil { + return BudgetUsageSnapshot{}, err + } + if decision.Allowed != expected.Allowed || decision.Reason != expected.Reason { + return BudgetUsageSnapshot{}, fmt.Errorf("budget decision does not match summary and quota") + } + snapshot := BudgetUsageSnapshot{ + TenantID: strings.TrimSpace(summary.TenantID), + AppID: strings.TrimSpace(summary.AppID), + PromptTokensUsed: summary.PromptTokens, + CompletionTokensUsed: summary.CompletionTokens, + TotalTokensUsed: effectiveTotalTokens, + CostUsed: summary.TotalCost, + MaxPromptTokens: quota.MaxPromptTokens, + MaxCompletionTokens: quota.MaxCompletionTokens, + MaxTotalTokens: quota.MaxTotalTokens, + MaxCost: quota.MaxCost, + PromptTokensRemaining: remainingInt(quota.MaxPromptTokens, summary.PromptTokens), + CompletionTokensRemaining: remainingInt(quota.MaxCompletionTokens, summary.CompletionTokens), + TotalTokensRemaining: remainingInt(quota.MaxTotalTokens, effectiveTotalTokens), + CostRemaining: remainingCost(quota.MaxCost, summary.TotalCost), + Decision: decision, + } + if err := snapshot.Validate(); err != nil { + return BudgetUsageSnapshot{}, err + } + return snapshot, nil +} + // Validate checks quota limits are non-negative. func (q TenantQuota) Validate() error { if q.MaxPromptTokens < 0 { @@ -133,3 +223,88 @@ func isFiniteNonNegative(value float64) bool { func maxInt() int { return int(^uint(0) >> 1) } + +func remainingInt(limit int, used int) int { + if limit <= 0 { + return 0 + } + if used >= limit { + return 0 + } + return limit - used +} + +func remainingCost(limit float64, used float64) float64 { + if limit <= 0 || used >= limit { + return 0 + } + return limit - used +} + +// Validate checks that the budget snapshot is internally consistent. +func (s BudgetUsageSnapshot) Validate() error { + if strings.TrimSpace(s.TenantID) == "" { + return ErrTenantIDRequired + } + if err := validateAuditRedactedText("app_id", s.AppID); err != nil { + return err + } + if s.PromptTokensUsed < 0 || + s.CompletionTokensUsed < 0 || + s.TotalTokensUsed < 0 { + return fmt.Errorf("budget usage values must be non-negative") + } + if !isFiniteNonNegative(s.CostUsed) { + return fmt.Errorf("budget cost used must be finite and non-negative") + } + if s.PromptTokensRemaining < 0 || + s.CompletionTokensRemaining < 0 || + s.TotalTokensRemaining < 0 { + return fmt.Errorf("budget remaining token values must be non-negative") + } + if !isFiniteNonNegative(s.CostRemaining) { + return fmt.Errorf("budget cost remaining must be finite and non-negative") + } + quota := TenantQuota{ + MaxPromptTokens: s.MaxPromptTokens, + MaxCompletionTokens: s.MaxCompletionTokens, + MaxTotalTokens: s.MaxTotalTokens, + MaxCost: s.MaxCost, + } + if err := quota.Validate(); err != nil { + return err + } + estimate := UsageEstimate{ + PromptTokens: s.PromptTokensUsed, + CompletionTokens: s.CompletionTokensUsed, + TotalTokens: s.TotalTokensUsed, + Cost: s.CostUsed, + } + effectiveTotalTokens, err := estimate.effectiveTotalTokens() + if err != nil { + return err + } + if s.TotalTokensUsed != effectiveTotalTokens { + return fmt.Errorf("total_tokens_used must match effective total tokens") + } + expected, err := quota.Check(estimate) + if err != nil { + return err + } + if s.Decision.Allowed != expected.Allowed || s.Decision.Reason != expected.Reason { + return fmt.Errorf("budget snapshot decision does not match quota") + } + if s.PromptTokensRemaining != remainingInt(s.MaxPromptTokens, s.PromptTokensUsed) { + return fmt.Errorf("prompt_tokens_remaining does not match quota") + } + if s.CompletionTokensRemaining != remainingInt(s.MaxCompletionTokens, s.CompletionTokensUsed) { + return fmt.Errorf("completion_tokens_remaining does not match quota") + } + if s.TotalTokensRemaining != remainingInt(s.MaxTotalTokens, s.TotalTokensUsed) { + return fmt.Errorf("total_tokens_remaining does not match quota") + } + if s.CostRemaining != remainingCost(s.MaxCost, s.CostUsed) { + return fmt.Errorf("cost_remaining does not match quota") + } + return nil +} diff --git a/platform/budget_test.go b/platform/budget_test.go index def373d0df..1bbb1491dc 100644 --- a/platform/budget_test.go +++ b/platform/budget_test.go @@ -10,6 +10,7 @@ package platform import ( "math" + "strings" "testing" ) @@ -44,6 +45,222 @@ func TestCheckTenantBudgetAllowsWithinQuota(t *testing.T) { } } +func TestCheckUsageSummaryBudgetBuildsAllowedSnapshot(t *testing.T) { + tenant := Tenant{ + TenantID: "tenant", + QuotaJSON: `{"max_prompt_tokens":100,"max_completion_tokens":50,"max_total_tokens":150,"max_cost":1.25}`, + } + summary := UsageSummary{ + TenantID: "tenant", + AppID: "app", + PromptTokens: 40, + CompletionTokens: 20, + TotalTokens: 60, + TotalCost: 0.75, + } + + snapshot, err := CheckUsageSummaryBudget(tenant, summary) + if err != nil { + t.Fatalf("CheckUsageSummaryBudget: %v", err) + } + + if !snapshot.Decision.Allowed || snapshot.Decision.Reason != "" { + t.Fatalf("expected allowed snapshot, got %+v", snapshot.Decision) + } + if snapshot.TenantID != "tenant" || snapshot.AppID != "app" { + t.Fatalf("unexpected snapshot scope: %+v", snapshot) + } + if snapshot.PromptTokensRemaining != 60 || + snapshot.CompletionTokensRemaining != 30 || + snapshot.TotalTokensRemaining != 90 { + t.Fatalf("unexpected token remaining values: %+v", snapshot) + } + assertFloat(t, "CostRemaining", snapshot.CostRemaining, 0.50) +} + +func TestCheckUsageSummaryBudgetBuildsDeniedSnapshot(t *testing.T) { + tenant := Tenant{ + TenantID: "tenant", + QuotaJSON: `{"max_total_tokens":100,"max_cost":1.00}`, + } + summary := UsageSummary{ + TenantID: "tenant", + PromptTokens: 80, + CompletionTokens: 30, + TotalTokens: 110, + TotalCost: 0.50, + } + + snapshot, err := CheckUsageSummaryBudget(tenant, summary) + if err != nil { + t.Fatalf("CheckUsageSummaryBudget: %v", err) + } + + if snapshot.Decision.Allowed || snapshot.Decision.Reason != "total_tokens_exceeded" { + t.Fatalf("expected total token denial, got %+v", snapshot.Decision) + } + if snapshot.TotalTokensRemaining != 0 { + t.Fatalf("expected no remaining total tokens, got %+v", snapshot) + } +} + +func TestCheckUsageSummaryBudgetUsesEffectiveTotalTokens(t *testing.T) { + tenant := Tenant{ + TenantID: "tenant", + QuotaJSON: `{"max_total_tokens":100}`, + } + summary := UsageSummary{ + TenantID: "tenant", + PromptTokens: 80, + CompletionTokens: 30, + TotalTokens: 1, + } + + snapshot, err := CheckUsageSummaryBudget(tenant, summary) + if err != nil { + t.Fatalf("CheckUsageSummaryBudget: %v", err) + } + + if snapshot.TotalTokensUsed != 110 || snapshot.TotalTokensRemaining != 0 { + t.Fatalf("expected effective total token usage, got %+v", snapshot) + } + if snapshot.Decision.Allowed || snapshot.Decision.Reason != "total_tokens_exceeded" { + t.Fatalf("expected effective total token denial, got %+v", snapshot.Decision) + } +} + +func TestBudgetUsageSnapshotAllowsUnlimitedQuotaRemaining(t *testing.T) { + snapshot, err := CheckUsageSummaryBudget( + Tenant{TenantID: "tenant"}, + UsageSummary{ + TenantID: "tenant", + PromptTokens: 10, + CompletionTokens: 5, + TotalTokens: 15, + TotalCost: 0.25, + }, + ) + if err != nil { + t.Fatalf("CheckUsageSummaryBudget: %v", err) + } + + if !snapshot.Decision.Allowed { + t.Fatalf("expected unlimited quota to allow usage, got %+v", snapshot.Decision) + } + if snapshot.PromptTokensRemaining != 0 || + snapshot.CompletionTokensRemaining != 0 || + snapshot.TotalTokensRemaining != 0 || + snapshot.CostRemaining != 0 { + t.Fatalf("expected unlimited remaining values to stay zero, got %+v", snapshot) + } +} + +func TestCheckUsageSummaryBudgetRejectsMismatchedTenant(t *testing.T) { + _, err := CheckUsageSummaryBudget( + Tenant{TenantID: "tenant-a"}, + UsageSummary{TenantID: "tenant-b"}, + ) + if err == nil || !strings.Contains(err.Error(), "tenant_id mismatch") { + t.Fatalf("expected tenant mismatch error, got %v", err) + } +} + +func TestBudgetUsageSnapshotRejectsIncorrectRemaining(t *testing.T) { + snapshot, err := CheckUsageSummaryBudget( + Tenant{ + TenantID: "tenant", + QuotaJSON: `{"max_prompt_tokens":100,"max_completion_tokens":50,"max_total_tokens":150,"max_cost":1.00}`, + }, + UsageSummary{ + TenantID: "tenant", + PromptTokens: 40, + CompletionTokens: 20, + TotalTokens: 60, + TotalCost: 0.25, + }, + ) + if err != nil { + t.Fatalf("CheckUsageSummaryBudget: %v", err) + } + + snapshot.TotalTokensRemaining = 123 + if err := snapshot.Validate(); err == nil || !strings.Contains(err.Error(), "total_tokens_remaining") { + t.Fatalf("expected remaining mismatch error, got %v", err) + } + + snapshot, err = CheckUsageSummaryBudget( + Tenant{TenantID: "tenant", QuotaJSON: `{"max_cost":1.00}`}, + UsageSummary{TenantID: "tenant", TotalCost: 0.25}, + ) + if err != nil { + t.Fatalf("CheckUsageSummaryBudget cost: %v", err) + } + snapshot.CostRemaining = 0.10 + if err := snapshot.Validate(); err == nil || !strings.Contains(err.Error(), "cost_remaining") { + t.Fatalf("expected cost remaining mismatch error, got %v", err) + } +} + +func TestBudgetUsageSnapshotRejectsNonCanonicalTotalTokens(t *testing.T) { + snapshot := BudgetUsageSnapshot{ + TenantID: "tenant", + PromptTokensUsed: 80, + CompletionTokensUsed: 30, + TotalTokensUsed: 1, + MaxTotalTokens: 100, + TotalTokensRemaining: 99, + Decision: BudgetDecision{Reason: "total_tokens_exceeded"}, + } + + if err := snapshot.Validate(); err == nil || !strings.Contains(err.Error(), "effective total tokens") { + t.Fatalf("expected canonical total token error, got %v", err) + } +} + +func TestCheckUsageSummaryBudgetRejectsInvalidSummaryValues(t *testing.T) { + _, err := CheckUsageSummaryBudget( + Tenant{TenantID: "tenant"}, + UsageSummary{TenantID: "tenant", PromptTokens: -1}, + ) + if err == nil || !strings.Contains(err.Error(), "usage estimate") { + t.Fatalf("expected negative usage error, got %v", err) + } + + _, err = CheckUsageSummaryBudget( + Tenant{TenantID: "tenant"}, + UsageSummary{TenantID: "tenant", TotalCost: math.Inf(1)}, + ) + if err == nil || !strings.Contains(err.Error(), "cost") { + t.Fatalf("expected non-finite cost error, got %v", err) + } +} + +func TestCheckUsageSummaryBudgetRejectsTokenOverflow(t *testing.T) { + max := int(^uint(0) >> 1) + _, err := CheckUsageSummaryBudget( + Tenant{TenantID: "tenant"}, + UsageSummary{ + TenantID: "tenant", + PromptTokens: max, + CompletionTokens: 1, + }, + ) + if err == nil || !strings.Contains(err.Error(), "overflow") { + t.Fatalf("expected overflow error, got %v", err) + } +} + +func TestBudgetUsageSnapshotRejectsInconsistentDecision(t *testing.T) { + _, err := NewBudgetUsageSnapshot( + UsageSummary{TenantID: "tenant", TotalTokens: 200}, + TenantQuota{MaxTotalTokens: 100}, + BudgetDecision{Allowed: true}, + ) + if err == nil || !strings.Contains(err.Error(), "does not match") { + t.Fatalf("expected decision mismatch error, got %v", err) + } +} + func TestCheckTenantBudgetDeniesExceededLimits(t *testing.T) { tests := []struct { name string From 86cabbb7f2ee7cb261b418c40032796dc193f963 Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 17:04:30 +0800 Subject: [PATCH 67/95] reduce budget snapshot validation complexity --- platform/budget.go | 42 ++++++++++++++++++++++++++++++++++++------ 1 file changed, 36 insertions(+), 6 deletions(-) diff --git a/platform/budget.go b/platform/budget.go index 61e1e631c2..2c58b0e778 100644 --- a/platform/budget.go +++ b/platform/budget.go @@ -249,6 +249,24 @@ func (s BudgetUsageSnapshot) Validate() error { if err := validateAuditRedactedText("app_id", s.AppID); err != nil { return err } + if err := s.validateUsageValues(); err != nil { + return err + } + quota := s.quota() + if err := quota.Validate(); err != nil { + return err + } + estimate, err := s.estimate() + if err != nil { + return err + } + if err := s.validateDecision(quota, estimate); err != nil { + return err + } + return s.validateRemaining() +} + +func (s BudgetUsageSnapshot) validateUsageValues() error { if s.PromptTokensUsed < 0 || s.CompletionTokensUsed < 0 || s.TotalTokensUsed < 0 { @@ -265,15 +283,19 @@ func (s BudgetUsageSnapshot) Validate() error { if !isFiniteNonNegative(s.CostRemaining) { return fmt.Errorf("budget cost remaining must be finite and non-negative") } - quota := TenantQuota{ + return nil +} + +func (s BudgetUsageSnapshot) quota() TenantQuota { + return TenantQuota{ MaxPromptTokens: s.MaxPromptTokens, MaxCompletionTokens: s.MaxCompletionTokens, MaxTotalTokens: s.MaxTotalTokens, MaxCost: s.MaxCost, } - if err := quota.Validate(); err != nil { - return err - } +} + +func (s BudgetUsageSnapshot) estimate() (UsageEstimate, error) { estimate := UsageEstimate{ PromptTokens: s.PromptTokensUsed, CompletionTokens: s.CompletionTokensUsed, @@ -282,11 +304,15 @@ func (s BudgetUsageSnapshot) Validate() error { } effectiveTotalTokens, err := estimate.effectiveTotalTokens() if err != nil { - return err + return UsageEstimate{}, err } if s.TotalTokensUsed != effectiveTotalTokens { - return fmt.Errorf("total_tokens_used must match effective total tokens") + return UsageEstimate{}, fmt.Errorf("total_tokens_used must match effective total tokens") } + return estimate, nil +} + +func (s BudgetUsageSnapshot) validateDecision(quota TenantQuota, estimate UsageEstimate) error { expected, err := quota.Check(estimate) if err != nil { return err @@ -294,6 +320,10 @@ func (s BudgetUsageSnapshot) Validate() error { if s.Decision.Allowed != expected.Allowed || s.Decision.Reason != expected.Reason { return fmt.Errorf("budget snapshot decision does not match quota") } + return nil +} + +func (s BudgetUsageSnapshot) validateRemaining() error { if s.PromptTokensRemaining != remainingInt(s.MaxPromptTokens, s.PromptTokensUsed) { return fmt.Errorf("prompt_tokens_remaining does not match quota") } From 4755f6491fee609d9e898c6b2f75b3acb3f3619c Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 08:58:48 +0800 Subject: [PATCH 68/95] feat(gateway): enforce text length channel limit --- platform/gateway/errors.go | 2 + platform/gateway/service.go | 19 ++++++++- platform/gateway/service_test.go | 69 ++++++++++++++++++++++++++++++++ 3 files changed, 89 insertions(+), 1 deletion(-) diff --git a/platform/gateway/errors.go b/platform/gateway/errors.go index 5fbacdf10b..f8347f2e01 100644 --- a/platform/gateway/errors.go +++ b/platform/gateway/errors.go @@ -27,6 +27,8 @@ var ( ErrUnsupportedMessageType = errors.New("gateway only supports text messages") // ErrEmptyText indicates that a text message does not contain usable text. ErrEmptyText = errors.New("gateway text content is required") + // ErrTextTooLong indicates that a text message exceeds the runtime binding limit. + ErrTextTooLong = errors.New("gateway text content exceeds channel limit") // ErrBudgetExceeded indicates that a runtime budget gate rejected the request. ErrBudgetExceeded = errors.New("gateway budget exceeded") // ErrRunnerResponseEmpty indicates that the runner completed without assistant text. diff --git a/platform/gateway/service.go b/platform/gateway/service.go index 4b1a3ca031..e798ba21a0 100644 --- a/platform/gateway/service.go +++ b/platform/gateway/service.go @@ -186,7 +186,7 @@ func (s *Service) HandleInbound( return Result{}, err } auditSink := s.auditSinkForRuntime(runtime) - text, err := s.validateInboundContent(ctx, routeSpan, msg, start, auditSink) + text, err := s.validateInboundContent(ctx, routeSpan, runtime, msg, start, auditSink) if err != nil { return Result{}, err } @@ -429,6 +429,7 @@ func validateRuntimeForMessage(runtime Runtime, msg platform.InboundMessage) err func (s *Service) validateInboundContent( ctx context.Context, routeSpan oteltrace.Span, + runtime Runtime, msg platform.InboundMessage, start time.Time, auditSink platform.AuditSink, @@ -439,6 +440,11 @@ func (s *Service) validateInboundContent( recordSpanError(routeSpan, err) return "", err } + if err := validateTextLimit(text, runtime.Binding.ChannelLimits); err != nil { + s.writeRejectAuditTo(ctx, auditSink, msg, start, err) + recordSpanError(routeSpan, err) + return "", err + } return text, nil } @@ -881,6 +887,16 @@ func inboundText(msg platform.InboundMessage) (string, error) { return text, nil } +func validateTextLimit(text string, limits platform.ChannelLimits) error { + if limits.MaxTextLength <= 0 { + return nil + } + if len([]rune(text)) > limits.MaxTextLength { + return ErrTextTooLong + } + return nil +} + func collectAssistantText(ctx context.Context, ch <-chan *event.Event) (string, error) { output, err := collectAssistantOutput(ctx, ch) if err != nil { @@ -1165,6 +1181,7 @@ var traceErrorTypes = []struct { {ErrBindingMentionRequired, "binding_mention_required"}, {ErrUnsupportedMessageType, "unsupported_message_type"}, {ErrEmptyText, "empty_text"}, + {ErrTextTooLong, "text_too_long"}, {ErrBudgetExceeded, "budget_exceeded"}, {ErrRunnerResponseEmpty, "runner_response_empty"}, } diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go index 51296ac0c6..05b6101f42 100644 --- a/platform/gateway/service_test.go +++ b/platform/gateway/service_test.go @@ -780,6 +780,75 @@ func TestServiceHandleInboundRejectsUnsupportedMessage(t *testing.T) { assert.NotEqual(t, "user-1", audit.Records()[0].UserID) } +func TestServiceHandleInboundRejectsTextOverChannelLimitBeforeIdempotency(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "unused"} + runtime := validRuntime("tenant-a", r) + runtime.Binding.ChannelLimits.MaxTextLength = 5 + require.NoError(t, registry.Register(runtime)) + idempotency := platform.NewInMemoryIdempotencyStore() + audit := platform.NewInMemoryAuditSink() + svc := NewService( + registry, + idempotency, + NewInMemoryOutboundStore(), + WithAuditSink(audit), + ) + + _, err := svc.HandleInbound(ctx, inbound("tenant-a", "msg-1", "user-1", "你好世界呀!")) + + require.ErrorIs(t, err, ErrTextTooLong) + assert.Empty(t, r.calls) + _, ok, getErr := idempotency.Get(ctx, platform.IdempotencyKey("tenant-a", "wecom", "acct", "msg-1")) + require.NoError(t, getErr) + assert.False(t, ok) + require.Len(t, audit.Records(), 1) + assert.Equal(t, "reject", audit.Records()[0].Decision) + assert.Equal(t, ErrTextTooLong.Error(), audit.Records()[0].DecisionReason) + assert.NotContains(t, audit.Records()[0].DecisionReason, "你好世界呀") +} + +func TestServiceHandleInboundAllowsTextAtChannelLimitBoundary(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "ok"} + runtime := validRuntime("tenant-a", r) + runtime.Binding.ChannelLimits.MaxTextLength = 5 + require.NoError(t, registry.Register(runtime)) + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + ) + + result, err := svc.HandleInbound(ctx, inbound("tenant-a", "msg-1", "user-1", "你好世界呀")) + + require.NoError(t, err) + assert.Equal(t, "ok", result.Outbound.Content) + assert.Len(t, r.calls, 1) +} + +func TestServiceHandleInboundAllowsTextWhenChannelLimitUnset(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "ok"} + runtime := validRuntime("tenant-a", r) + runtime.Binding.ChannelLimits.MaxTextLength = 0 + require.NoError(t, registry.Register(runtime)) + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + ) + + result, err := svc.HandleInbound(ctx, inbound("tenant-a", "msg-1", "user-1", strings.Repeat("x", 8192))) + + require.NoError(t, err) + assert.Equal(t, "ok", result.Outbound.Content) + assert.Len(t, r.calls, 1) +} + func TestServiceHandleInboundRejectsDisallowedUser(t *testing.T) { ctx := context.Background() registry := NewInMemoryRegistry() From 63d7f0fa38b6e2598af3d3cf7307faebc01c6cbd Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 09:06:31 +0800 Subject: [PATCH 69/95] feat(gateway): enforce file size channel limit --- platform/gateway/errors.go | 2 + platform/gateway/service.go | 33 ++++++++++ platform/gateway/service_test.go | 107 +++++++++++++++++++++++++++++++ 3 files changed, 142 insertions(+) diff --git a/platform/gateway/errors.go b/platform/gateway/errors.go index f8347f2e01..0b22c97c8d 100644 --- a/platform/gateway/errors.go +++ b/platform/gateway/errors.go @@ -29,6 +29,8 @@ var ( ErrEmptyText = errors.New("gateway text content is required") // ErrTextTooLong indicates that a text message exceeds the runtime binding limit. ErrTextTooLong = errors.New("gateway text content exceeds channel limit") + // ErrFileTooLarge indicates that an inbound file part exceeds the runtime binding limit. + ErrFileTooLarge = errors.New("gateway file content exceeds channel limit") // ErrBudgetExceeded indicates that a runtime budget gate rejected the request. ErrBudgetExceeded = errors.New("gateway budget exceeded") // ErrRunnerResponseEmpty indicates that the runner completed without assistant text. diff --git a/platform/gateway/service.go b/platform/gateway/service.go index e798ba21a0..f091eebf34 100644 --- a/platform/gateway/service.go +++ b/platform/gateway/service.go @@ -434,6 +434,11 @@ func (s *Service) validateInboundContent( start time.Time, auditSink platform.AuditSink, ) (string, error) { + if err := validateFileLimits(msg, runtime.Binding.ChannelLimits); err != nil { + s.writeRejectAuditTo(ctx, auditSink, msg, start, err) + recordSpanError(routeSpan, err) + return "", err + } text, err := inboundText(msg) if err != nil { s.writeRejectAuditTo(ctx, auditSink, msg, start, err) @@ -897,6 +902,33 @@ func validateTextLimit(text string, limits platform.ChannelLimits) error { return nil } +func validateFileLimits(msg platform.InboundMessage, limits platform.ChannelLimits) error { + if limits.FileMaxBytes <= 0 { + return nil + } + for _, part := range msg.ContentParts { + if !contentPartHasFile(part) { + continue + } + if part.SizeBytes > limits.FileMaxBytes { + return ErrFileTooLarge + } + } + return nil +} + +func contentPartHasFile(part platform.ContentPart) bool { + switch part.Type { + case platform.ContentPartTypeImage, + platform.ContentPartTypeFile, + platform.ContentPartTypeAudio, + platform.ContentPartTypeVideo: + return true + default: + return false + } +} + func collectAssistantText(ctx context.Context, ch <-chan *event.Event) (string, error) { output, err := collectAssistantOutput(ctx, ch) if err != nil { @@ -1182,6 +1214,7 @@ var traceErrorTypes = []struct { {ErrUnsupportedMessageType, "unsupported_message_type"}, {ErrEmptyText, "empty_text"}, {ErrTextTooLong, "text_too_long"}, + {ErrFileTooLarge, "file_too_large"}, {ErrBudgetExceeded, "budget_exceeded"}, {ErrRunnerResponseEmpty, "runner_response_empty"}, } diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go index 05b6101f42..c2f01945ae 100644 --- a/platform/gateway/service_test.go +++ b/platform/gateway/service_test.go @@ -849,6 +849,113 @@ func TestServiceHandleInboundAllowsTextWhenChannelLimitUnset(t *testing.T) { assert.Len(t, r.calls, 1) } +func TestServiceHandleInboundRejectsFileOverChannelLimitBeforeIdempotency(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "unused"} + runtime := validRuntime("tenant-a", r) + runtime.Binding.ChannelLimits.FileMaxBytes = 10 + require.NoError(t, registry.Register(runtime)) + idempotency := platform.NewInMemoryIdempotencyStore() + audit := platform.NewInMemoryAuditSink() + svc := NewService( + registry, + idempotency, + NewInMemoryOutboundStore(), + WithAuditSink(audit), + ) + msg := inbound("tenant-a", "msg-1", "user-1", "hello") + msg.MessageType = platform.MessageTypeFile + msg.ContentParts = []platform.ContentPart{ + { + Type: platform.ContentPartTypeFile, + FileRef: "artifact://file@1", + MIMEType: "application/pdf", + SizeBytes: 11, + }, + } + + _, err := svc.HandleInbound(ctx, msg) + + require.ErrorIs(t, err, ErrFileTooLarge) + assert.Empty(t, r.calls) + _, ok, getErr := idempotency.Get(ctx, platform.IdempotencyKey("tenant-a", "wecom", "acct", "msg-1")) + require.NoError(t, getErr) + assert.False(t, ok) + require.Len(t, audit.Records(), 1) + assert.Equal(t, "reject", audit.Records()[0].Decision) + assert.Equal(t, ErrFileTooLarge.Error(), audit.Records()[0].DecisionReason) + assert.NotContains(t, audit.Records()[0].DecisionReason, "artifact://file@1") +} + +func TestServiceHandleInboundRejectsUnsupportedFileAtChannelLimitBoundary(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "unused"} + runtime := validRuntime("tenant-a", r) + runtime.Binding.ChannelLimits.FileMaxBytes = 10 + require.NoError(t, registry.Register(runtime)) + audit := platform.NewInMemoryAuditSink() + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + WithAuditSink(audit), + ) + msg := inbound("tenant-a", "msg-1", "user-1", "hello") + msg.MessageType = platform.MessageTypeFile + msg.ContentParts = []platform.ContentPart{ + { + Type: platform.ContentPartTypeFile, + FileRef: "artifact://file@1", + MIMEType: "application/pdf", + SizeBytes: 10, + }, + } + + _, err := svc.HandleInbound(ctx, msg) + + require.ErrorIs(t, err, ErrUnsupportedMessageType) + assert.Empty(t, r.calls) + require.Len(t, audit.Records(), 1) + assert.Equal(t, "reject", audit.Records()[0].Decision) + assert.Equal(t, ErrUnsupportedMessageType.Error(), audit.Records()[0].DecisionReason) +} + +func TestServiceHandleInboundSkipsFileLimitWhenChannelLimitUnset(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "unused"} + runtime := validRuntime("tenant-a", r) + runtime.Binding.ChannelLimits.FileMaxBytes = 0 + require.NoError(t, registry.Register(runtime)) + audit := platform.NewInMemoryAuditSink() + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + WithAuditSink(audit), + ) + msg := inbound("tenant-a", "msg-1", "user-1", "hello") + msg.MessageType = platform.MessageTypeFile + msg.ContentParts = []platform.ContentPart{ + { + Type: platform.ContentPartTypeFile, + FileRef: "artifact://file@1", + MIMEType: "application/pdf", + SizeBytes: 1 << 30, + }, + } + + _, err := svc.HandleInbound(ctx, msg) + + require.ErrorIs(t, err, ErrUnsupportedMessageType) + assert.Empty(t, r.calls) + require.Len(t, audit.Records(), 1) + assert.Equal(t, "reject", audit.Records()[0].Decision) + assert.Equal(t, ErrUnsupportedMessageType.Error(), audit.Records()[0].DecisionReason) +} + func TestServiceHandleInboundRejectsDisallowedUser(t *testing.T) { ctx := context.Background() registry := NewInMemoryRegistry() From 53ddb3a10ce36a20a21a1fa6145d85db2a317de2 Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 09:19:02 +0800 Subject: [PATCH 70/95] feat(gateway): enforce channel rate limits --- platform/gateway/errors.go | 2 + platform/gateway/service.go | 137 ++++++++++++++++++++++ platform/gateway/service_test.go | 191 +++++++++++++++++++++++++++++++ 3 files changed, 330 insertions(+) diff --git a/platform/gateway/errors.go b/platform/gateway/errors.go index 0b22c97c8d..f46fbfc1d6 100644 --- a/platform/gateway/errors.go +++ b/platform/gateway/errors.go @@ -31,6 +31,8 @@ var ( ErrTextTooLong = errors.New("gateway text content exceeds channel limit") // ErrFileTooLarge indicates that an inbound file part exceeds the runtime binding limit. ErrFileTooLarge = errors.New("gateway file content exceeds channel limit") + // ErrRateLimited indicates that a channel binding rate gate rejected the request. + ErrRateLimited = errors.New("gateway channel rate limit exceeded") // ErrBudgetExceeded indicates that a runtime budget gate rejected the request. ErrBudgetExceeded = errors.New("gateway budget exceeded") // ErrRunnerResponseEmpty indicates that the runner completed without assistant text. diff --git a/platform/gateway/service.go b/platform/gateway/service.go index f091eebf34..c1afe0aaa9 100644 --- a/platform/gateway/service.go +++ b/platform/gateway/service.go @@ -14,6 +14,7 @@ import ( "fmt" "reflect" "strings" + "sync" "time" "go.opentelemetry.io/otel/attribute" @@ -39,6 +40,7 @@ type Service struct { messageEventSink platform.MessageEventSink usageSink platform.UsageSink budgetEstimator BudgetEstimator + rateLimiter RateLimiter now func() time.Time } @@ -80,6 +82,29 @@ func (f BudgetEstimatorFunc) EstimateBudget( return f(ctx, request) } +// RateLimitRequest contains safe request metadata for gateway rate checks. +type RateLimitRequest struct { + Key string + Limits platform.ChannelLimits + Now time.Time +} + +// RateLimiter decides whether a gateway request may proceed under channel limits. +type RateLimiter interface { + Allow(ctx context.Context, request RateLimitRequest) (bool, error) +} + +// RateLimiterFunc adapts a function into a RateLimiter. +type RateLimiterFunc func(ctx context.Context, request RateLimitRequest) (bool, error) + +// Allow implements RateLimiter. +func (f RateLimiterFunc) Allow(ctx context.Context, request RateLimitRequest) (bool, error) { + if f == nil { + return true, nil + } + return f(ctx, request) +} + // WithAuditSink sets the audit sink used by the service. func WithAuditSink(sink platform.AuditSink) Option { return func(s *Service) { @@ -124,6 +149,15 @@ func WithBudgetEstimator(estimator BudgetEstimator) Option { } } +// WithRateLimiter sets the limiter used for channel binding rate checks. +func WithRateLimiter(limiter RateLimiter) Option { + return func(s *Service) { + if limiter != nil { + s.rateLimiter = limiter + } + } +} + // NewService creates a gateway service. func NewService( registry Registry, @@ -136,6 +170,7 @@ func NewService( idempotencyStore: idempotencyStore, outboundStore: outboundStore, leaseStore: NewInMemorySessionLeaseStore(), + rateLimiter: NewInMemoryRateLimiter(), now: time.Now, } for _, opt := range opts { @@ -198,6 +233,9 @@ func (s *Service) HandleInbound( internalUserID := platform.InternalUserID(msg.TenantID, msg.Channel, msg.ExternalUserID) setInboundTraceAttributes(callbackSpan, msg, sessionID, requestID, internalUserID) setInboundTraceAttributes(routeSpan, msg, sessionID, requestID, internalUserID) + if err := s.checkRateLimit(ctx, routeSpan, runtime, auditSink, msg, start); err != nil { + return Result{}, err + } if err := s.checkBudget( routeCtx, ctx, @@ -311,6 +349,36 @@ func (s *Service) lookupRuntime( return runtime, nil } +func (s *Service) checkRateLimit( + ctx context.Context, + routeSpan oteltrace.Span, + runtime Runtime, + auditSink platform.AuditSink, + msg platform.InboundMessage, + start time.Time, +) error { + limits := runtime.Binding.ChannelLimits + if limits.RateLimitQPS <= 0 { + return nil + } + allowed, err := s.rateLimiter.Allow(ctx, RateLimitRequest{ + Key: rateLimitKey(runtime.Binding), + Limits: limits, + Now: start, + }) + if err != nil { + s.writeRejectAuditTo(ctx, auditSink, msg, start, err) + recordSpanError(routeSpan, err) + return err + } + if allowed { + return nil + } + s.writeRejectAuditTo(ctx, auditSink, msg, start, ErrRateLimited) + recordSpanError(routeSpan, ErrRateLimited) + return ErrRateLimited +} + func (s *Service) checkBudget( routeCtx context.Context, auditCtx context.Context, @@ -809,6 +877,9 @@ func (s *Service) validateService() error { if s.outboundStore == nil { return fmt.Errorf("gateway outbound store is required") } + if s.rateLimiter == nil { + return fmt.Errorf("gateway rate limiter is required") + } if s.leaseStore == nil { return fmt.Errorf("gateway session lease store is required") } @@ -1215,6 +1286,7 @@ var traceErrorTypes = []struct { {ErrEmptyText, "empty_text"}, {ErrTextTooLong, "text_too_long"}, {ErrFileTooLarge, "file_too_large"}, + {ErrRateLimited, "rate_limited"}, {ErrBudgetExceeded, "budget_exceeded"}, {ErrRunnerResponseEmpty, "runner_response_empty"}, } @@ -1233,3 +1305,68 @@ func estimatedTotalTokens(estimate platform.UsageEstimate) int { func maxInt() int { return int(^uint(0) >> 1) } + +type rateLimitBucket struct { + tokens float64 + at time.Time +} + +// InMemoryRateLimiter is a process-local token bucket limiter for gateway bindings. +type InMemoryRateLimiter struct { + mu sync.Mutex + buckets map[string]rateLimitBucket +} + +// NewInMemoryRateLimiter creates a process-local token bucket limiter. +func NewInMemoryRateLimiter() *InMemoryRateLimiter { + return &InMemoryRateLimiter{buckets: make(map[string]rateLimitBucket)} +} + +// Allow implements RateLimiter. +func (l *InMemoryRateLimiter) Allow(ctx context.Context, request RateLimitRequest) (bool, error) { + if err := ctx.Err(); err != nil { + return false, err + } + qps := request.Limits.RateLimitQPS + if qps <= 0 { + return true, nil + } + burst := request.Limits.Burst + if burst <= 0 { + burst = qps + } + now := request.Now + if now.IsZero() { + now = time.Now() + } + l.mu.Lock() + defer l.mu.Unlock() + bucket := l.buckets[request.Key] + if bucket.at.IsZero() { + bucket.tokens = float64(burst) + bucket.at = now + } else if elapsed := now.Sub(bucket.at); elapsed > 0 { + bucket.tokens += elapsed.Seconds() * float64(qps) + if bucket.tokens > float64(burst) { + bucket.tokens = float64(burst) + } + bucket.at = now + } + if bucket.tokens < 1 { + l.buckets[request.Key] = bucket + return false, nil + } + bucket.tokens-- + l.buckets[request.Key] = bucket + return true, nil +} + +func rateLimitKey(binding platform.ChannelBinding) string { + return strings.Join([]string{ + binding.TenantID, + binding.AppID, + binding.BindingID, + binding.Channel, + binding.AccountID, + }, "\x00") +} diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go index c2f01945ae..62c0f6f39b 100644 --- a/platform/gateway/service_test.go +++ b/platform/gateway/service_test.go @@ -956,6 +956,197 @@ func TestServiceHandleInboundSkipsFileLimitWhenChannelLimitUnset(t *testing.T) { assert.Equal(t, ErrUnsupportedMessageType.Error(), audit.Records()[0].DecisionReason) } +func TestServiceHandleInboundRejectsRateLimitedBeforeBudgetAndIdempotency(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "unused"} + runtime := validRuntime("tenant-a", r) + runtime.Binding.ChannelLimits.RateLimitQPS = 1 + runtime.Binding.ChannelLimits.Burst = 1 + require.NoError(t, registry.Register(runtime)) + idempotency := platform.NewInMemoryIdempotencyStore() + audit := platform.NewInMemoryAuditSink() + estimateCalls := 0 + svc := NewService( + registry, + idempotency, + NewInMemoryOutboundStore(), + WithAuditSink(audit), + WithBudgetEstimator(BudgetEstimatorFunc(func( + ctx context.Context, + request BudgetEstimateRequest, + ) (platform.UsageEstimate, error) { + estimateCalls++ + return platform.UsageEstimate{}, nil + })), + ) + first := inbound("tenant-a", "msg-1", "user-1", "hello") + second := inbound("tenant-a", "msg-2", "user-1", "again") + + firstResult, err := svc.HandleInbound(ctx, first) + require.NoError(t, err) + _, err = svc.HandleInbound(ctx, second) + + require.ErrorIs(t, err, ErrRateLimited) + assert.Equal(t, "unused", firstResult.Outbound.Content) + assert.Len(t, r.calls, 1) + assert.Equal(t, 1, estimateCalls) + _, ok, getErr := idempotency.Get(ctx, platform.IdempotencyKey("tenant-a", "wecom", "acct", "msg-2")) + require.NoError(t, getErr) + assert.False(t, ok) + require.Len(t, audit.Records(), 2) + assert.Equal(t, "completed", audit.Records()[0].Decision) + assert.Equal(t, "reject", audit.Records()[1].Decision) + assert.Equal(t, ErrRateLimited.Error(), audit.Records()[1].DecisionReason) +} + +func TestServiceHandleInboundRateLimitRefillsOverTime(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "ok"} + runtime := validRuntime("tenant-a", r) + runtime.Binding.ChannelLimits.RateLimitQPS = 2 + runtime.Binding.ChannelLimits.Burst = 2 + require.NoError(t, registry.Register(runtime)) + now := time.Unix(1000, 0) + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + WithNow(func() time.Time { return now }), + ) + + _, err := svc.HandleInbound(ctx, inbound("tenant-a", "msg-1", "user-1", "first")) + require.NoError(t, err) + _, err = svc.HandleInbound(ctx, inbound("tenant-a", "msg-2", "user-1", "second")) + require.NoError(t, err) + _, err = svc.HandleInbound(ctx, inbound("tenant-a", "msg-3", "user-1", "third")) + require.ErrorIs(t, err, ErrRateLimited) + + now = now.Add(500 * time.Millisecond) + _, err = svc.HandleInbound(ctx, inbound("tenant-a", "msg-4", "user-1", "fourth")) + + require.NoError(t, err) + assert.Len(t, r.calls, 3) +} + +func TestServiceHandleInboundSkipsRateLimitWhenUnset(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "ok"} + runtime := validRuntime("tenant-a", r) + runtime.Binding.ChannelLimits.RateLimitQPS = 0 + runtime.Binding.ChannelLimits.Burst = 1 + require.NoError(t, registry.Register(runtime)) + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + ) + + _, err := svc.HandleInbound(ctx, inbound("tenant-a", "msg-1", "user-1", "first")) + require.NoError(t, err) + _, err = svc.HandleInbound(ctx, inbound("tenant-a", "msg-2", "user-1", "second")) + + require.NoError(t, err) + assert.Len(t, r.calls, 2) +} + +func TestServiceHandleInboundRateLimitKeepsDefaultLimiterOnNilOption(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "ok"} + runtime := validRuntime("tenant-a", r) + runtime.Binding.ChannelLimits.RateLimitQPS = 1 + runtime.Binding.ChannelLimits.Burst = 1 + require.NoError(t, registry.Register(runtime)) + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + WithRateLimiter(nil), + ) + + _, err := svc.HandleInbound(ctx, inbound("tenant-a", "msg-1", "user-1", "first")) + require.NoError(t, err) + _, err = svc.HandleInbound(ctx, inbound("tenant-a", "msg-2", "user-1", "second")) + + require.ErrorIs(t, err, ErrRateLimited) + assert.Len(t, r.calls, 1) +} + +func TestServiceHandleInboundRateLimitUsesQPSAsDefaultBurst(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "ok"} + runtime := validRuntime("tenant-a", r) + runtime.Binding.ChannelLimits.RateLimitQPS = 2 + runtime.Binding.ChannelLimits.Burst = 0 + require.NoError(t, registry.Register(runtime)) + now := time.Unix(1000, 0) + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + WithNow(func() time.Time { return now }), + ) + + _, err := svc.HandleInbound(ctx, inbound("tenant-a", "msg-1", "user-1", "first")) + require.NoError(t, err) + _, err = svc.HandleInbound(ctx, inbound("tenant-a", "msg-2", "user-1", "second")) + require.NoError(t, err) + _, err = svc.HandleInbound(ctx, inbound("tenant-a", "msg-3", "user-1", "third")) + + require.ErrorIs(t, err, ErrRateLimited) + assert.Len(t, r.calls, 2) +} + +func TestServiceHandleInboundRateLimitIsolatesBindings(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "ok"} + firstRuntime := validRuntime("tenant-a", r) + firstRuntime.Binding.ChannelLimits.RateLimitQPS = 1 + firstRuntime.Binding.ChannelLimits.Burst = 1 + require.NoError(t, registry.Register(firstRuntime)) + secondRuntime := validRuntimeForBinding( + "tenant-a", + "app-alt", + "binding-alt", + "wecom", + "acct-alt", + r, + ) + secondRuntime.Binding.ChannelLimits.RateLimitQPS = 1 + secondRuntime.Binding.ChannelLimits.Burst = 1 + require.NoError(t, registry.Register(secondRuntime)) + now := time.Unix(1000, 0) + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + WithNow(func() time.Time { return now }), + ) + + _, err := svc.HandleInbound(ctx, inbound("tenant-a", "msg-1", "user-1", "first")) + require.NoError(t, err) + _, err = svc.HandleInbound(ctx, inboundForRuntime( + "tenant-a", + "app-alt", + "binding-alt", + "wecom", + "acct-alt", + "msg-2", + "user-1", + "second", + )) + require.NoError(t, err) + _, err = svc.HandleInbound(ctx, inbound("tenant-a", "msg-3", "user-1", "third")) + + require.ErrorIs(t, err, ErrRateLimited) + assert.Len(t, r.calls, 2) +} + func TestServiceHandleInboundRejectsDisallowedUser(t *testing.T) { ctx := context.Background() registry := NewInMemoryRegistry() From 06cf699e3159faf856783dcd3a8956ab8174151d Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 09:29:15 +0800 Subject: [PATCH 71/95] feat(gateway): enforce mime type channel limit --- platform/gateway/errors.go | 2 + platform/gateway/service.go | 28 +++++- platform/gateway/service_test.go | 141 +++++++++++++++++++++++++++++++ platform/types.go | 1 + 4 files changed, 170 insertions(+), 2 deletions(-) diff --git a/platform/gateway/errors.go b/platform/gateway/errors.go index f46fbfc1d6..8ea65eaac9 100644 --- a/platform/gateway/errors.go +++ b/platform/gateway/errors.go @@ -31,6 +31,8 @@ var ( ErrTextTooLong = errors.New("gateway text content exceeds channel limit") // ErrFileTooLarge indicates that an inbound file part exceeds the runtime binding limit. ErrFileTooLarge = errors.New("gateway file content exceeds channel limit") + // ErrMIMETypeNotAllowed indicates that an inbound file part has a disallowed MIME type. + ErrMIMETypeNotAllowed = errors.New("gateway mime type is not allowed") // ErrRateLimited indicates that a channel binding rate gate rejected the request. ErrRateLimited = errors.New("gateway channel rate limit exceeded") // ErrBudgetExceeded indicates that a runtime budget gate rejected the request. diff --git a/platform/gateway/service.go b/platform/gateway/service.go index c1afe0aaa9..4314361a6c 100644 --- a/platform/gateway/service.go +++ b/platform/gateway/service.go @@ -974,20 +974,43 @@ func validateTextLimit(text string, limits platform.ChannelLimits) error { } func validateFileLimits(msg platform.InboundMessage, limits platform.ChannelLimits) error { - if limits.FileMaxBytes <= 0 { + if limits.FileMaxBytes <= 0 && len(limits.AllowedMIMETypes) == 0 { return nil } for _, part := range msg.ContentParts { if !contentPartHasFile(part) { continue } - if part.SizeBytes > limits.FileMaxBytes { + if limits.FileMaxBytes > 0 && part.SizeBytes > limits.FileMaxBytes { return ErrFileTooLarge } + if !mimeTypeAllowed(part.MIMEType, limits.AllowedMIMETypes) { + return ErrMIMETypeNotAllowed + } } return nil } +func mimeTypeAllowed(mimeType string, allowed []string) bool { + if len(allowed) == 0 { + return true + } + mimeType = strings.ToLower(strings.TrimSpace(mimeType)) + if mimeType == "" { + return false + } + for _, candidate := range allowed { + candidate = strings.ToLower(strings.TrimSpace(candidate)) + if candidate == "" { + continue + } + if candidate == mimeType { + return true + } + } + return false +} + func contentPartHasFile(part platform.ContentPart) bool { switch part.Type { case platform.ContentPartTypeImage, @@ -1286,6 +1309,7 @@ var traceErrorTypes = []struct { {ErrEmptyText, "empty_text"}, {ErrTextTooLong, "text_too_long"}, {ErrFileTooLarge, "file_too_large"}, + {ErrMIMETypeNotAllowed, "mime_type_not_allowed"}, {ErrRateLimited, "rate_limited"}, {ErrBudgetExceeded, "budget_exceeded"}, {ErrRunnerResponseEmpty, "runner_response_empty"}, diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go index 62c0f6f39b..b3446b8f01 100644 --- a/platform/gateway/service_test.go +++ b/platform/gateway/service_test.go @@ -956,6 +956,147 @@ func TestServiceHandleInboundSkipsFileLimitWhenChannelLimitUnset(t *testing.T) { assert.Equal(t, ErrUnsupportedMessageType.Error(), audit.Records()[0].DecisionReason) } +func TestServiceHandleInboundRejectsDisallowedMIMEBeforeIdempotency(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "unused"} + runtime := validRuntime("tenant-a", r) + runtime.Binding.ChannelLimits.AllowedMIMETypes = []string{"image/png"} + require.NoError(t, registry.Register(runtime)) + idempotency := platform.NewInMemoryIdempotencyStore() + audit := platform.NewInMemoryAuditSink() + svc := NewService( + registry, + idempotency, + NewInMemoryOutboundStore(), + WithAuditSink(audit), + ) + msg := inbound("tenant-a", "msg-1", "user-1", "hello") + msg.MessageType = platform.MessageTypeFile + msg.ContentParts = []platform.ContentPart{ + { + Type: platform.ContentPartTypeFile, + FileRef: "artifact://file@1", + MIMEType: "application/pdf", + SizeBytes: 10, + }, + } + + _, err := svc.HandleInbound(ctx, msg) + + require.ErrorIs(t, err, ErrMIMETypeNotAllowed) + assert.Empty(t, r.calls) + _, ok, getErr := idempotency.Get(ctx, platform.IdempotencyKey("tenant-a", "wecom", "acct", "msg-1")) + require.NoError(t, getErr) + assert.False(t, ok) + require.Len(t, audit.Records(), 1) + assert.Equal(t, "reject", audit.Records()[0].Decision) + assert.Equal(t, ErrMIMETypeNotAllowed.Error(), audit.Records()[0].DecisionReason) + assert.NotContains(t, audit.Records()[0].DecisionReason, "application/pdf") +} + +func TestServiceHandleInboundRejectsMissingMIMEWhenAllowlistConfigured(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "unused"} + runtime := validRuntime("tenant-a", r) + runtime.Binding.ChannelLimits.AllowedMIMETypes = []string{" ", "image/png"} + require.NoError(t, registry.Register(runtime)) + audit := platform.NewInMemoryAuditSink() + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + WithAuditSink(audit), + ) + msg := inbound("tenant-a", "msg-1", "user-1", "hello") + msg.MessageType = platform.MessageTypeFile + msg.ContentParts = []platform.ContentPart{ + { + Type: platform.ContentPartTypeFile, + FileRef: "artifact://file@1", + MIMEType: "", + SizeBytes: 10, + }, + } + + _, err := svc.HandleInbound(ctx, msg) + + require.ErrorIs(t, err, ErrMIMETypeNotAllowed) + assert.Empty(t, r.calls) + require.Len(t, audit.Records(), 1) + assert.Equal(t, "reject", audit.Records()[0].Decision) + assert.Equal(t, ErrMIMETypeNotAllowed.Error(), audit.Records()[0].DecisionReason) +} + +func TestServiceHandleInboundAllowsMIMECaseInsensitiveBeforeUnsupportedFile(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "unused"} + runtime := validRuntime("tenant-a", r) + runtime.Binding.ChannelLimits.AllowedMIMETypes = []string{" Application/PDF "} + require.NoError(t, registry.Register(runtime)) + audit := platform.NewInMemoryAuditSink() + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + WithAuditSink(audit), + ) + msg := inbound("tenant-a", "msg-1", "user-1", "hello") + msg.MessageType = platform.MessageTypeFile + msg.ContentParts = []platform.ContentPart{ + { + Type: platform.ContentPartTypeFile, + FileRef: "artifact://file@1", + MIMEType: "application/pdf", + SizeBytes: 10, + }, + } + + _, err := svc.HandleInbound(ctx, msg) + + require.ErrorIs(t, err, ErrUnsupportedMessageType) + assert.Empty(t, r.calls) + require.Len(t, audit.Records(), 1) + assert.Equal(t, "reject", audit.Records()[0].Decision) + assert.Equal(t, ErrUnsupportedMessageType.Error(), audit.Records()[0].DecisionReason) +} + +func TestServiceHandleInboundSkipsMIMEFilterWhenAllowlistUnset(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "unused"} + runtime := validRuntime("tenant-a", r) + runtime.Binding.ChannelLimits.AllowedMIMETypes = nil + require.NoError(t, registry.Register(runtime)) + audit := platform.NewInMemoryAuditSink() + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + WithAuditSink(audit), + ) + msg := inbound("tenant-a", "msg-1", "user-1", "hello") + msg.MessageType = platform.MessageTypeFile + msg.ContentParts = []platform.ContentPart{ + { + Type: platform.ContentPartTypeFile, + FileRef: "artifact://file@1", + MIMEType: "application/x-custom", + SizeBytes: 10, + }, + } + + _, err := svc.HandleInbound(ctx, msg) + + require.ErrorIs(t, err, ErrUnsupportedMessageType) + assert.Empty(t, r.calls) + require.Len(t, audit.Records(), 1) + assert.Equal(t, "reject", audit.Records()[0].Decision) + assert.Equal(t, ErrUnsupportedMessageType.Error(), audit.Records()[0].DecisionReason) +} + func TestServiceHandleInboundRejectsRateLimitedBeforeBudgetAndIdempotency(t *testing.T) { ctx := context.Background() registry := NewInMemoryRegistry() diff --git a/platform/types.go b/platform/types.go index 8daa5d52d5..b232587346 100644 --- a/platform/types.go +++ b/platform/types.go @@ -293,6 +293,7 @@ type ChannelLimits struct { MaxTextLength int CallbackACKTimeout time.Duration FileMaxBytes int64 + AllowedMIMETypes []string RateLimitQPS int Burst int SupportsAsyncReply bool From b93061884843e2eb34dab3948bc524d832ba1be6 Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 09:45:46 +0800 Subject: [PATCH 72/95] feat(gateway): enforce user concurrency limits --- platform/gateway/service.go | 163 ++++++++++++++++++++++++++++++- platform/gateway/service_test.go | 101 +++++++++++++++++++ platform/types.go | 23 ++--- 3 files changed, 274 insertions(+), 13 deletions(-) diff --git a/platform/gateway/service.go b/platform/gateway/service.go index 4314361a6c..497428728d 100644 --- a/platform/gateway/service.go +++ b/platform/gateway/service.go @@ -36,6 +36,7 @@ type Service struct { idempotencyStore platform.IdempotencyStore outboundStore OutboundStore leaseStore SessionLeaseStore + userGate UserConcurrencyGate auditSink platform.AuditSink messageEventSink platform.MessageEventSink usageSink platform.UsageSink @@ -82,6 +83,22 @@ func (f BudgetEstimatorFunc) EstimateBudget( return f(ctx, request) } +// UserConcurrencyRequest identifies one user-scoped gateway execution slot. +type UserConcurrencyRequest struct { + Key string + Limit int +} + +// UserConcurrencyGate limits concurrent gateway executions for a user scope. +type UserConcurrencyGate interface { + Acquire(ctx context.Context, request UserConcurrencyRequest) (UserConcurrencyLease, bool, error) +} + +// UserConcurrencyLease releases one acquired user concurrency slot. +type UserConcurrencyLease interface { + Release(ctx context.Context) error +} + // RateLimitRequest contains safe request metadata for gateway rate checks. type RateLimitRequest struct { Key string @@ -138,7 +155,18 @@ func WithNow(now func() time.Time) Option { // WithSessionLeaseStore sets the lease store used to serialize same-session runs. func WithSessionLeaseStore(store SessionLeaseStore) Option { return func(s *Service) { - s.leaseStore = store + if store != nil { + s.leaseStore = store + } + } +} + +// WithUserConcurrencyGate sets the gate used for per-user concurrency checks. +func WithUserConcurrencyGate(gate UserConcurrencyGate) Option { + return func(s *Service) { + if gate != nil { + s.userGate = gate + } } } @@ -170,6 +198,7 @@ func NewService( idempotencyStore: idempotencyStore, outboundStore: outboundStore, leaseStore: NewInMemorySessionLeaseStore(), + userGate: NewInMemoryUserConcurrencyGate(), rateLimiter: NewInMemoryRateLimiter(), now: time.Now, } @@ -265,6 +294,7 @@ func (s *Service) HandleInbound( requestID, internalUserID, key, + runtime.Binding.ChannelLimits.MaxConcurrentPerUser, ) if err != nil { return Result{}, err @@ -273,6 +303,7 @@ func (s *Service) HandleInbound( return result, nil } defer s.releaseSessionLease(ctx, record.SessionLease) + defer s.releaseUserConcurrency(ctx, record.UserLease) return s.runAndReply( routeCtx, @@ -295,6 +326,7 @@ func (s *Service) HandleInbound( type inboundRunRecord struct { Record platform.IdempotencyRecord SessionLease SessionLease + UserLease UserConcurrencyLease } type inboundRunInput struct { @@ -529,6 +561,7 @@ func (s *Service) startInboundRun( requestID string, internalUserID string, key string, + userConcurrencyLimit int, ) (inboundRunRecord, bool, Result, error) { idempotencyCtx, idempotencySpan := telemetrytrace.Tracer.Start(routeCtx, "gateway.idempotency") defer idempotencySpan.End() @@ -542,6 +575,17 @@ func (s *Service) startInboundRun( result, err := s.duplicateResult(resultCtx, existing) return inboundRunRecord{}, true, result, err } + userLease, handled, result, err := s.acquireUserConcurrency( + routeCtx, + msg, + sessionID, + requestID, + internalUserID, + userConcurrencyLimit, + ) + if err != nil || handled { + return inboundRunRecord{}, handled, result, err + } return s.acquireSessionLeaseAndStart( routeCtx, resultCtx, @@ -552,6 +596,7 @@ func (s *Service) startInboundRun( requestID, internalUserID, key, + userLease, ) } @@ -565,6 +610,7 @@ func (s *Service) acquireSessionLeaseAndStart( requestID string, internalUserID string, key string, + userLease UserConcurrencyLease, ) (inboundRunRecord, bool, Result, error) { lease, handled, result, err := s.acquireSessionLease( routeCtx, @@ -574,6 +620,7 @@ func (s *Service) acquireSessionLeaseAndStart( internalUserID, ) if err != nil || handled { + s.releaseUserConcurrency(resultCtx, userLease) return inboundRunRecord{}, handled, result, err } record, started, err := s.idempotencyStore.Start(idempotencyCtx, platform.IdempotencyRecord{ @@ -587,15 +634,50 @@ func (s *Service) acquireSessionLeaseAndStart( }) if err != nil { s.releaseSessionLease(resultCtx, lease) + s.releaseUserConcurrency(resultCtx, userLease) recordSpanError(idempotencySpan, err) return inboundRunRecord{}, false, Result{}, err } if !started { s.releaseSessionLease(resultCtx, lease) + s.releaseUserConcurrency(resultCtx, userLease) result, err := s.duplicateResult(resultCtx, record) return inboundRunRecord{}, true, result, err } - return inboundRunRecord{Record: record, SessionLease: lease}, false, Result{}, nil + return inboundRunRecord{Record: record, SessionLease: lease, UserLease: userLease}, false, Result{}, nil +} + +func (s *Service) acquireUserConcurrency( + routeCtx context.Context, + msg platform.InboundMessage, + sessionID string, + requestID string, + internalUserID string, + userConcurrencyLimit int, +) (UserConcurrencyLease, bool, Result, error) { + if userConcurrencyLimit <= 0 { + return nil, false, Result{}, nil + } + gateCtx, gateSpan := telemetrytrace.Tracer.Start(routeCtx, "gateway.user_concurrency") + defer gateSpan.End() + setInboundTraceAttributes(gateSpan, msg, sessionID, requestID, internalUserID) + lease, acquired, err := s.userGate.Acquire(gateCtx, UserConcurrencyRequest{ + Key: userConcurrencyKey(msg), + Limit: userConcurrencyLimit, + }) + if err != nil { + recordSpanError(gateSpan, err) + return nil, false, Result{}, err + } + if acquired { + return lease, false, Result{}, nil + } + return nil, true, Result{ + RequestID: requestID, + SessionID: sessionID, + Status: platform.IdempotencyStatusProcessing, + Processing: true, + }, nil } func (s *Service) acquireSessionLease( @@ -629,6 +711,18 @@ func (s *Service) acquireSessionLease( } func (s *Service) releaseSessionLease(ctx context.Context, lease SessionLease) { + if lease == nil { + return + } + cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + defer cancel() + _ = lease.Release(cleanupCtx) +} + +func (s *Service) releaseUserConcurrency(ctx context.Context, lease UserConcurrencyLease) { + if lease == nil { + return + } cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) defer cancel() _ = lease.Release(cleanupCtx) @@ -877,6 +971,9 @@ func (s *Service) validateService() error { if s.outboundStore == nil { return fmt.Errorf("gateway outbound store is required") } + if s.userGate == nil { + return fmt.Errorf("gateway user concurrency gate is required") + } if s.rateLimiter == nil { return fmt.Errorf("gateway rate limiter is required") } @@ -1394,3 +1491,65 @@ func rateLimitKey(binding platform.ChannelBinding) string { binding.AccountID, }, "\x00") } + +type inMemoryUserConcurrencyLease struct { + gate *InMemoryUserConcurrencyGate + key string + once sync.Once +} + +// InMemoryUserConcurrencyGate is a process-local user concurrency gate. +type InMemoryUserConcurrencyGate struct { + mu sync.Mutex + active map[string]int +} + +// NewInMemoryUserConcurrencyGate creates a process-local user concurrency gate. +func NewInMemoryUserConcurrencyGate() *InMemoryUserConcurrencyGate { + return &InMemoryUserConcurrencyGate{active: make(map[string]int)} +} + +// Acquire implements UserConcurrencyGate. +func (g *InMemoryUserConcurrencyGate) Acquire( + ctx context.Context, + request UserConcurrencyRequest, +) (UserConcurrencyLease, bool, error) { + if err := ctx.Err(); err != nil { + return nil, false, err + } + if request.Limit <= 0 { + return nil, true, nil + } + g.mu.Lock() + defer g.mu.Unlock() + if g.active[request.Key] >= request.Limit { + return nil, false, nil + } + g.active[request.Key]++ + return &inMemoryUserConcurrencyLease{gate: g, key: request.Key}, true, nil +} + +func (l *inMemoryUserConcurrencyLease) Release(ctx context.Context) error { + l.once.Do(func() { + l.gate.mu.Lock() + defer l.gate.mu.Unlock() + count := l.gate.active[l.key] + if count <= 1 { + delete(l.gate.active, l.key) + return + } + l.gate.active[l.key] = count - 1 + }) + return ctx.Err() +} + +func userConcurrencyKey(msg platform.InboundMessage) string { + return strings.Join([]string{ + msg.TenantID, + msg.AppID, + msg.BindingID, + msg.Channel, + msg.ChannelAccountID, + platform.UserIDHash(msg.TenantID, msg.Channel, msg.ExternalUserID), + }, "\x00") +} diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go index b3446b8f01..01f6c0f9d9 100644 --- a/platform/gateway/service_test.go +++ b/platform/gateway/service_test.go @@ -675,6 +675,107 @@ func TestServiceHandleInboundAllowsDifferentSessions(t *testing.T) { assert.NotEqual(t, r.calls[0].sessionID, r.calls[1].sessionID) } +func TestServiceHandleInboundRejectsUserConcurrencyBeforeIdempotency(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + registry := NewInMemoryRegistry() + r := &hangingFirstRunner{ + started: make(chan struct{}), + response: "done", + } + runtime := validRuntime("tenant-a", r) + runtime.Binding.ChannelLimits.MaxConcurrentPerUser = 1 + require.NoError(t, registry.Register(runtime)) + idempotency := platform.NewInMemoryIdempotencyStore() + svc := NewService(registry, idempotency, NewInMemoryOutboundStore()) + first := inbound("tenant-a", "msg-1", "user-1", "first") + first.ConversationType = platform.ConversationTypeGroup + first.ExternalGroupID = "group-1" + second := inbound("tenant-a", "msg-2", "user-1", "second") + second.ConversationType = platform.ConversationTypeGroup + second.ExternalGroupID = "group-2" + errCh := make(chan error, 1) + go func() { + _, err := svc.HandleInbound(ctx, first) + errCh <- err + }() + <-r.started + + busy, err := svc.HandleInbound(context.Background(), second) + + require.NoError(t, err) + assert.False(t, busy.Duplicate) + assert.True(t, busy.Processing) + assert.Equal(t, platform.IdempotencyStatusProcessing, busy.Status) + assert.Len(t, r.calls, 1) + _, ok, getErr := idempotency.Get( + context.Background(), + platform.IdempotencyKey("tenant-a", "wecom", "acct", "msg-2"), + ) + require.NoError(t, getErr) + assert.False(t, ok) + cancel() + require.ErrorIs(t, <-errCh, context.Canceled) +} + +func TestServiceHandleInboundUserConcurrencyIsolatesUsers(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + registry := NewInMemoryRegistry() + r := &hangingFirstRunner{ + started: make(chan struct{}), + response: "done", + } + runtime := validRuntime("tenant-a", r) + runtime.Binding.ChannelLimits.MaxConcurrentPerUser = 1 + require.NoError(t, registry.Register(runtime)) + svc := NewService(registry, platform.NewInMemoryIdempotencyStore(), NewInMemoryOutboundStore()) + first := inbound("tenant-a", "msg-1", "user-1", "first") + first.ConversationType = platform.ConversationTypeGroup + first.ExternalGroupID = "group-1" + second := inbound("tenant-a", "msg-2", "user-2", "second") + second.ConversationType = platform.ConversationTypeGroup + second.ExternalGroupID = "group-2" + errCh := make(chan error, 1) + go func() { + _, err := svc.HandleInbound(ctx, first) + errCh <- err + }() + <-r.started + + result, err := svc.HandleInbound(context.Background(), second) + + require.NoError(t, err) + assert.False(t, result.Processing) + assert.Equal(t, "done", result.Outbound.Content) + assert.Len(t, r.calls, 2) + cancel() + require.ErrorIs(t, <-errCh, context.Canceled) +} + +func TestServiceHandleInboundUserConcurrencyReleasesAfterCompletion(t *testing.T) { + ctx := context.Background() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "done"} + runtime := validRuntime("tenant-a", r) + runtime.Binding.ChannelLimits.MaxConcurrentPerUser = 1 + require.NoError(t, registry.Register(runtime)) + svc := NewService(registry, platform.NewInMemoryIdempotencyStore(), NewInMemoryOutboundStore()) + first := inbound("tenant-a", "msg-1", "user-1", "first") + first.ConversationType = platform.ConversationTypeGroup + first.ExternalGroupID = "group-1" + second := inbound("tenant-a", "msg-2", "user-1", "second") + second.ConversationType = platform.ConversationTypeGroup + second.ExternalGroupID = "group-2" + + _, err := svc.HandleInbound(ctx, first) + require.NoError(t, err) + _, err = svc.HandleInbound(ctx, second) + + require.NoError(t, err) + assert.Len(t, r.calls, 2) +} + func TestServiceHandleInboundReleaseIgnoresCanceledRequestContext(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) registry := NewInMemoryRegistry() diff --git a/platform/types.go b/platform/types.go index b232587346..f1ceb2bb46 100644 --- a/platform/types.go +++ b/platform/types.go @@ -290,17 +290,18 @@ type ToolPolicy struct { // ChannelLimits stores configurable channel capability and limit values. type ChannelLimits struct { - MaxTextLength int - CallbackACKTimeout time.Duration - FileMaxBytes int64 - AllowedMIMETypes []string - RateLimitQPS int - Burst int - SupportsAsyncReply bool - SupportsEdit bool - SupportsCardUpdate bool - RetryMaxAttempts int - RetryBackoff string + MaxTextLength int + CallbackACKTimeout time.Duration + FileMaxBytes int64 + AllowedMIMETypes []string + RateLimitQPS int + Burst int + MaxConcurrentPerUser int + SupportsAsyncReply bool + SupportsEdit bool + SupportsCardUpdate bool + RetryMaxAttempts int + RetryBackoff string } // ChannelBinding maps one external IM account to one tenant app. From 38232d97a8e6e7afa62ccbecfd9211d1f2177e79 Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 10:12:37 +0800 Subject: [PATCH 73/95] feat(gateway): enrich reject audit context --- platform/gateway/service.go | 127 +++++++++++++++++++++++++------ platform/gateway/service_test.go | 40 ++++++++++ 2 files changed, 143 insertions(+), 24 deletions(-) diff --git a/platform/gateway/service.go b/platform/gateway/service.go index 497428728d..137551a293 100644 --- a/platform/gateway/service.go +++ b/platform/gateway/service.go @@ -241,28 +241,30 @@ func (s *Service) HandleInbound( return Result{}, err } requestID := requestIDFor(msg) - setInboundTraceAttributes(callbackSpan, msg, "", requestID, "") - routeCtx, routeSpan := telemetrytrace.Tracer.Start(ctx, "gateway.route") - defer routeSpan.End() - setInboundTraceAttributes(routeSpan, msg, "", requestID, "") - runtime, err := s.lookupRuntime(routeCtx, ctx, routeSpan, msg, start) + sessionID, err := platform.SessionIDForInbound(msg) if err != nil { + recordSpanError(callbackSpan, err) return Result{}, err } - auditSink := s.auditSinkForRuntime(runtime) - text, err := s.validateInboundContent(ctx, routeSpan, runtime, msg, start, auditSink) + internalUserID := auditInternalUserID(msg) + auditContext := rejectAuditContext{ + SessionID: sessionID, + InternalUserID: internalUserID, + } + setInboundTraceAttributes(callbackSpan, msg, sessionID, requestID, internalUserID) + routeCtx, routeSpan := telemetrytrace.Tracer.Start(ctx, "gateway.route") + defer routeSpan.End() + setInboundTraceAttributes(routeSpan, msg, sessionID, requestID, internalUserID) + runtime, err := s.lookupRuntime(routeCtx, ctx, routeSpan, msg, start, auditContext) if err != nil { return Result{}, err } - sessionID, err := platform.SessionIDForInbound(msg) + auditSink := s.auditSinkForRuntime(runtime) + text, err := s.validateInboundContent(ctx, routeSpan, runtime, msg, start, auditSink, auditContext) if err != nil { - recordSpanError(routeSpan, err) return Result{}, err } - internalUserID := platform.InternalUserID(msg.TenantID, msg.Channel, msg.ExternalUserID) - setInboundTraceAttributes(callbackSpan, msg, sessionID, requestID, internalUserID) - setInboundTraceAttributes(routeSpan, msg, sessionID, requestID, internalUserID) - if err := s.checkRateLimit(ctx, routeSpan, runtime, auditSink, msg, start); err != nil { + if err := s.checkRateLimit(ctx, routeSpan, runtime, auditSink, msg, start, auditContext); err != nil { return Result{}, err } if err := s.checkBudget( @@ -277,6 +279,7 @@ func (s *Service) HandleInbound( requestID, internalUserID, start, + auditContext, ); err != nil { return Result{}, err } @@ -350,6 +353,7 @@ func (s *Service) lookupRuntime( routeSpan oteltrace.Span, msg platform.InboundMessage, start time.Time, + auditContext rejectAuditContext, ) (Runtime, error) { runtime, ok, err := s.registry.Lookup(routeCtx, msg) if err != nil { @@ -358,22 +362,23 @@ func (s *Service) lookupRuntime( } if !ok { err := ErrRuntimeNotFound - s.writeRejectAudit(auditCtx, msg, start, err) + s.writeRejectAuditWithContext(auditCtx, msg, start, err, auditContext) recordSpanError(routeSpan, err) return Runtime{}, err } if err := validateRuntimeForMessage(runtime, msg); err != nil { - s.writeRejectAudit(auditCtx, msg, start, err) + s.writeRejectAuditWithContext(auditCtx, msg, start, err, auditContext) recordSpanError(routeSpan, err) return Runtime{}, err } if err := authorizeBinding(runtime.Binding, msg); err != nil { - s.writeRejectAuditTo( + s.writeRejectAuditWithContextTo( auditCtx, s.auditSinkForRuntime(runtime), msg, start, err, + auditContext, ) recordSpanError(routeSpan, err) return Runtime{}, err @@ -388,6 +393,7 @@ func (s *Service) checkRateLimit( auditSink platform.AuditSink, msg platform.InboundMessage, start time.Time, + auditContext rejectAuditContext, ) error { limits := runtime.Binding.ChannelLimits if limits.RateLimitQPS <= 0 { @@ -399,14 +405,14 @@ func (s *Service) checkRateLimit( Now: start, }) if err != nil { - s.writeRejectAuditTo(ctx, auditSink, msg, start, err) + s.writeRejectAuditWithContextTo(ctx, auditSink, msg, start, err, auditContext) recordSpanError(routeSpan, err) return err } if allowed { return nil } - s.writeRejectAuditTo(ctx, auditSink, msg, start, ErrRateLimited) + s.writeRejectAuditWithContextTo(ctx, auditSink, msg, start, ErrRateLimited, auditContext) recordSpanError(routeSpan, ErrRateLimited) return ErrRateLimited } @@ -423,6 +429,7 @@ func (s *Service) checkBudget( requestID string, internalUserID string, start time.Time, + auditContext rejectAuditContext, ) error { if s.budgetEstimator == nil { return nil @@ -432,7 +439,7 @@ func (s *Service) checkBudget( setInboundTraceAttributes(budgetSpan, msg, sessionID, requestID, internalUserID) quota, err := platform.ParseTenantQuota(runtime.Tenant) if err != nil { - s.writeRejectAuditTo(auditCtx, auditSink, msg, start, err) + s.writeRejectAuditWithContextTo(auditCtx, auditSink, msg, start, err, auditContext) recordSpanError(routeSpan, err) recordSpanError(budgetSpan, err) return err @@ -449,14 +456,14 @@ func (s *Service) checkBudget( }, ) if err != nil { - s.writeRejectAuditTo(auditCtx, auditSink, msg, start, err) + s.writeRejectAuditWithContextTo(auditCtx, auditSink, msg, start, err, auditContext) recordSpanError(routeSpan, err) recordSpanError(budgetSpan, err) return err } decision, err := quota.Check(estimate) if err != nil { - s.writeRejectAuditTo(auditCtx, auditSink, msg, start, err) + s.writeRejectAuditWithContextTo(auditCtx, auditSink, msg, start, err, auditContext) recordSpanError(routeSpan, err) recordSpanError(budgetSpan, err) return err @@ -477,6 +484,8 @@ func (s *Service) checkBudget( decision, estimate, quota, + msg, + auditContext, start, ) err = fmt.Errorf("%w: %s", ErrBudgetExceeded, decision.Reason) @@ -493,6 +502,8 @@ func (s *Service) writeBudgetDeniedAudit( decision platform.BudgetDecision, estimate platform.UsageEstimate, quota platform.TenantQuota, + msg platform.InboundMessage, + auditContext rejectAuditContext, start time.Time, ) { record, err := platform.NewBudgetDecisionAuditRecord(platform.BudgetDecisionAuditInput{ @@ -513,6 +524,17 @@ func (s *Service) writeBudgetDeniedAudit( }, start, err) return } + record.Channel = msg.Channel + record.BindingID = msg.BindingID + record.UserID = platform.UserIDHash(msg.TenantID, msg.Channel, msg.ExternalUserID) + record.InternalUserID = auditContext.InternalUserID + record.UserIDHash = platform.UserIDHash(msg.TenantID, msg.Channel, msg.ExternalUserID) + record.SessionID = auditContext.SessionID + record.MessageID = msg.PlatformMessageID + if err := record.Validate(); err != nil { + s.writeRejectAuditWithContextTo(ctx, auditSink, msg, start, err, auditContext) + return + } s.writeAuditTo(ctx, auditSink, record) } @@ -533,20 +555,21 @@ func (s *Service) validateInboundContent( msg platform.InboundMessage, start time.Time, auditSink platform.AuditSink, + auditContext rejectAuditContext, ) (string, error) { if err := validateFileLimits(msg, runtime.Binding.ChannelLimits); err != nil { - s.writeRejectAuditTo(ctx, auditSink, msg, start, err) + s.writeRejectAuditWithContextTo(ctx, auditSink, msg, start, err, auditContext) recordSpanError(routeSpan, err) return "", err } text, err := inboundText(msg) if err != nil { - s.writeRejectAuditTo(ctx, auditSink, msg, start, err) + s.writeRejectAuditWithContextTo(ctx, auditSink, msg, start, err, auditContext) recordSpanError(routeSpan, err) return "", err } if err := validateTextLimit(text, runtime.Binding.ChannelLimits); err != nil { - s.writeRejectAuditTo(ctx, auditSink, msg, start, err) + s.writeRejectAuditWithContextTo(ctx, auditSink, msg, start, err, auditContext) recordSpanError(routeSpan, err) return "", err } @@ -947,6 +970,27 @@ func (s *Service) writeRejectAudit( s.writeAudit(ctx, auditFromMessage(msg, "", "", "reject", err.Error(), start, err)) } +func (s *Service) writeRejectAuditWithContext( + ctx context.Context, + msg platform.InboundMessage, + start time.Time, + err error, + auditContext rejectAuditContext, +) { + s.writeAudit( + ctx, + auditFromMessage( + msg, + auditContext.SessionID, + auditContext.InternalUserID, + "reject", + err.Error(), + start, + err, + ), + ) +} + func (s *Service) writeRejectAuditTo( ctx context.Context, auditSink platform.AuditSink, @@ -961,6 +1005,34 @@ func (s *Service) writeRejectAuditTo( ) } +type rejectAuditContext struct { + SessionID string + InternalUserID string +} + +func (s *Service) writeRejectAuditWithContextTo( + ctx context.Context, + auditSink platform.AuditSink, + msg platform.InboundMessage, + start time.Time, + err error, + auditContext rejectAuditContext, +) { + s.writeAuditTo( + ctx, + auditSink, + auditFromMessage( + msg, + auditContext.SessionID, + auditContext.InternalUserID, + "reject", + err.Error(), + start, + err, + ), + ) +} + func (s *Service) validateService() error { if s.registry == nil { return fmt.Errorf("gateway registry is required") @@ -1228,6 +1300,13 @@ func auditFromMessage( return record } +func auditInternalUserID(msg platform.InboundMessage) string { + if strings.TrimSpace(msg.ExternalUserID) == "" { + return "" + } + return platform.InternalUserID(msg.TenantID, msg.Channel, msg.ExternalUserID) +} + func redactAuditReason(reason string) string { if reason == "" { return "" diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go index 01f6c0f9d9..a197046780 100644 --- a/platform/gateway/service_test.go +++ b/platform/gateway/service_test.go @@ -406,6 +406,32 @@ func TestServiceHandleInboundUsesFallbackAuditForInvalidRuntime(t *testing.T) { assert.Empty(t, runtimeAudit.Records()) require.Len(t, fallbackAudit.Records(), 1) assert.Equal(t, "reject", fallbackAudit.Records()[0].Decision) + assert.NotEmpty(t, fallbackAudit.Records()[0].SessionID) + assert.NotEmpty(t, fallbackAudit.Records()[0].InternalUserID) +} + +func TestServiceHandleInboundRuntimeNotFoundAuditsSessionWithoutSyntheticUser(t *testing.T) { + ctx := context.Background() + audit := platform.NewInMemoryAuditSink() + svc := NewService( + NewInMemoryRegistry(), + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + WithAuditSink(audit), + ) + msg := inbound("tenant-a", "msg-runtime-missing", "", "") + msg.MessageType = platform.MessageTypeEvent + msg.RawEventType = "member_joined" + msg.ConversationType = "" + msg.ContentParts = nil + + _, err := svc.HandleInbound(ctx, msg) + + require.ErrorIs(t, err, ErrRuntimeNotFound) + require.Len(t, audit.Records(), 1) + assert.Equal(t, "reject", audit.Records()[0].Decision) + assert.NotEmpty(t, audit.Records()[0].SessionID) + assert.Empty(t, audit.Records()[0].InternalUserID) } func TestServiceHandleInboundFallsBackFromTypedNilRuntimeAudit(t *testing.T) { @@ -462,6 +488,8 @@ func TestServiceHandleInboundUsesRuntimeAuditForBindingRejection(t *testing.T) { require.ErrorIs(t, err, ErrBindingAccessDenied) require.Len(t, runtimeAudit.Records(), 1) assert.Equal(t, "reject", runtimeAudit.Records()[0].Decision) + assert.NotEmpty(t, runtimeAudit.Records()[0].SessionID) + assert.NotEmpty(t, runtimeAudit.Records()[0].InternalUserID) assert.Empty(t, fallbackAudit.Records()) } @@ -879,6 +907,8 @@ func TestServiceHandleInboundRejectsUnsupportedMessage(t *testing.T) { assert.NotEmpty(t, audit.Records()[0].AuditID) assert.Equal(t, "reject", audit.Records()[0].Decision) assert.NotEqual(t, "user-1", audit.Records()[0].UserID) + assert.NotEmpty(t, audit.Records()[0].SessionID) + assert.NotEmpty(t, audit.Records()[0].InternalUserID) } func TestServiceHandleInboundRejectsTextOverChannelLimitBeforeIdempotency(t *testing.T) { @@ -907,6 +937,8 @@ func TestServiceHandleInboundRejectsTextOverChannelLimitBeforeIdempotency(t *tes require.Len(t, audit.Records(), 1) assert.Equal(t, "reject", audit.Records()[0].Decision) assert.Equal(t, ErrTextTooLong.Error(), audit.Records()[0].DecisionReason) + assert.NotEmpty(t, audit.Records()[0].SessionID) + assert.NotEmpty(t, audit.Records()[0].InternalUserID) assert.NotContains(t, audit.Records()[0].DecisionReason, "你好世界呀") } @@ -1240,6 +1272,8 @@ func TestServiceHandleInboundRejectsRateLimitedBeforeBudgetAndIdempotency(t *tes assert.Equal(t, "completed", audit.Records()[0].Decision) assert.Equal(t, "reject", audit.Records()[1].Decision) assert.Equal(t, ErrRateLimited.Error(), audit.Records()[1].DecisionReason) + assert.NotEmpty(t, audit.Records()[1].SessionID) + assert.NotEmpty(t, audit.Records()[1].InternalUserID) } func TestServiceHandleInboundRateLimitRefillsOverTime(t *testing.T) { @@ -1507,6 +1541,12 @@ func TestServiceHandleInboundRejectsBudgetExceededBeforeIdempotency(t *testing.T assert.Equal(t, "total_tokens_exceeded", record.DecisionReason) assert.Equal(t, "req-budget", record.RequestID) assert.Equal(t, "req-budget", record.TraceID) + assert.Equal(t, msg.Channel, record.Channel) + assert.Equal(t, msg.BindingID, record.BindingID) + assert.Equal(t, msg.PlatformMessageID, record.MessageID) + assert.NotEmpty(t, record.SessionID) + assert.NotEmpty(t, record.InternalUserID) + assert.Equal(t, platform.UserIDHash(msg.TenantID, msg.Channel, msg.ExternalUserID), record.UserIDHash) assert.Contains(t, record.TokenUsageJSON, "prompt_tokens:8") assert.Contains(t, record.TokenUsageJSON, "completion_tokens:5") assert.Contains(t, record.TokenUsageJSON, "total_tokens:13") From 5afad209c28578b6f4d2c5c9a0f0eb86b070f5df Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 10:34:27 +0800 Subject: [PATCH 74/95] feat(toolpolicy): enrich audit runtime context --- platform/gateway/service.go | 10 +++ platform/toolpolicy/policy.go | 64 ++++++++++++- platform/toolpolicy/policy_test.go | 138 +++++++++++++++++++++++++++++ platform/validation.go | 14 ++- 4 files changed, 222 insertions(+), 4 deletions(-) diff --git a/platform/gateway/service.go b/platform/gateway/service.go index 137551a293..38f1cf4b36 100644 --- a/platform/gateway/service.go +++ b/platform/gateway/service.go @@ -27,6 +27,7 @@ import ( "trpc.group/trpc-go/trpc-agent-go/model" "trpc.group/trpc-go/trpc-agent-go/platform" "trpc.group/trpc-go/trpc-agent-go/platform/channeladapter" + "trpc.group/trpc-go/trpc-agent-go/platform/toolpolicy" telemetrytrace "trpc.group/trpc-go/trpc-agent-go/telemetry/trace" ) @@ -797,6 +798,15 @@ func (s *Service) runGatewayRunner( runnerCtx, runnerSpan := telemetrytrace.Tracer.Start(routeCtx, "runner.run") defer runnerSpan.End() runnerCtx = platform.ContextWithStorageFencingToken(runnerCtx, input.FencingToken) + runnerCtx = toolpolicy.ContextWithAuditContext(runnerCtx, toolpolicy.AuditContext{ + Channel: msg.Channel, + BindingID: msg.BindingID, + SessionID: input.SessionID, + InternalUserID: input.InternalUserID, + UserIDHash: platform.UserIDHash(msg.TenantID, msg.Channel, msg.ExternalUserID), + RequestID: input.RequestID, + AgentName: runtime.App.AgentName, + }) setInboundTraceAttributes(runnerSpan, msg, input.SessionID, input.RequestID, input.InternalUserID) if input.FencingToken > 0 { runnerSpan.SetAttributes(attribute.Int64("storage.fencing_token", input.FencingToken)) diff --git a/platform/toolpolicy/policy.go b/platform/toolpolicy/policy.go index 58e334398e..7e8fc0bc7e 100644 --- a/platform/toolpolicy/policy.go +++ b/platform/toolpolicy/policy.go @@ -18,6 +18,8 @@ import ( "strings" "time" + oteltrace "go.opentelemetry.io/otel/trace" + "trpc.group/trpc-go/trpc-agent-go/agent" "trpc.group/trpc-go/trpc-agent-go/platform" "trpc.group/trpc-go/trpc-agent-go/plugin" "trpc.group/trpc-go/trpc-agent-go/plugin/guardrail/approval" @@ -56,6 +58,35 @@ type ApprovalSummary struct { CreatedAt time.Time } +type auditContextKey struct{} + +// AuditContext carries trusted platform identity for tool governance audit. +// Tool policy deliberately does not infer user identity from generic agent +// sessions because session.UserID is not guaranteed to be a platform-derived +// internal user id outside the gateway runtime. +type AuditContext struct { + Channel string + BindingID string + SessionID string + InternalUserID string + UserIDHash string + RequestID string + AgentName string +} + +// ContextWithAuditContext attaches trusted platform audit context. +func ContextWithAuditContext(ctx context.Context, auditCtx AuditContext) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, auditContextKey{}, auditCtx) +} + +func auditContextFromContext(ctx context.Context) (AuditContext, bool) { + auditCtx, ok := ctx.Value(auditContextKey{}).(AuditContext) + return auditCtx, ok +} + // Option configures Policy. type Option func(*Policy) @@ -599,7 +630,7 @@ func (p *Policy) writeAudit(ctx context.Context, summary ApprovalSummary) error return nil } detailRef := summary.DetailRef() - if err := p.audit.WriteAudit(ctx, platform.AuditRecord{ + record := platform.AuditRecord{ AuditID: platform.AuditID(summary.TenantID, summary.AppID, summary.ToolName, summary.ToolCallID, string(summary.Decision), detailRef), TenantID: summary.TenantID, AppID: summary.AppID, @@ -609,12 +640,41 @@ func (p *Policy) writeAudit(ctx context.Context, summary ApprovalSummary) error RedactedDetailRef: detailRef, RedactionVersion: summary.RedactionVersion, CreatedAt: summary.CreatedAt, - }); err != nil { + } + applyAuditContext(ctx, &record) + if err := record.Validate(); err != nil { + return fmt.Errorf("tool policy audit record: %w", err) + } + if err := p.audit.WriteAudit(ctx, record); err != nil { return fmt.Errorf("write tool policy audit: %w", err) } return nil } +func applyAuditContext(ctx context.Context, record *platform.AuditRecord) { + if record == nil { + return + } + inv, ok := agent.InvocationFromContext(ctx) + if ok && inv != nil { + record.RequestID = strings.TrimSpace(inv.RunOptions.RequestID) + } + if auditCtx, ok := auditContextFromContext(ctx); ok { + if requestID := strings.TrimSpace(auditCtx.RequestID); requestID != "" { + record.RequestID = requestID + } + record.Channel = strings.TrimSpace(auditCtx.Channel) + record.BindingID = strings.TrimSpace(auditCtx.BindingID) + record.SessionID = strings.TrimSpace(auditCtx.SessionID) + record.InternalUserID = strings.TrimSpace(auditCtx.InternalUserID) + record.UserIDHash = strings.TrimSpace(auditCtx.UserIDHash) + record.AgentName = strings.TrimSpace(auditCtx.AgentName) + } + if spanCtx := oteltrace.SpanContextFromContext(ctx); spanCtx.IsValid() { + record.TraceID = spanCtx.TraceID().String() + } +} + // DetailRef returns compact non-secret detail that can be stored in audit logs. func (s ApprovalSummary) DetailRef() string { parts := []string{ diff --git a/platform/toolpolicy/policy_test.go b/platform/toolpolicy/policy_test.go index 41762ac246..b854e76e2a 100644 --- a/platform/toolpolicy/policy_test.go +++ b/platform/toolpolicy/policy_test.go @@ -16,11 +16,14 @@ import ( "time" "github.com/stretchr/testify/require" + oteltrace "go.opentelemetry.io/otel/trace" + "trpc.group/trpc-go/trpc-agent-go/agent" "trpc.group/trpc-go/trpc-agent-go/platform" "trpc.group/trpc-go/trpc-agent-go/plugin" "trpc.group/trpc-go/trpc-agent-go/plugin/guardrail/approval" "trpc.group/trpc-go/trpc-agent-go/plugin/guardrail/approval/review" + "trpc.group/trpc-go/trpc-agent-go/session" "trpc.group/trpc-go/trpc-agent-go/tool" ) @@ -128,6 +131,141 @@ func TestPolicyAllowsHighRiskWithAuditAndRedactsArguments(t *testing.T) { } } +func TestPolicyAuditIncludesInvocationContext(t *testing.T) { + audit := platform.NewInMemoryAuditSink() + p := newPolicy( + t, + platform.ToolPolicy{ + TenantID: "tenant", + AppID: "app", + DangerousToolAction: platform.DangerousToolActionAllowWithAudit, + HighRiskTools: []string{"http_post"}, + }, + WithAuditSink(audit), + ) + inv := agent.NewInvocation( + agent.WithInvocationRunOptions(agent.RunOptions{RequestID: "request-1"}), + ) + inv.AgentName = "assistant" + traceID := oteltrace.TraceID{ + 0x01, 0x02, 0x03, 0x04, + 0x05, 0x06, 0x07, 0x08, + 0x09, 0x0a, 0x0b, 0x0c, + 0x0d, 0x0e, 0x0f, 0x10, + } + spanID := oteltrace.SpanID{0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08} + ctx := oteltrace.ContextWithSpanContext( + ContextWithAuditContext( + agent.NewInvocationContext(context.Background(), inv), + AuditContext{ + Channel: "wecom", + BindingID: "binding-1", + SessionID: "session-1", + InternalUserID: "usr_internal", + UserIDHash: platform.UserIDHash("tenant", "wecom", "external-user"), + RequestID: "request-carrier", + AgentName: "assistant-carrier", + }, + ), + oteltrace.NewSpanContext(oteltrace.SpanContextConfig{ + TraceID: traceID, + SpanID: spanID, + }), + ) + + _, err := p.CheckToolPermission( + ctx, + request("http_post", tool.ToolMetadata{}, []byte(`{"Authorization":"Bearer raw-token"}`)), + ) + require.NoError(t, err) + + records := audit.Records() + require.Len(t, records, 1) + record := records[0] + require.Equal(t, "request-carrier", record.RequestID) + require.Equal(t, "wecom", record.Channel) + require.Equal(t, "binding-1", record.BindingID) + require.Equal(t, "session-1", record.SessionID) + require.Equal(t, "usr_internal", record.InternalUserID) + require.Equal(t, "assistant-carrier", record.AgentName) + require.Equal(t, traceID.String(), record.TraceID) + require.Equal(t, platform.UserIDHash("tenant", "wecom", "external-user"), record.UserIDHash) + require.NotContains(t, record.RedactedDetailRef, "raw-token") +} + +func TestPolicyAuditDoesNotInferUserIdentityFromInvocationSession(t *testing.T) { + audit := platform.NewInMemoryAuditSink() + p := newPolicy( + t, + platform.ToolPolicy{ + TenantID: "tenant", + AppID: "app", + DangerousToolAction: platform.DangerousToolActionAllowWithAudit, + HighRiskTools: []string{"http_post"}, + }, + WithAuditSink(audit), + ) + inv := agent.NewInvocation( + agent.WithInvocationSession(session.NewSession("app", "external-user-raw", "session-raw")), + agent.WithInvocationRunOptions(agent.RunOptions{RequestID: "request-1"}), + ) + + _, err := p.CheckToolPermission( + agent.NewInvocationContext(context.Background(), inv), + request("http_post", tool.ToolMetadata{}, []byte(`{"Authorization":"Bearer raw-token"}`)), + ) + require.NoError(t, err) + + records := audit.Records() + require.Len(t, records, 1) + record := records[0] + require.Equal(t, "request-1", record.RequestID) + require.Empty(t, record.SessionID) + require.Empty(t, record.InternalUserID) + require.Empty(t, record.UserIDHash) + require.NotContains(t, record.RedactedDetailRef, "raw-token") +} + +func TestPolicyAuditRejectsUnsafeContextFields(t *testing.T) { + tests := map[string]AuditContext{ + "session": { + SessionID: "Authorization: Bearer raw-token", + }, + "internal user": { + InternalUserID: "sk-raw-secret", + }, + "user hash": { + UserIDHash: "Authorization: Bearer raw-token", + }, + "agent name": { + AgentName: "Authorization: Bearer raw-token", + }, + } + for name, auditCtx := range tests { + t.Run(name, func(t *testing.T) { + audit := platform.NewInMemoryAuditSink() + p := newPolicy( + t, + platform.ToolPolicy{ + TenantID: "tenant", + AppID: "app", + DangerousToolAction: platform.DangerousToolActionAllowWithAudit, + HighRiskTools: []string{"http_post"}, + }, + WithAuditSink(audit), + ) + + _, err := p.CheckToolPermission( + ContextWithAuditContext(context.Background(), auditCtx), + request("http_post", tool.ToolMetadata{}, []byte(`{"Authorization":"Bearer raw-token"}`)), + ) + require.Error(t, err) + require.Contains(t, err.Error(), "tool policy audit record") + require.Empty(t, audit.Records()) + }) + } +} + func TestPolicyBuildsApprovalSummaryWithoutRawArguments(t *testing.T) { now := time.Unix(200, 0) p := newPolicy( diff --git a/platform/validation.go b/platform/validation.go index 5fb4fb7995..2bce624687 100644 --- a/platform/validation.go +++ b/platform/validation.go @@ -377,8 +377,18 @@ func (r AuditRecord) Validate() error { return fmt.Errorf("cost must be greater than or equal to 0") } for field, value := range map[string]string{ - "request_id": r.RequestID, - "trace_id": r.TraceID, + "channel": r.Channel, + "binding_id": r.BindingID, + "user_id": r.UserID, + "internal_user_id": r.InternalUserID, + "user_id_hash": r.UserIDHash, + "session_id": r.SessionID, + "message_id": r.MessageID, + "request_id": r.RequestID, + "agent_name": r.AgentName, + "model_name": r.ModelName, + "tool_name": r.ToolName, + "trace_id": r.TraceID, } { if err := validateAuditRedactedText(field, value); err != nil { return err From 5e64885445a45c7be66273b155ac213790aad722 Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 10:39:45 +0800 Subject: [PATCH 75/95] feat(platform): expand audit query dimensions --- platform/audit_query.go | 46 +++++++++++++++-------- platform/audit_query_test.go | 72 +++++++++++++++++++++++++++++++++--- 2 files changed, 97 insertions(+), 21 deletions(-) diff --git a/platform/audit_query.go b/platform/audit_query.go index 0d91cb450b..def6894f64 100644 --- a/platform/audit_query.go +++ b/platform/audit_query.go @@ -16,21 +16,25 @@ import ( // AuditQueryFilter scopes audit retrieval to one tenant and optional safe dimensions. type AuditQueryFilter struct { - TenantID string - AppID string - AuditID string - Channel string - BindingID string - UserIDHash string - SessionID string - RequestID string - MessageID string - ToolName string - Decision string - TraceID string - CreatedFrom time.Time - CreatedTo time.Time - Limit int + TenantID string + AppID string + AuditID string + Channel string + BindingID string + UserIDHash string + SessionID string + RequestID string + MessageID string + AgentName string + ModelName string + ToolName string + Decision string + ErrorType string + TraceID string + RedactionVersion string + CreatedFrom time.Time + CreatedTo time.Time + Limit int } // QueryAudit returns audit records matching one tenant-scoped filter. @@ -76,9 +80,13 @@ func (f AuditQueryFilter) normalize() (AuditQueryFilter, error) { safeTextField{"session_id", f.SessionID}, safeTextField{"request_id", f.RequestID}, safeTextField{"message_id", f.MessageID}, + safeTextField{"agent_name", f.AgentName}, + safeTextField{"model_name", f.ModelName}, safeTextField{"tool_name", f.ToolName}, safeTextField{"decision", f.Decision}, + safeTextField{"error_type", f.ErrorType}, safeTextField{"trace_id", f.TraceID}, + safeTextField{"redaction_version", f.RedactionVersion}, ); err != nil { return AuditQueryFilter{}, err } @@ -96,9 +104,13 @@ func (f AuditQueryFilter) normalize() (AuditQueryFilter, error) { f.SessionID = strings.TrimSpace(f.SessionID) f.RequestID = strings.TrimSpace(f.RequestID) f.MessageID = strings.TrimSpace(f.MessageID) + f.AgentName = strings.TrimSpace(f.AgentName) + f.ModelName = strings.TrimSpace(f.ModelName) f.ToolName = strings.TrimSpace(f.ToolName) f.Decision = strings.TrimSpace(f.Decision) + f.ErrorType = strings.TrimSpace(f.ErrorType) f.TraceID = strings.TrimSpace(f.TraceID) + f.RedactionVersion = strings.TrimSpace(f.RedactionVersion) return f, nil } @@ -115,8 +127,12 @@ func (f AuditQueryFilter) matches(record AuditRecord) bool { !matchOptional(f.SessionID, record.SessionID) || !matchOptional(f.RequestID, record.RequestID) || !matchOptional(f.MessageID, record.MessageID) || + !matchOptional(f.AgentName, record.AgentName) || + !matchOptional(f.ModelName, record.ModelName) || !matchOptional(f.ToolName, record.ToolName) || !matchOptional(f.Decision, record.Decision) || + !matchOptional(f.ErrorType, record.ErrorType) || + !matchOptional(f.RedactionVersion, record.RedactionVersion) || !matchOptional(f.TraceID, record.TraceID) { return false } diff --git a/platform/audit_query_test.go b/platform/audit_query_test.go index d0e85b5bd3..98f0089605 100644 --- a/platform/audit_query_test.go +++ b/platform/audit_query_test.go @@ -75,6 +75,47 @@ func TestQueryAuditSupportsAuditIDAndLimit(t *testing.T) { } } +func TestQueryAuditFiltersRuntimeAndRedactionDimensions(t *testing.T) { + baseTime := time.Date(2026, 7, 8, 10, 0, 0, 0, time.UTC) + target := auditRecordForQuery("tenant", "audit-1", "app", "wecom", "binding", "session", "request", "message", "workspace_write", "deny", "trace", baseTime) + target.AgentName = "assistant" + target.ModelName = "gpt-test" + target.ErrorType = "permission_denied" + target.RedactionVersion = "platform-toolpolicy-v1" + otherAgent := target + otherAgent.AuditID = "audit-2" + otherAgent.AgentName = "other" + otherModel := target + otherModel.AuditID = "audit-3" + otherModel.ModelName = "gpt-other" + otherError := target + otherError.AuditID = "audit-4" + otherError.ErrorType = "runner_error" + otherRedaction := target + otherRedaction.AuditID = "audit-5" + otherRedaction.RedactionVersion = "platform-budget-v1" + + matches, err := QueryAudit([]AuditRecord{ + otherAgent, + otherModel, + otherError, + otherRedaction, + target, + }, AuditQueryFilter{ + TenantID: "tenant", + AgentName: "assistant", + ModelName: "gpt-test", + ErrorType: "permission_denied", + RedactionVersion: "platform-toolpolicy-v1", + }) + if err != nil { + t.Fatalf("query runtime dimensions: %v", err) + } + if len(matches) != 1 || matches[0].AuditID != "audit-1" { + t.Fatalf("expected only audit-1, got %+v", matches) + } +} + func TestAuditSinkQueryUsesSnapshotAndReturnsCopies(t *testing.T) { sink := NewInMemoryAuditSink() record := auditRecordForQuery("tenant", "audit-1", "app", "telegram", "binding", "session", "request", "message", "tool", "allow", "trace", time.Now()) @@ -107,12 +148,31 @@ func TestQueryAuditRequiresTenant(t *testing.T) { } func TestQueryAuditRejectsUnsafeFilterValues(t *testing.T) { - _, err := QueryAudit(nil, AuditQueryFilter{ - TenantID: "tenant", - ToolName: "workspace_exec Authorization: Bearer raw-token", - }) - if err == nil || !strings.Contains(err.Error(), "tool_name") { - t.Fatalf("expected unsafe tool filter error, got %v", err) + tests := map[string]AuditQueryFilter{ + "tool_name": { + ToolName: "workspace_exec Authorization: Bearer raw-token", + }, + "agent_name": { + AgentName: "assistant Authorization: Bearer raw-token", + }, + "model_name": { + ModelName: "gpt-test password=plain", + }, + "error_type": { + ErrorType: "runner_error sk-1234567890abcdef", + }, + "redaction_version": { + RedactionVersion: "platform-v1 Authorization: Bearer raw-token", + }, + } + for name, filter := range tests { + t.Run(name, func(t *testing.T) { + filter.TenantID = "tenant" + _, err := QueryAudit(nil, filter) + if err == nil || !strings.Contains(err.Error(), name) { + t.Fatalf("expected unsafe %s filter error, got %v", name, err) + } + }) } } From 04f34db05f20897772b60bc4ae48a98f0daafb31 Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 17:07:55 +0800 Subject: [PATCH 76/95] reduce audit query filter complexity --- platform/audit_query.go | 51 ++++++++++++++++++++++++++--------------- 1 file changed, 33 insertions(+), 18 deletions(-) diff --git a/platform/audit_query.go b/platform/audit_query.go index def6894f64..6d4be5358d 100644 --- a/platform/audit_query.go +++ b/platform/audit_query.go @@ -119,27 +119,42 @@ func (f AuditQueryFilter) matchesScope(record AuditRecord) bool { } func (f AuditQueryFilter) matches(record AuditRecord) bool { - if !matchOptional(f.AppID, record.AppID) || - !matchOptional(f.AuditID, record.AuditID) || - !matchOptional(f.Channel, record.Channel) || - !matchOptional(f.BindingID, record.BindingID) || - !matchOptional(f.UserIDHash, record.UserIDHash) || - !matchOptional(f.SessionID, record.SessionID) || - !matchOptional(f.RequestID, record.RequestID) || - !matchOptional(f.MessageID, record.MessageID) || - !matchOptional(f.AgentName, record.AgentName) || - !matchOptional(f.ModelName, record.ModelName) || - !matchOptional(f.ToolName, record.ToolName) || - !matchOptional(f.Decision, record.Decision) || - !matchOptional(f.ErrorType, record.ErrorType) || - !matchOptional(f.RedactionVersion, record.RedactionVersion) || - !matchOptional(f.TraceID, record.TraceID) { - return false + return f.matchesOptionalFields(record) && f.matchesCreatedAt(record.CreatedAt) +} + +func (f AuditQueryFilter) matchesOptionalFields(record AuditRecord) bool { + for _, field := range []struct { + want string + got string + }{ + {f.AppID, record.AppID}, + {f.AuditID, record.AuditID}, + {f.Channel, record.Channel}, + {f.BindingID, record.BindingID}, + {f.UserIDHash, record.UserIDHash}, + {f.SessionID, record.SessionID}, + {f.RequestID, record.RequestID}, + {f.MessageID, record.MessageID}, + {f.AgentName, record.AgentName}, + {f.ModelName, record.ModelName}, + {f.ToolName, record.ToolName}, + {f.Decision, record.Decision}, + {f.ErrorType, record.ErrorType}, + {f.RedactionVersion, record.RedactionVersion}, + {f.TraceID, record.TraceID}, + } { + if !matchOptional(field.want, field.got) { + return false + } } - if !f.CreatedFrom.IsZero() && record.CreatedAt.Before(f.CreatedFrom) { + return true +} + +func (f AuditQueryFilter) matchesCreatedAt(createdAt time.Time) bool { + if !f.CreatedFrom.IsZero() && createdAt.Before(f.CreatedFrom) { return false } - if !f.CreatedTo.IsZero() && record.CreatedAt.After(f.CreatedTo) { + if !f.CreatedTo.IsZero() && createdAt.After(f.CreatedTo) { return false } return true From 7d40a54d27c7eaa9805e547f26db4e4c16b055c8 Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 10:49:54 +0800 Subject: [PATCH 77/95] feat(toolpolicy): audit budget remaining summary --- platform/toolpolicy/policy.go | 106 +++++++++++++++++++---------- platform/toolpolicy/policy_test.go | 45 ++++++++++++ 2 files changed, 115 insertions(+), 36 deletions(-) diff --git a/platform/toolpolicy/policy.go b/platform/toolpolicy/policy.go index 7e8fc0bc7e..467a40747f 100644 --- a/platform/toolpolicy/policy.go +++ b/platform/toolpolicy/policy.go @@ -38,24 +38,26 @@ type Policy struct { // ApprovalSummary is the safe approval-facing summary of one tool call. type ApprovalSummary struct { - TenantID string - AppID string - PolicyID string - ToolName string - ToolCallID string - Decision tool.PermissionAction - Reason string - ArgumentsDigest string - ArgumentsBytes int - RequiresApproval bool - ReadOnly bool - Destructive bool - OpenWorld bool - ConcurrencySafe bool - SearchOrRead bool - MaxResultSize int - RedactionVersion string - CreatedAt time.Time + TenantID string + AppID string + PolicyID string + ToolName string + ToolCallID string + Decision tool.PermissionAction + Reason string + ArgumentsDigest string + ArgumentsBytes int + ToolBudgetRemainingDigest string + ToolBudgetRemainingBytes int + RequiresApproval bool + ReadOnly bool + Destructive bool + OpenWorld bool + ConcurrencySafe bool + SearchOrRead bool + MaxResultSize int + RedactionVersion string + CreatedAt time.Time } type auditContextKey struct{} @@ -506,25 +508,28 @@ func (p *Policy) ApprovalSummary( return ApprovalSummary{}, err } argumentsDigest, argumentsBytes := argumentDigest(req.Arguments) + toolBudgetRemainingDigest, toolBudgetRemainingBytes := policyJSONDigest(p.policy.ToolBudgetRemainingJSON) summary := ApprovalSummary{ - TenantID: strings.TrimSpace(p.policy.TenantID), - AppID: strings.TrimSpace(p.policy.AppID), - PolicyID: strings.TrimSpace(p.policy.PolicyID), - ToolName: name, - ToolCallID: strings.TrimSpace(req.ToolCallID), - Decision: decision.Action, - Reason: reason, - ArgumentsDigest: argumentsDigest, - ArgumentsBytes: argumentsBytes, - RequiresApproval: decision.Action == tool.PermissionActionAsk, - ReadOnly: req.Metadata.ReadOnly, - Destructive: req.Metadata.Destructive, - OpenWorld: req.Metadata.OpenWorld, - ConcurrencySafe: req.Metadata.ConcurrencySafe, - SearchOrRead: req.Metadata.SearchOrRead, - MaxResultSize: req.Metadata.MaxResultSize, - RedactionVersion: "platform-toolpolicy-v1", - CreatedAt: p.now(), + TenantID: strings.TrimSpace(p.policy.TenantID), + AppID: strings.TrimSpace(p.policy.AppID), + PolicyID: strings.TrimSpace(p.policy.PolicyID), + ToolName: name, + ToolCallID: strings.TrimSpace(req.ToolCallID), + Decision: decision.Action, + Reason: reason, + ArgumentsDigest: argumentsDigest, + ArgumentsBytes: argumentsBytes, + ToolBudgetRemainingDigest: toolBudgetRemainingDigest, + ToolBudgetRemainingBytes: toolBudgetRemainingBytes, + RequiresApproval: decision.Action == tool.PermissionActionAsk, + ReadOnly: req.Metadata.ReadOnly, + Destructive: req.Metadata.Destructive, + OpenWorld: req.Metadata.OpenWorld, + ConcurrencySafe: req.Metadata.ConcurrencySafe, + SearchOrRead: req.Metadata.SearchOrRead, + MaxResultSize: req.Metadata.MaxResultSize, + RedactionVersion: "platform-toolpolicy-v1", + CreatedAt: p.now(), } if err := summary.Validate(); err != nil { return ApprovalSummary{}, err @@ -540,6 +545,9 @@ func (s ApprovalSummary) Validate() error { if err := s.validateArguments(); err != nil { return err } + if err := s.validateToolBudgetRemaining(); err != nil { + return err + } if err := s.validateDecision(); err != nil { return err } @@ -603,6 +611,20 @@ func (s ApprovalSummary) validateArguments() error { return nil } +func (s ApprovalSummary) validateToolBudgetRemaining() error { + if s.ToolBudgetRemainingBytes < 0 { + return fmt.Errorf("tool_budget_remaining_bytes must be greater than or equal to 0") + } + if s.ToolBudgetRemainingBytes == 0 { + if s.ToolBudgetRemainingDigest != "" { + return fmt.Errorf("tool_budget_remaining_digest must be empty when tool_budget_remaining_bytes is 0") + } + } else if !validSHA256Digest(s.ToolBudgetRemainingDigest) { + return fmt.Errorf("tool_budget_remaining_digest must be sha256 followed by a 64 character hex digest") + } + return nil +} + func (s ApprovalSummary) validateDecision() error { switch s.Decision { case tool.PermissionActionAllow: @@ -688,6 +710,10 @@ func (s ApprovalSummary) DetailRef() string { parts = append(parts, "args:"+s.ArgumentsDigest) parts = append(parts, "args_bytes:"+strconv.Itoa(s.ArgumentsBytes)) } + if s.ToolBudgetRemainingDigest != "" { + parts = append(parts, "tool_budget_remaining:"+s.ToolBudgetRemainingDigest) + parts = append(parts, "tool_budget_remaining_bytes:"+strconv.Itoa(s.ToolBudgetRemainingBytes)) + } if s.RequiresApproval { parts = append(parts, "requires_approval:true") } @@ -711,6 +737,14 @@ func argumentDigest(args []byte) (string, int) { return "sha256:" + hex.EncodeToString(sum[:]), len(args) } +func policyJSONDigest(value string) (string, int) { + value = strings.TrimSpace(value) + if value == "" { + return "", 0 + } + return argumentDigest([]byte(value)) +} + var sha256DigestPattern = regexp.MustCompile(`^sha256:[a-f0-9]{64}$`) func validSHA256Digest(value string) bool { diff --git a/platform/toolpolicy/policy_test.go b/platform/toolpolicy/policy_test.go index b854e76e2a..72f4d1ca3c 100644 --- a/platform/toolpolicy/policy_test.go +++ b/platform/toolpolicy/policy_test.go @@ -328,6 +328,36 @@ func TestPolicyBuildsApprovalSummaryWithoutRawArguments(t *testing.T) { } } +func TestApprovalSummaryIncludesSafeToolBudgetRemainingDigest(t *testing.T) { + p := newPolicy( + t, + platform.ToolPolicy{ + ToolBudgetRemainingJSON: `{"tenant":"tenant","remaining_calls":1,"token":"sk-secret"}`, + }, + ) + req := request("workspace_write", tool.ToolMetadata{}, []byte(`{"path":"/private/file"}`)) + req.ToolCallID = "call-1" + + summary, err := p.ApprovalSummary(req, tool.AllowPermission(), "") + if err != nil { + t.Fatalf("ApprovalSummary: %v", err) + } + if summary.ToolBudgetRemainingBytes == 0 || + !strings.HasPrefix(summary.ToolBudgetRemainingDigest, "sha256:") { + t.Fatalf("expected tool budget remaining digest, got %+v", summary) + } + detail := summary.DetailRef() + if !strings.Contains(detail, "tool_budget_remaining:sha256:") || + !strings.Contains(detail, "tool_budget_remaining_bytes:") { + t.Fatalf("expected budget digest in detail, got %q", detail) + } + if strings.Contains(detail, "remaining_calls") || + strings.Contains(detail, "sk-secret") || + strings.Contains(detail, "tenant") { + t.Fatalf("summary detail leaked raw budget remaining JSON: %q", detail) + } +} + func TestApprovalSummaryValidationRejectsUnsafeOrInconsistentFields(t *testing.T) { now := time.Unix(300, 0) valid := ApprovalSummary{ @@ -390,6 +420,21 @@ func TestApprovalSummaryValidationRejectsUnsafeOrInconsistentFields(t *testing.T t.Fatalf("expected empty-arguments digest rejection, got %v", err) } + invalidBudgetDigest := valid + invalidBudgetDigest.ToolBudgetRemainingDigest = "raw-json" + invalidBudgetDigest.ToolBudgetRemainingBytes = 16 + if err := invalidBudgetDigest.Validate(); err == nil || + !strings.Contains(err.Error(), "tool_budget_remaining_digest") { + t.Fatalf("expected budget digest rejection, got %v", err) + } + + noBudgetBytes := valid + noBudgetBytes.ToolBudgetRemainingDigest = "sha256:9f86d081884c7d659a2feaa0c55ad015a3bf4f1b2b0b822cd15d6c15b0f00a08" + if err := noBudgetBytes.Validate(); err == nil || + !strings.Contains(err.Error(), "tool_budget_remaining_digest") { + t.Fatalf("expected empty-budget digest rejection, got %v", err) + } + allowNeedsApproval := valid allowNeedsApproval.Decision = tool.PermissionActionAllow if err := allowNeedsApproval.Validate(); err == nil || !strings.Contains(err.Error(), "requires_approval") { From 762411db6dee5f12bb9c0b9be3e396567ed9a921 Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 10:56:06 +0800 Subject: [PATCH 78/95] feat(platform): query redacted audit details --- platform/audit_query.go | 42 +++++++++++++++++++---------------- platform/audit_query_test.go | 43 ++++++++++++++++++++++++++++++++++++ 2 files changed, 66 insertions(+), 19 deletions(-) diff --git a/platform/audit_query.go b/platform/audit_query.go index 6d4be5358d..1bcfd863c3 100644 --- a/platform/audit_query.go +++ b/platform/audit_query.go @@ -16,25 +16,26 @@ import ( // AuditQueryFilter scopes audit retrieval to one tenant and optional safe dimensions. type AuditQueryFilter struct { - TenantID string - AppID string - AuditID string - Channel string - BindingID string - UserIDHash string - SessionID string - RequestID string - MessageID string - AgentName string - ModelName string - ToolName string - Decision string - ErrorType string - TraceID string - RedactionVersion string - CreatedFrom time.Time - CreatedTo time.Time - Limit int + TenantID string + AppID string + AuditID string + Channel string + BindingID string + UserIDHash string + SessionID string + RequestID string + MessageID string + AgentName string + ModelName string + ToolName string + Decision string + ErrorType string + TraceID string + RedactedDetailRef string + RedactionVersion string + CreatedFrom time.Time + CreatedTo time.Time + Limit int } // QueryAudit returns audit records matching one tenant-scoped filter. @@ -86,6 +87,7 @@ func (f AuditQueryFilter) normalize() (AuditQueryFilter, error) { safeTextField{"decision", f.Decision}, safeTextField{"error_type", f.ErrorType}, safeTextField{"trace_id", f.TraceID}, + safeTextField{"redacted_detail_ref", f.RedactedDetailRef}, safeTextField{"redaction_version", f.RedactionVersion}, ); err != nil { return AuditQueryFilter{}, err @@ -110,6 +112,7 @@ func (f AuditQueryFilter) normalize() (AuditQueryFilter, error) { f.Decision = strings.TrimSpace(f.Decision) f.ErrorType = strings.TrimSpace(f.ErrorType) f.TraceID = strings.TrimSpace(f.TraceID) + f.RedactedDetailRef = strings.TrimSpace(f.RedactedDetailRef) f.RedactionVersion = strings.TrimSpace(f.RedactionVersion) return f, nil } @@ -140,6 +143,7 @@ func (f AuditQueryFilter) matchesOptionalFields(record AuditRecord) bool { {f.ToolName, record.ToolName}, {f.Decision, record.Decision}, {f.ErrorType, record.ErrorType}, + {f.RedactedDetailRef, record.RedactedDetailRef}, {f.RedactionVersion, record.RedactionVersion}, {f.TraceID, record.TraceID}, } { diff --git a/platform/audit_query_test.go b/platform/audit_query_test.go index 98f0089605..b848d5c7c0 100644 --- a/platform/audit_query_test.go +++ b/platform/audit_query_test.go @@ -116,6 +116,46 @@ func TestQueryAuditFiltersRuntimeAndRedactionDimensions(t *testing.T) { } } +func TestQueryAuditFiltersBudgetDecisionByRedactedDetailRef(t *testing.T) { + now := time.Date(2026, 7, 8, 10, 0, 0, 0, time.UTC) + quota := TenantQuota{MaxCost: 1.00} + estimate := UsageEstimate{PromptTokens: 10, CompletionTokens: 5, Cost: 2.00} + decision, err := quota.Check(estimate) + if err != nil { + t.Fatalf("quota check: %v", err) + } + target, err := NewBudgetDecisionAuditRecord(BudgetDecisionAuditInput{ + TenantID: "tenant", + AppID: "app", + RequestID: "request-1", + TraceID: "trace-1", + Decision: decision, + Estimate: estimate, + Quota: quota, + CreatedAt: now, + }) + if err != nil { + t.Fatalf("NewBudgetDecisionAuditRecord: %v", err) + } + other := target + other.AuditID = "other-budget-audit" + other.RedactedDetailRef = strings.ReplaceAll(target.RedactedDetailRef, "estimated_cost:2.000000", "estimated_cost:3.000000") + + matches, err := QueryAudit([]AuditRecord{other, target}, AuditQueryFilter{ + TenantID: "tenant", + ToolName: "budget:tenant", + Decision: string(BudgetDecisionOutcomeDeny), + RedactionVersion: "platform-budget-decision-v1", + RedactedDetailRef: " " + target.RedactedDetailRef + " ", + }) + if err != nil { + t.Fatalf("query budget audit detail: %v", err) + } + if len(matches) != 1 || matches[0].AuditID != target.AuditID { + t.Fatalf("expected target budget audit, got %+v", matches) + } +} + func TestAuditSinkQueryUsesSnapshotAndReturnsCopies(t *testing.T) { sink := NewInMemoryAuditSink() record := auditRecordForQuery("tenant", "audit-1", "app", "telegram", "binding", "session", "request", "message", "tool", "allow", "trace", time.Now()) @@ -164,6 +204,9 @@ func TestQueryAuditRejectsUnsafeFilterValues(t *testing.T) { "redaction_version": { RedactionVersion: "platform-v1 Authorization: Bearer raw-token", }, + "redacted_detail_ref": { + RedactedDetailRef: "outcome:deny Authorization: Bearer raw-token", + }, } for name, filter := range tests { t.Run(name, func(t *testing.T) { From ae9a3a803d743d60b6924858fb64c8c2743f4b1b Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 11:02:49 +0800 Subject: [PATCH 79/95] feat(gateway): enrich budget audit runtime context --- platform/gateway/service.go | 2 ++ platform/gateway/service_test.go | 19 +++++++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/platform/gateway/service.go b/platform/gateway/service.go index 38f1cf4b36..fc7219466b 100644 --- a/platform/gateway/service.go +++ b/platform/gateway/service.go @@ -532,6 +532,8 @@ func (s *Service) writeBudgetDeniedAudit( record.UserIDHash = platform.UserIDHash(msg.TenantID, msg.Channel, msg.ExternalUserID) record.SessionID = auditContext.SessionID record.MessageID = msg.PlatformMessageID + record.AgentName = runtime.App.AgentName + record.ModelName = usageModelName(runtime) if err := record.Validate(); err != nil { s.writeRejectAuditWithContextTo(ctx, auditSink, msg, start, err, auditContext) return diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go index a197046780..f2ba8453f0 100644 --- a/platform/gateway/service_test.go +++ b/platform/gateway/service_test.go @@ -1506,6 +1506,13 @@ func TestServiceHandleInboundRejectsBudgetExceededBeforeIdempotency(t *testing.T registry := NewInMemoryRegistry() r := &recordingRunner{response: "unused"} runtime := validRuntime("tenant-a", r) + runtime.App.AgentName = "budget-agent" + runtime.App.ModelProfileID = "profile-gpt" + runtime.ModelProfile = platform.ModelProfile{ + TenantID: "tenant-a", + ProfileID: "profile-gpt", + Model: "gpt-budget", + } runtime.Tenant.QuotaJSON = `{"max_total_tokens":10}` require.NoError(t, registry.Register(runtime)) audit := platform.NewInMemoryAuditSink() @@ -1544,6 +1551,8 @@ func TestServiceHandleInboundRejectsBudgetExceededBeforeIdempotency(t *testing.T assert.Equal(t, msg.Channel, record.Channel) assert.Equal(t, msg.BindingID, record.BindingID) assert.Equal(t, msg.PlatformMessageID, record.MessageID) + assert.Equal(t, "budget-agent", record.AgentName) + assert.Equal(t, "gpt-budget", record.ModelName) assert.NotEmpty(t, record.SessionID) assert.NotEmpty(t, record.InternalUserID) assert.Equal(t, platform.UserIDHash(msg.TenantID, msg.Channel, msg.ExternalUserID), record.UserIDHash) @@ -1555,6 +1564,16 @@ func TestServiceHandleInboundRejectsBudgetExceededBeforeIdempotency(t *testing.T assert.Equal(t, "req-budget", estimateRequest.RequestID) assert.NotEmpty(t, estimateRequest.SessionID) assert.NotEmpty(t, estimateRequest.InternalUserID) + + matches, queryErr := audit.Query(platform.AuditQueryFilter{ + TenantID: "tenant-a", + ToolName: "budget:tenant", + AgentName: "budget-agent", + ModelName: "gpt-budget", + }) + require.NoError(t, queryErr) + require.Len(t, matches, 1) + assert.Equal(t, record.AuditID, matches[0].AuditID) } func TestServiceHandleInboundAllowsWithinBudget(t *testing.T) { From cbdb32d5bf9a9e5ece738fc01d8c394ed78d3b51 Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 11:16:32 +0800 Subject: [PATCH 80/95] feat(gateway): record redaction failed audit --- platform/gateway/service.go | 24 +++++++++++++++++++++ platform/gateway/service_test.go | 37 ++++++++++++++++++++++++++++++++ 2 files changed, 61 insertions(+) diff --git a/platform/gateway/service.go b/platform/gateway/service.go index fc7219466b..2259decf96 100644 --- a/platform/gateway/service.go +++ b/platform/gateway/service.go @@ -1342,9 +1342,33 @@ func (s *Service) writeAuditTo( if isNilInterfaceValue(auditSink) { return } + if err := record.Validate(); err != nil { + if !isAuditRedactionFailure(err) { + return + } + record = redactionFailedAuditRecord(record, err) + } _ = auditSink.WriteAudit(ctx, record) } +func isAuditRedactionFailure(err error) bool { + return err != nil && strings.Contains(err.Error(), "contains unredacted sensitive content") +} + +func redactionFailedAuditRecord(record platform.AuditRecord, err error) platform.AuditRecord { + return platform.AuditRecord{ + AuditID: platform.AuditID(record.TenantID, record.AppID, record.RequestID, record.TraceID, record.MessageID, "redaction_failed"), + TenantID: record.TenantID, + AppID: record.AppID, + Decision: "redaction_failed", + DecisionReason: "audit redaction failed", + ErrorType: "redaction_failed", + RedactedDetailRef: "failed_audit:" + platform.AuditID(record.AuditID, record.ToolName, record.Decision, fmt.Sprintf("%T", err)), + RedactionVersion: "platform-gateway-redaction-failed-v1", + CreatedAt: record.CreatedAt, + } +} + func (s *Service) auditSinkForRuntime(runtime Runtime) platform.AuditSink { if !isNilInterfaceValue(runtime.Audit) { return runtime.Audit diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go index f2ba8453f0..290aa6d3f6 100644 --- a/platform/gateway/service_test.go +++ b/platform/gateway/service_test.go @@ -1501,6 +1501,43 @@ func TestServiceHandleInboundRejectsMissingRequiredMention(t *testing.T) { assert.Equal(t, ErrBindingMentionRequired.Error(), audit.Records()[0].DecisionReason) } +func TestServiceWriteAuditRecordsRedactionFailureFallback(t *testing.T) { + audit := platform.NewInMemoryAuditSink() + svc := NewService( + NewInMemoryRegistry(), + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + ) + unsafe := platform.AuditRecord{ + AuditID: "audit-unsafe", + TenantID: "tenant-a", + AppID: "app", + RequestID: "request-1", + MessageID: "msg-1", + ToolName: "workspace_write", + Decision: "reject", + DecisionReason: "Authorization: Bearer raw-token", + CreatedAt: time.Unix(1500, 0), + } + + svc.writeAuditTo(context.Background(), audit, unsafe) + + records := audit.Records() + require.Len(t, records, 1) + record := records[0] + assert.Equal(t, "tenant-a", record.TenantID) + assert.Equal(t, "app", record.AppID) + assert.Equal(t, "redaction_failed", record.Decision) + assert.Equal(t, "audit redaction failed", record.DecisionReason) + assert.Equal(t, "redaction_failed", record.ErrorType) + assert.Equal(t, "platform-gateway-redaction-failed-v1", record.RedactionVersion) + assert.NotEmpty(t, record.AuditID) + assert.Contains(t, record.RedactedDetailRef, "failed_audit:audit_") + assert.NotContains(t, record.DecisionReason, "raw-token") + assert.NotContains(t, record.RedactedDetailRef, "raw-token") + assert.NotContains(t, record.RedactedDetailRef, "Authorization") +} + func TestServiceHandleInboundRejectsBudgetExceededBeforeIdempotency(t *testing.T) { ctx := context.Background() registry := NewInMemoryRegistry() From 455255504e31ecdd4ed77b43b10f16f36e2bf229 Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 11:23:23 +0800 Subject: [PATCH 81/95] feat(platform): add tool approval audit records --- platform/tool_approval_audit.go | 155 +++++++++++++++++++++++++ platform/tool_approval_audit_test.go | 166 +++++++++++++++++++++++++++ 2 files changed, 321 insertions(+) create mode 100644 platform/tool_approval_audit.go create mode 100644 platform/tool_approval_audit_test.go diff --git a/platform/tool_approval_audit.go b/platform/tool_approval_audit.go new file mode 100644 index 0000000000..b9a44d2f5e --- /dev/null +++ b/platform/tool_approval_audit.go @@ -0,0 +1,155 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "crypto/sha256" + "encoding/hex" + "fmt" + "strings" + "time" +) + +// ToolApprovalDecision is the externally visible tool approval audit decision. +type ToolApprovalDecision string + +const ( + // ToolApprovalDecisionRequested records that a tool call requires approval. + ToolApprovalDecisionRequested ToolApprovalDecision = "approval_requested" + // ToolApprovalDecisionApproved records that a tool approval was granted. + ToolApprovalDecisionApproved ToolApprovalDecision = "approval_approved" + // ToolApprovalDecisionRejected records that a tool approval was rejected. + ToolApprovalDecisionRejected ToolApprovalDecision = "approval_rejected" +) + +// ToolApprovalAuditInput contains safe dimensions for a tool approval boundary. +type ToolApprovalAuditInput struct { + TenantID string + AppID string + ToolName string + ToolCallID string + Decision ToolApprovalDecision + DecisionReason string + ApproverUserID string + RequestID string + TraceID string + ArgumentSummaryRef string + CreatedAt time.Time +} + +// NewToolApprovalAuditRecord maps one tool approval boundary into a safe audit record. +func NewToolApprovalAuditRecord(input ToolApprovalAuditInput) (AuditRecord, error) { + normalized, err := input.normalize() + if err != nil { + return AuditRecord{}, err + } + record := AuditRecord{ + TenantID: normalized.TenantID, + AppID: normalized.AppID, + AuditID: normalized.auditID(), + RequestID: normalized.RequestID, + TraceID: normalized.TraceID, + UserIDHash: normalized.approverHash(), + ToolName: normalized.ToolName, + Decision: string(normalized.Decision), + DecisionReason: normalized.DecisionReason, + RedactedDetailRef: normalized.detailRef(), + RedactionVersion: "platform-tool-approval-v1", + CreatedAt: normalized.CreatedAt, + } + if err := record.Validate(); err != nil { + return AuditRecord{}, err + } + return record, nil +} + +func (i ToolApprovalAuditInput) normalize() (ToolApprovalAuditInput, error) { + i.TenantID = strings.TrimSpace(i.TenantID) + if i.TenantID == "" { + return ToolApprovalAuditInput{}, ErrTenantIDRequired + } + i.AppID = strings.TrimSpace(i.AppID) + i.ToolName = strings.TrimSpace(i.ToolName) + if i.ToolName == "" { + return ToolApprovalAuditInput{}, fmt.Errorf("tool_name is required") + } + i.ToolCallID = strings.TrimSpace(i.ToolCallID) + if i.ToolCallID == "" { + return ToolApprovalAuditInput{}, fmt.Errorf("tool_call_id is required") + } + i.Decision = ToolApprovalDecision(strings.TrimSpace(string(i.Decision))) + if !i.Decision.valid() { + return ToolApprovalAuditInput{}, fmt.Errorf("invalid tool approval decision %q", i.Decision) + } + i.DecisionReason = strings.TrimSpace(i.DecisionReason) + i.ApproverUserID = strings.TrimSpace(i.ApproverUserID) + if i.Decision != ToolApprovalDecisionRequested && i.ApproverUserID == "" { + return ToolApprovalAuditInput{}, fmt.Errorf("approver_user_id is required for decided approvals") + } + i.RequestID = strings.TrimSpace(i.RequestID) + i.TraceID = strings.TrimSpace(i.TraceID) + i.ArgumentSummaryRef = strings.TrimSpace(i.ArgumentSummaryRef) + if err := validateAuditRedactedFields( + safeTextField{"app_id", i.AppID}, + safeTextField{"tool_name", i.ToolName}, + safeTextField{"tool_call_id", i.ToolCallID}, + safeTextField{"decision", string(i.Decision)}, + safeTextField{"decision_reason", i.DecisionReason}, + safeTextField{"request_id", i.RequestID}, + safeTextField{"trace_id", i.TraceID}, + safeTextField{"argument_summary_ref", i.ArgumentSummaryRef}, + ); err != nil { + return ToolApprovalAuditInput{}, err + } + return i, nil +} + +func (i ToolApprovalAuditInput) auditID() string { + return AuditID( + i.TenantID, + i.AppID, + i.ToolName, + i.ToolCallID, + string(i.Decision), + i.ApproverUserID, + ) +} + +func (i ToolApprovalAuditInput) approverHash() string { + if i.ApproverUserID == "" { + return "" + } + return UserIDHash(i.TenantID, "approval", i.ApproverUserID) +} + +func (i ToolApprovalAuditInput) detailRef() string { + parts := []string{ + "tool_call_id:" + i.ToolCallID, + } + if i.ArgumentSummaryRef != "" { + sum := sha256.Sum256([]byte(i.ArgumentSummaryRef)) + parts = append(parts, "args_ref_sha256:"+hex.EncodeToString(sum[:])) + parts = append(parts, fmt.Sprintf("args_ref_bytes:%d", len(i.ArgumentSummaryRef))) + } + if approverHash := i.approverHash(); approverHash != "" { + parts = append(parts, "approver_hash:"+approverHash) + } + return strings.Join(parts, " ") +} + +func (d ToolApprovalDecision) valid() bool { + switch d { + case ToolApprovalDecisionRequested, + ToolApprovalDecisionApproved, + ToolApprovalDecisionRejected: + return true + default: + return false + } +} diff --git a/platform/tool_approval_audit_test.go b/platform/tool_approval_audit_test.go new file mode 100644 index 0000000000..4d3d30aaa8 --- /dev/null +++ b/platform/tool_approval_audit_test.go @@ -0,0 +1,166 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package platform + +import ( + "errors" + "strings" + "testing" + "time" +) + +func TestNewToolApprovalAuditRecordBuildsRequestedRecord(t *testing.T) { + createdAt := time.Date(2026, 7, 8, 10, 0, 0, 0, time.UTC) + input := ToolApprovalAuditInput{ + TenantID: " tenant ", + AppID: " app ", + ToolName: " workspace_write ", + ToolCallID: " call-1 ", + Decision: ToolApprovalDecisionRequested, + DecisionReason: "high-risk tool requires approval", + RequestID: " request-1 ", + TraceID: " trace-1 ", + ArgumentSummaryRef: `args:sha256:0123456789abcdef args_bytes:128`, + CreatedAt: createdAt, + } + + record, err := NewToolApprovalAuditRecord(input) + if err != nil { + t.Fatalf("new tool approval audit: %v", err) + } + if record.TenantID != "tenant" || + record.AppID != "app" || + record.ToolName != "workspace_write" || + record.Decision != "approval_requested" || + record.DecisionReason != "high-risk tool requires approval" || + record.RequestID != "request-1" || + record.TraceID != "trace-1" || + !record.CreatedAt.Equal(createdAt) { + t.Fatalf("unexpected record: %+v", record) + } + if record.UserIDHash != "" { + t.Fatalf("requested approval should not require approver hash, got %+v", record) + } + if record.AuditID == "" || record.RedactionVersion != "platform-tool-approval-v1" { + t.Fatalf("expected audit id and redaction version, got %+v", record) + } + if !strings.Contains(record.RedactedDetailRef, "tool_call_id:call-1") || + !strings.Contains(record.RedactedDetailRef, "args_ref_sha256:") || + !strings.Contains(record.RedactedDetailRef, "args_ref_bytes:") { + t.Fatalf("unexpected redacted detail ref: %q", record.RedactedDetailRef) + } + if strings.Contains(record.RedactedDetailRef, "args:sha256:") || + strings.Contains(record.RedactedDetailRef, "args_bytes:128") { + t.Fatalf("approval audit leaked raw argument summary: %q", record.RedactedDetailRef) + } + + again, err := NewToolApprovalAuditRecord(input) + if err != nil { + t.Fatalf("new duplicate tool approval audit: %v", err) + } + if record.AuditID != again.AuditID { + t.Fatalf("expected stable audit id, got %q and %q", record.AuditID, again.AuditID) + } +} + +func TestNewToolApprovalAuditRecordBuildsDecidedRecordWithApproverHash(t *testing.T) { + record, err := NewToolApprovalAuditRecord(ToolApprovalAuditInput{ + TenantID: "tenant", + AppID: "app", + ToolName: "workspace_write", + ToolCallID: "call-1", + Decision: ToolApprovalDecisionApproved, + DecisionReason: "approved by security reviewer", + ApproverUserID: "security@example.com", + RequestID: "request-1", + TraceID: "trace-1", + CreatedAt: time.Unix(200, 0), + }) + if err != nil { + t.Fatalf("new decided tool approval audit: %v", err) + } + if record.Decision != "approval_approved" || + record.UserIDHash == "" || + !strings.HasPrefix(record.UserIDHash, "user_hash_") || + !strings.Contains(record.RedactedDetailRef, "approver_hash:user_hash_") { + t.Fatalf("expected decided approval with approver hash, got %+v", record) + } + if strings.Contains(record.UserIDHash, "security@example.com") || + strings.Contains(record.RedactedDetailRef, "security@example.com") { + t.Fatalf("approval audit leaked raw approver id: %+v", record) + } +} + +func TestNewToolApprovalAuditRecordRejectsInvalidInputs(t *testing.T) { + base := validToolApprovalAuditInput() + + missingTenant := base + missingTenant.TenantID = " " + if _, err := NewToolApprovalAuditRecord(missingTenant); !errors.Is(err, ErrTenantIDRequired) { + t.Fatalf("expected tenant requirement, got %v", err) + } + + missingTool := base + missingTool.ToolName = " " + if _, err := NewToolApprovalAuditRecord(missingTool); err == nil || + !strings.Contains(err.Error(), "tool_name") { + t.Fatalf("expected tool name requirement, got %v", err) + } + + missingCall := base + missingCall.ToolCallID = " " + if _, err := NewToolApprovalAuditRecord(missingCall); err == nil || + !strings.Contains(err.Error(), "tool_call_id") { + t.Fatalf("expected tool call id requirement, got %v", err) + } + + unknownDecision := base + unknownDecision.Decision = "bypassed" + if _, err := NewToolApprovalAuditRecord(unknownDecision); err == nil || + !strings.Contains(err.Error(), "invalid tool approval decision") { + t.Fatalf("expected decision validation, got %v", err) + } + + missingApprover := base + missingApprover.Decision = ToolApprovalDecisionRejected + if _, err := NewToolApprovalAuditRecord(missingApprover); err == nil || + !strings.Contains(err.Error(), "approver_user_id") { + t.Fatalf("expected approver requirement, got %v", err) + } +} + +func TestNewToolApprovalAuditRecordRejectsSensitivePublicFields(t *testing.T) { + input := validToolApprovalAuditInput() + input.DecisionReason = "Authorization: Bearer raw-token" + if _, err := NewToolApprovalAuditRecord(input); err == nil || + !strings.Contains(err.Error(), "decision_reason") { + t.Fatalf("expected sensitive decision reason rejection, got %v", err) + } + + input = validToolApprovalAuditInput() + input.ArgumentSummaryRef = "token=sk-secret" + if _, err := NewToolApprovalAuditRecord(input); err == nil || + !strings.Contains(err.Error(), "argument_summary_ref") { + t.Fatalf("expected sensitive argument summary rejection, got %v", err) + } +} + +func validToolApprovalAuditInput() ToolApprovalAuditInput { + return ToolApprovalAuditInput{ + TenantID: "tenant", + AppID: "app", + ToolName: "workspace_write", + ToolCallID: "call-1", + Decision: ToolApprovalDecisionRequested, + DecisionReason: "approval required", + RequestID: "request", + TraceID: "trace", + CreatedAt: time.Unix(100, 0), + } +} From 6849a5c7f32623072268ba64c057289737084dbb Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 11:39:46 +0800 Subject: [PATCH 82/95] feat(approval): audit tool approval runtime decisions --- plugin/guardrail/approval/approval.go | 62 ++++++++++ plugin/guardrail/approval/approval_test.go | 124 +++++++++++++++++++ plugin/guardrail/approval/audit.go | 131 +++++++++++++++++++++ plugin/guardrail/approval/option.go | 32 ++++- 4 files changed, 348 insertions(+), 1 deletion(-) create mode 100644 plugin/guardrail/approval/audit.go diff --git a/plugin/guardrail/approval/approval.go b/plugin/guardrail/approval/approval.go index 4cfca29ac9..dbf231b7de 100644 --- a/plugin/guardrail/approval/approval.go +++ b/plugin/guardrail/approval/approval.go @@ -13,9 +13,11 @@ import ( "context" "fmt" "strings" + "time" "trpc.group/trpc-go/trpc-agent-go/log" "trpc.group/trpc-go/trpc-agent-go/model" + "trpc.group/trpc-go/trpc-agent-go/platform" "trpc.group/trpc-go/trpc-agent-go/plugin" "trpc.group/trpc-go/trpc-agent-go/plugin/guardrail/approval/review" "trpc.group/trpc-go/trpc-agent-go/tool" @@ -28,6 +30,9 @@ type Plugin struct { defaultToolPolicy ToolPolicy toolPolicies map[string]ToolPolicy tokenCounter model.TokenCounter + auditSink platform.AuditSink + approverUserID string + now func() time.Time } // New creates a new approval plugin. @@ -47,12 +52,18 @@ func New(options ...Option) (*Plugin, error) { if requiresReviewer(opts) && opts.reviewer == nil { return nil, fmt.Errorf("newing approval plugin: reviewer is nil") } + if opts.auditSink != nil && requiresReviewer(opts) && strings.TrimSpace(opts.approverUserID) == "" { + return nil, fmt.Errorf("newing approval plugin: approver user id is required when approval audit is enabled") + } return &Plugin{ name: opts.name, reviewer: opts.reviewer, defaultToolPolicy: opts.defaultToolPolicy, toolPolicies: opts.toolPolicies, tokenCounter: model.NewSimpleTokenCounter(), + auditSink: opts.auditSink, + approverUserID: strings.TrimSpace(opts.approverUserID), + now: opts.now, }, nil } @@ -95,6 +106,23 @@ func (p *Plugin) beforeTool() tool.BeforeToolCallbackStructured { CustomResult: fmt.Sprintf("approval review failed for tool %q: %v", args.ToolName, err), }, nil } + if err := p.writeApprovalAudit( + ctx, + args, + platform.ToolApprovalDecisionRequested, + approvalAuditDecisionReason(platform.ToolApprovalDecisionRequested), + "", + ); err != nil { + log.ErrorfContext( + ctx, + "Automatic approval review denied: approval audit failed for tool %q: %v", + args.ToolName, + err, + ) + return &tool.BeforeToolResult{ + CustomResult: fmt.Sprintf("approval audit failed for tool %q: %v", args.ToolName, err), + }, nil + } decision, err := p.reviewer.Review(ctx, req) if err != nil { log.ErrorfContext( @@ -122,6 +150,23 @@ func (p *Plugin) beforeTool() tool.BeforeToolCallbackStructured { riskLevel := strings.TrimSpace(decision.RiskLevel) reason := strings.TrimSpace(decision.Reason) if decision.Approved { + if err := p.writeApprovalAudit( + ctx, + args, + platform.ToolApprovalDecisionApproved, + approvalAuditDecisionReason(platform.ToolApprovalDecisionApproved), + p.auditApproverUserID(), + ); err != nil { + log.ErrorfContext( + ctx, + "Automatic approval review denied: approval audit failed for tool %q: %v", + args.ToolName, + err, + ) + return &tool.BeforeToolResult{ + CustomResult: fmt.Sprintf("approval audit failed for tool %q: %v", args.ToolName, err), + }, nil + } log.InfofContext( ctx, "Automatic approval review approved (risk: %s): %s", @@ -135,6 +180,23 @@ func (p *Plugin) beforeTool() tool.BeforeToolCallbackStructured { riskLevel, reason, ) + if err := p.writeApprovalAudit( + ctx, + args, + platform.ToolApprovalDecisionRejected, + approvalAuditDecisionReason(platform.ToolApprovalDecisionRejected), + p.auditApproverUserID(), + ); err != nil { + log.ErrorfContext( + ctx, + "Automatic approval review denied: approval audit failed for tool %q: %v", + args.ToolName, + err, + ) + return &tool.BeforeToolResult{ + CustomResult: fmt.Sprintf("approval audit failed for tool %q: %v", args.ToolName, err), + }, nil + } log.WarnContext(ctx, denyMessage) return &tool.BeforeToolResult{ CustomResult: denyMessage, diff --git a/plugin/guardrail/approval/approval_test.go b/plugin/guardrail/approval/approval_test.go index db25fe2937..123b61bb10 100644 --- a/plugin/guardrail/approval/approval_test.go +++ b/plugin/guardrail/approval/approval_test.go @@ -13,6 +13,7 @@ import ( "errors" "fmt" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -20,6 +21,7 @@ import ( "trpc.group/trpc-go/trpc-agent-go/event" approvallog "trpc.group/trpc-go/trpc-agent-go/log" "trpc.group/trpc-go/trpc-agent-go/model" + "trpc.group/trpc-go/trpc-agent-go/platform" "trpc.group/trpc-go/trpc-agent-go/plugin" approvalreview "trpc.group/trpc-go/trpc-agent-go/plugin/guardrail/approval/review" guardtranscript "trpc.group/trpc-go/trpc-agent-go/plugin/guardrail/internal/transcript" @@ -72,6 +74,15 @@ func TestNew_EmptyToolPolicyName(t *testing.T) { require.Contains(t, err.Error(), "tool policy name is empty") } +func TestNew_RequiresApproverUserIDWhenAuditEnabledForApproval(t *testing.T) { + _, err := New( + WithReviewer(&stubReviewer{}), + WithAuditSink(platform.NewInMemoryAuditSink()), + ) + require.Error(t, err) + require.Contains(t, err.Error(), "approver user id is required") +} + func TestNew_WithName(t *testing.T) { p, err := New(WithReviewer(&stubReviewer{}), WithName("tool-approval")) require.NoError(t, err) @@ -260,6 +271,119 @@ func TestBeforeTool_ReviewerApprovedLogsInfo(t *testing.T) { ) } +func TestBeforeTool_RequireApprovalWritesApprovedAuditRecords(t *testing.T) { + now := time.Date(2026, 7, 11, 10, 30, 0, 0, time.UTC) + audit := platform.NewInMemoryAuditSink() + p, err := New( + WithReviewer(&stubReviewer{ + reviewFn: func(ctx context.Context, req *approvalreview.Request) (*approvalreview.Decision, error) { + return &approvalreview.Decision{ + Approved: true, + RiskScore: 18, + RiskLevel: "low", + Reason: "Approved command: git status --short.", + }, nil + }, + }), + WithAuditSink(audit), + WithApproverUserID("security@example.com"), + withNow(func() time.Time { return now }), + ) + require.NoError(t, err) + + callbacks := registeredToolCallbacks(t, p) + ctx := ContextWithAuditContext(context.Background(), AuditContext{ + TenantID: "tenant", + AppID: "app", + RequestID: "request-1", + TraceID: "trace-1", + }) + result, runErr := callbacks.RunBeforeTool(ctx, &tool.BeforeToolArgs{ + ToolName: "shell", + ToolCallID: "call-1", + Arguments: []byte(`{"command":"git status --short"}`), + }) + require.NoError(t, runErr) + require.Nil(t, result) + + records := audit.Records() + require.Len(t, records, 2) + assert.Equal(t, "approval_requested", records[0].Decision) + assert.Equal(t, "approval_approved", records[1].Decision) + for _, record := range records { + assert.Equal(t, "tenant", record.TenantID) + assert.Equal(t, "app", record.AppID) + assert.Equal(t, "request-1", record.RequestID) + assert.Equal(t, "trace-1", record.TraceID) + assert.Equal(t, "shell", record.ToolName) + assert.True(t, record.CreatedAt.Equal(now)) + assert.Contains(t, record.RedactedDetailRef, "tool_call_id:call-1") + assert.Contains(t, record.RedactedDetailRef, "args_ref_sha256:") + assert.NotContains(t, record.RedactedDetailRef, "git status --short") + assert.NotContains(t, record.RedactedDetailRef, "security@example.com") + } + assert.Empty(t, records[0].UserIDHash) + assert.NotEmpty(t, records[1].UserIDHash) + assert.Equal(t, "tool approval approved", records[1].DecisionReason) +} + +func TestBeforeTool_RequireApprovalWritesRejectedAuditRecords(t *testing.T) { + now := time.Date(2026, 7, 11, 11, 0, 0, 0, time.UTC) + audit := platform.NewInMemoryAuditSink() + p, err := New( + WithReviewer(&stubReviewer{ + reviewFn: func(ctx context.Context, req *approvalreview.Request) (*approvalreview.Decision, error) { + return &approvalreview.Decision{ + Approved: false, + RiskScore: 95, + RiskLevel: "high", + Reason: "Command rm -rf workspace can delete workspace files.", + }, nil + }, + }), + WithAuditSink(audit), + WithApproverUserID("security@example.com"), + withNow(func() time.Time { return now }), + ) + require.NoError(t, err) + + callbacks := registeredToolCallbacks(t, p) + ctx := ContextWithAuditContext(context.Background(), AuditContext{ + TenantID: "tenant", + AppID: "app", + RequestID: "request-2", + TraceID: "trace-2", + }) + result, runErr := callbacks.RunBeforeTool(ctx, &tool.BeforeToolArgs{ + ToolName: "shell", + ToolCallID: "call-2", + Arguments: []byte(`{"command":"rm -rf workspace"}`), + }) + require.NoError(t, runErr) + require.NotNil(t, result) + require.Equal(t, "Automatic approval review denied (risk: high): Command rm -rf workspace can delete workspace files.", result.CustomResult) + + records := audit.Records() + require.Len(t, records, 2) + assert.Equal(t, "approval_requested", records[0].Decision) + assert.Equal(t, "approval_rejected", records[1].Decision) + for _, record := range records { + assert.Equal(t, "tenant", record.TenantID) + assert.Equal(t, "app", record.AppID) + assert.Equal(t, "request-2", record.RequestID) + assert.Equal(t, "trace-2", record.TraceID) + assert.Equal(t, "shell", record.ToolName) + assert.True(t, record.CreatedAt.Equal(now)) + assert.Contains(t, record.RedactedDetailRef, "tool_call_id:call-2") + assert.Contains(t, record.RedactedDetailRef, "args_ref_sha256:") + assert.NotContains(t, record.RedactedDetailRef, "rm -rf workspace") + assert.NotContains(t, record.RedactedDetailRef, "security@example.com") + } + assert.Empty(t, records[0].UserIDHash) + assert.NotEmpty(t, records[1].UserIDHash) + assert.Equal(t, "tool approval rejected", records[1].DecisionReason) +} + func TestBeforeTool_ReviewerErrorFailsClosed(t *testing.T) { original := approvallog.ErrorfContext var errorLog string diff --git a/plugin/guardrail/approval/audit.go b/plugin/guardrail/approval/audit.go new file mode 100644 index 0000000000..dd5ba76ef5 --- /dev/null +++ b/plugin/guardrail/approval/audit.go @@ -0,0 +1,131 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package approval + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "strconv" + "strings" + "time" + + "trpc.group/trpc-go/trpc-agent-go/agent" + "trpc.group/trpc-go/trpc-agent-go/platform" + "trpc.group/trpc-go/trpc-agent-go/tool" +) + +// AuditContext carries trusted platform context for tool approval audit. +type AuditContext struct { + TenantID string + AppID string + RequestID string + TraceID string +} + +type auditContextKey struct{} + +// ContextWithAuditContext attaches trusted platform audit context. +func ContextWithAuditContext(ctx context.Context, auditCtx AuditContext) context.Context { + if ctx == nil { + ctx = context.Background() + } + return context.WithValue(ctx, auditContextKey{}, auditCtx) +} + +func (p *Plugin) writeApprovalAudit( + ctx context.Context, + args *tool.BeforeToolArgs, + decision platform.ToolApprovalDecision, + reason string, + approverUserID string, +) error { + if p == nil || p.auditSink == nil || args == nil { + return nil + } + auditCtx := approvalAuditContextFrom(ctx) + record, err := platform.NewToolApprovalAuditRecord(platform.ToolApprovalAuditInput{ + TenantID: auditCtx.TenantID, + AppID: auditCtx.AppID, + ToolName: args.ToolName, + ToolCallID: args.ToolCallID, + Decision: decision, + DecisionReason: reason, + ApproverUserID: approverUserID, + RequestID: auditCtx.RequestID, + TraceID: auditCtx.TraceID, + ArgumentSummaryRef: argumentSummaryRef(args.Arguments), + CreatedAt: p.auditNow(), + }) + if err != nil { + return err + } + return p.auditSink.WriteAudit(ctx, record) +} + +func (p *Plugin) auditApproverUserID() string { + if p == nil { + return "" + } + return p.approverUserID +} + +func (p *Plugin) auditNow() time.Time { + if p == nil || p.now == nil { + return time.Now() + } + return p.now() +} + +func approvalAuditContextFrom(ctx context.Context) AuditContext { + var auditCtx AuditContext + if ctx != nil { + if value, ok := ctx.Value(auditContextKey{}).(AuditContext); ok { + auditCtx = AuditContext{ + TenantID: strings.TrimSpace(value.TenantID), + AppID: strings.TrimSpace(value.AppID), + RequestID: strings.TrimSpace(value.RequestID), + TraceID: strings.TrimSpace(value.TraceID), + } + } + } + if invocation, ok := agent.InvocationFromContext(ctx); ok && invocation != nil { + if auditCtx.AppID == "" && invocation.Session != nil { + auditCtx.AppID = strings.TrimSpace(invocation.Session.AppName) + } + if auditCtx.RequestID == "" { + auditCtx.RequestID = strings.TrimSpace(invocation.RunOptions.RequestID) + } + } + if auditCtx.TraceID == "" { + auditCtx.TraceID = auditCtx.RequestID + } + return auditCtx +} + +func argumentSummaryRef(args []byte) string { + if len(args) == 0 { + return "" + } + sum := sha256.Sum256(args) + return "args:sha256:" + hex.EncodeToString(sum[:]) + " args_bytes:" + strconv.Itoa(len(args)) +} + +func approvalAuditDecisionReason(decision platform.ToolApprovalDecision) string { + switch decision { + case platform.ToolApprovalDecisionRequested: + return "tool approval requested" + case platform.ToolApprovalDecisionApproved: + return "tool approval approved" + case platform.ToolApprovalDecisionRejected: + return "tool approval rejected" + default: + return "" + } +} diff --git a/plugin/guardrail/approval/option.go b/plugin/guardrail/approval/option.go index 41aebfff33..1ef7bf7b8b 100644 --- a/plugin/guardrail/approval/option.go +++ b/plugin/guardrail/approval/option.go @@ -8,7 +8,12 @@ package approval -import "trpc.group/trpc-go/trpc-agent-go/plugin/guardrail/approval/review" +import ( + "time" + + "trpc.group/trpc-go/trpc-agent-go/platform" + "trpc.group/trpc-go/trpc-agent-go/plugin/guardrail/approval/review" +) const defaultPluginName = "approval" @@ -20,6 +25,9 @@ type options struct { reviewer review.Reviewer defaultToolPolicy ToolPolicy toolPolicies map[string]ToolPolicy + auditSink platform.AuditSink + approverUserID string + now func() time.Time } func newOptions(opts ...Option) *options { @@ -66,3 +74,25 @@ func WithToolPolicy(name string, policy ToolPolicy) Option { opts.toolPolicies[name] = policy } } + +// WithAuditSink records approval request and decision boundaries to audit. +func WithAuditSink(sink platform.AuditSink) Option { + return func(opts *options) { + opts.auditSink = sink + } +} + +// WithApproverUserID sets the stable identity used for automatic reviewer decisions. +func WithApproverUserID(userID string) Option { + return func(opts *options) { + opts.approverUserID = userID + } +} + +func withNow(now func() time.Time) Option { + return func(opts *options) { + if now != nil { + opts.now = now + } + } +} From c050a3a6699af22d54cb91edf7f4989501589e06 Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 11:46:11 +0800 Subject: [PATCH 83/95] feat(gateway): propagate approval audit context --- platform/gateway/service.go | 7 ++ platform/gateway/service_test.go | 111 +++++++++++++++++++++++++++++++ 2 files changed, 118 insertions(+) diff --git a/platform/gateway/service.go b/platform/gateway/service.go index 2259decf96..334fc93886 100644 --- a/platform/gateway/service.go +++ b/platform/gateway/service.go @@ -28,6 +28,7 @@ import ( "trpc.group/trpc-go/trpc-agent-go/platform" "trpc.group/trpc-go/trpc-agent-go/platform/channeladapter" "trpc.group/trpc-go/trpc-agent-go/platform/toolpolicy" + "trpc.group/trpc-go/trpc-agent-go/plugin/guardrail/approval" telemetrytrace "trpc.group/trpc-go/trpc-agent-go/telemetry/trace" ) @@ -809,6 +810,12 @@ func (s *Service) runGatewayRunner( RequestID: input.RequestID, AgentName: runtime.App.AgentName, }) + runnerCtx = approval.ContextWithAuditContext(runnerCtx, approval.AuditContext{ + TenantID: runtime.Tenant.TenantID, + AppID: runtime.App.AppID, + RequestID: input.RequestID, + TraceID: input.RequestID, + }) setInboundTraceAttributes(runnerSpan, msg, input.SessionID, input.RequestID, input.InternalUserID) if input.FencingToken > 0 { runnerSpan.SetAttributes(attribute.Int64("storage.fencing_token", input.FencingToken)) diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go index 290aa6d3f6..1bd393db46 100644 --- a/platform/gateway/service_test.go +++ b/platform/gateway/service_test.go @@ -27,7 +27,11 @@ import ( "trpc.group/trpc-go/trpc-agent-go/model" "trpc.group/trpc-go/trpc-agent-go/platform" "trpc.group/trpc-go/trpc-agent-go/platform/channeladapter" + "trpc.group/trpc-go/trpc-agent-go/plugin" + "trpc.group/trpc-go/trpc-agent-go/plugin/guardrail/approval" + approvalreview "trpc.group/trpc-go/trpc-agent-go/plugin/guardrail/approval/review" telemetrytrace "trpc.group/trpc-go/trpc-agent-go/telemetry/trace" + "trpc.group/trpc-go/trpc-agent-go/tool" ) func TestServiceHandleInboundIsolatesTenants(t *testing.T) { @@ -846,6 +850,56 @@ func TestServiceHandleInboundPropagatesLeaseFencingToken(t *testing.T) { assert.Equal(t, int64(42), r.calls[0].fencingToken) } +func TestServiceHandleInboundPropagatesApprovalAuditContext(t *testing.T) { + ctx := context.Background() + audit := platform.NewInMemoryAuditSink() + approvalPlugin, err := approval.New( + approval.WithReviewer(approvalReviewerFunc(func(ctx context.Context, req *approvalreview.Request) (*approvalreview.Decision, error) { + return &approvalreview.Decision{ + Approved: true, + RiskScore: 10, + RiskLevel: "low", + Reason: "approved", + }, nil + })), + approval.WithAuditSink(audit), + approval.WithApproverUserID("security@example.com"), + ) + require.NoError(t, err) + pluginManager := plugin.MustNewManager(approvalPlugin) + r := &approvalCallbackRunner{ + response: "ok", + callbacks: pluginManager.ToolCallbacks(), + } + registry := NewInMemoryRegistry() + require.NoError(t, registry.Register(validRuntime("tenant-a", r))) + svc := NewService( + registry, + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + ) + + result, err := svc.HandleInbound(ctx, inbound("tenant-a", "msg-1", "user-1", "hello")) + + require.NoError(t, err) + require.Equal(t, "ok", result.Outbound.Content) + records := audit.Records() + require.Len(t, records, 2) + assert.Equal(t, "approval_requested", records[0].Decision) + assert.Equal(t, "approval_approved", records[1].Decision) + for _, record := range records { + assert.Equal(t, "tenant-a", record.TenantID) + assert.Equal(t, "app", record.AppID) + assert.Equal(t, result.RequestID, record.RequestID) + assert.Equal(t, result.RequestID, record.TraceID) + assert.Equal(t, "shell", record.ToolName) + assert.Contains(t, record.RedactedDetailRef, "tool_call_id:call-approval") + assert.NotContains(t, record.RedactedDetailRef, "rm -rf workspace") + } + assert.Empty(t, records[0].UserIDHash) + assert.NotEmpty(t, records[1].UserIDHash) +} + func TestServiceHandleInboundCancellationDuringEventCollectionReleasesSessionLease(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) registry := NewInMemoryRegistry() @@ -2281,6 +2335,21 @@ type cancelingRunner struct { calls []runnerCall } +type approvalReviewerFunc func(context.Context, *approvalreview.Request) (*approvalreview.Decision, error) + +func (f approvalReviewerFunc) Review( + ctx context.Context, + req *approvalreview.Request, +) (*approvalreview.Decision, error) { + return f(ctx, req) +} + +type approvalCallbackRunner struct { + response string + callbacks *tool.Callbacks + calls []runnerCall +} + type staticRegistry struct { runtime Runtime } @@ -2362,6 +2431,48 @@ func (r *cancelingRunner) Close() error { return nil } +func (r *approvalCallbackRunner) Run( + ctx context.Context, + userID string, + sessionID string, + message model.Message, + runOpts ...agent.RunOption, +) (<-chan *event.Event, error) { + runOptions := runOptionsFromOptions(runOpts...) + r.calls = append(r.calls, runnerCall{ + userID: userID, + sessionID: sessionID, + message: message, + requestID: runOptions.RequestID, + runOptions: runOptions, + }) + if r.callbacks != nil { + result, err := r.callbacks.RunBeforeTool(ctx, &tool.BeforeToolArgs{ + ToolCallID: "call-approval", + ToolName: "shell", + Declaration: &tool.Declaration{ + Name: "shell", + Description: "Runs shell commands.", + }, + Arguments: []byte(`{"command":"rm -rf workspace"}`), + }) + if err != nil { + return nil, err + } + if result != nil && result.CustomResult != nil { + return nil, errors.New("approval callback blocked tool") + } + } + out := make(chan *event.Event, 1) + out <- responseEvent(r.response, true) + close(out) + return out, nil +} + +func (r *approvalCallbackRunner) Close() error { + return nil +} + type hangingFirstRunner struct { mu sync.Mutex started chan struct{} From 1b9d19c713e122aeb9e014ed76b29633f813a53d Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 12:23:24 +0800 Subject: [PATCH 84/95] feat(worker): wire approval guardrail plugin --- platform/toolpolicy/policy.go | 35 ++++ platform/worker/builder.go | 107 ++++++++++- platform/worker/governance_test.go | 201 +++++++++++++++++++++ plugin/guardrail/approval/approval.go | 7 +- plugin/guardrail/approval/approval_test.go | 24 ++- plugin/guardrail/approval/audit.go | 63 ++++++- 6 files changed, 428 insertions(+), 9 deletions(-) diff --git a/platform/toolpolicy/policy.go b/platform/toolpolicy/policy.go index 467a40747f..8644e19a24 100644 --- a/platform/toolpolicy/policy.go +++ b/platform/toolpolicy/policy.go @@ -198,6 +198,11 @@ func (p *Policy) CheckToolPermission( name = strings.TrimSpace(req.Declaration.Name) } decision, reason, audit := p.decide(req, name) + if decision.Action == tool.PermissionActionAsk && approvedToolCall(ctx, req) { + decision = tool.AllowPermission() + reason = "" + audit = false + } if audit { summary, err := p.ApprovalSummary(req, decision, reason) if err != nil { @@ -210,6 +215,36 @@ func (p *Policy) CheckToolPermission( return decision, nil } +func approvedToolCall(ctx context.Context, req *tool.PermissionRequest) bool { + if req == nil || strings.TrimSpace(req.ToolCallID) == "" { + return false + } + fingerprint, ok := approval.ApprovedToolCallFromContext(ctx) + if !ok { + return false + } + if fingerprint.ToolCallID != strings.TrimSpace(req.ToolCallID) { + return false + } + name := strings.TrimSpace(req.ToolName) + if name == "" && req.Declaration != nil { + name = strings.TrimSpace(req.Declaration.Name) + } + if fingerprint.ToolName != name { + return false + } + return fingerprint.ArgumentsHash == argumentsHash(req.Arguments) && + fingerprint.ArgumentsBytes == len(req.Arguments) +} + +func argumentsHash(args []byte) string { + if len(args) == 0 { + return "" + } + sum := sha256.Sum256(args) + return hex.EncodeToString(sum[:]) +} + func (p *Policy) beforeTool() tool.BeforeToolCallbackStructured { return func(ctx context.Context, args *tool.BeforeToolArgs) (*tool.BeforeToolResult, error) { if args == nil { diff --git a/platform/worker/builder.go b/platform/worker/builder.go index daf23d4a2a..3501e021e7 100644 --- a/platform/worker/builder.go +++ b/platform/worker/builder.go @@ -21,6 +21,9 @@ import ( "trpc.group/trpc-go/trpc-agent-go/platform" "trpc.group/trpc-go/trpc-agent-go/platform/gateway" "trpc.group/trpc-go/trpc-agent-go/platform/storagerouter" + "trpc.group/trpc-go/trpc-agent-go/platform/toolpolicy" + "trpc.group/trpc-go/trpc-agent-go/plugin" + "trpc.group/trpc-go/trpc-agent-go/plugin/guardrail/approval" "trpc.group/trpc-go/trpc-agent-go/runner" "trpc.group/trpc-go/trpc-agent-go/session" "trpc.group/trpc-go/trpc-agent-go/tool" @@ -44,6 +47,8 @@ type AgentDependencies struct { ToolFilter tool.FilterFunc // ToolPermissionPolicy enforces tool-call authorization before execution. ToolPermissionPolicy tool.PermissionPolicy + // Plugins contains runner-scoped plugins assembled by worker governance. + Plugins []plugin.Plugin } // AgentFactory builds an agent for one tenant app runtime. @@ -162,6 +167,10 @@ func (b *RuntimeBuilder) Build( if permissionPolicy != nil { toolFilter = permissionPolicy.ToolFilter() } + plugins, err := buildToolGovernancePlugins(resolvedPolicy, auditSink) + if err != nil { + return gateway.Runtime{}, err + } dependencies := AgentDependencies{ Tenant: tenant, @@ -176,6 +185,7 @@ func (b *RuntimeBuilder) Build( ToolPolicy: resolvedPolicy, ToolFilter: toolFilter, ToolPermissionPolicy: permissionPolicy, + Plugins: plugins, } ag, err := b.factory.BuildAgent(ctx, dependencies) if err != nil { @@ -188,6 +198,15 @@ func (b *RuntimeBuilder) Build( return gateway.Runtime{}, ErrAgentNameMismatch } + runnerOptions := []runner.Option{ + runner.WithSessionService(sessionService), + runner.WithMemoryService(memoryService), + runner.WithArtifactService(artifactService), + } + if len(plugins) > 0 { + runnerOptions = append(runnerOptions, runner.WithPlugins(plugins...)) + } + runtime := gateway.Runtime{ Tenant: tenant, App: app, @@ -195,9 +214,7 @@ func (b *RuntimeBuilder) Build( Runner: runner.NewRunner( storage.Scope().ScopedAppName(app.AppID), ag, - runner.WithSessionService(sessionService), - runner.WithMemoryService(memoryService), - runner.WithArtifactService(artifactService), + runnerOptions..., ), Audit: auditSink, ToolFilter: toolFilter, @@ -250,6 +267,90 @@ func validateRuntimeConfig( return nil } +func buildToolGovernancePlugins( + policy platform.ToolPolicy, + auditSink platform.AuditSink, +) ([]plugin.Plugin, error) { + if strings.TrimSpace(policy.PolicyID) == "" { + return nil, nil + } + opts, approvalRequired := toolApprovalOptions(policy) + if !approvalRequired { + return nil, nil + } + reviewer, err := toolpolicy.NewReviewer(policy) + if err != nil { + return nil, fmt.Errorf("build tool approval reviewer: %w", err) + } + opts = append( + opts, + approval.WithReviewer(reviewer), + approval.WithAuditSink(auditSink), + approval.WithApproverUserID("platform-tool-policy-reviewer"), + ) + approvalPlugin, err := approval.New(opts...) + if err != nil { + return nil, fmt.Errorf("build tool approval plugin: %w", err) + } + return []plugin.Plugin{approvalPlugin}, nil +} + +func toolApprovalOptions(policy platform.ToolPolicy) ([]approval.Option, bool) { + opts := []approval.Option{ + approval.WithDefaultToolPolicy(approval.ToolPolicySkipApproval), + } + if policy.DangerousToolAction != platform.DangerousToolActionAsk { + return opts, false + } + whitelist := normalizedToolNames(policy.ToolWhitelist) + denied := normalizedToolNames(policy.ToolDenylist, policy.PlatformDenylist) + hasWhitelist := len(whitelist) > 0 + approvalRequired := false + for _, name := range normalizedToolNames(policy.HighRiskTools) { + if hasWhitelist && !containsToolName(whitelist, name) { + continue + } + if containsToolName(denied, name) { + continue + } + opts = append(opts, approval.WithToolPolicy(name, approval.ToolPolicyRequireApproval)) + approvalRequired = true + } + return opts, approvalRequired +} + +func normalizedToolNames(lists ...[]string) []string { + seen := make(map[string]struct{}) + var names []string + for _, list := range lists { + for _, raw := range list { + name := strings.TrimSpace(raw) + if name == "" { + continue + } + if _, ok := seen[name]; ok { + continue + } + seen[name] = struct{}{} + names = append(names, name) + } + } + return names +} + +func containsToolName(names []string, target string) bool { + target = strings.TrimSpace(target) + if target == "" { + return false + } + for _, name := range names { + if name == target { + return true + } + } + return false +} + func isNilDependency(value any) bool { if value == nil { return true diff --git a/platform/worker/governance_test.go b/platform/worker/governance_test.go index 71114d874b..ff2bd0e784 100644 --- a/platform/worker/governance_test.go +++ b/platform/worker/governance_test.go @@ -10,6 +10,8 @@ package worker import ( "context" + "encoding/json" + "reflect" "strings" "testing" "time" @@ -25,6 +27,8 @@ import ( "trpc.group/trpc-go/trpc-agent-go/platform/artifactstore" "trpc.group/trpc-go/trpc-agent-go/platform/gateway" "trpc.group/trpc-go/trpc-agent-go/platform/storagerouter" + "trpc.group/trpc-go/trpc-agent-go/plugin" + "trpc.group/trpc-go/trpc-agent-go/plugin/guardrail/approval" sessioninmemory "trpc.group/trpc-go/trpc-agent-go/session/inmemory" "trpc.group/trpc-go/trpc-agent-go/tool" ) @@ -102,6 +106,122 @@ func TestRuntimeBuilderAppliesToolGovernanceEndToEnd(t *testing.T) { assert.Equal(t, "completed", records[1].Decision) } +func TestRuntimeBuilderRunsApprovalPluginBeforeMandatoryPolicy(t *testing.T) { + ctx := context.Background() + router, auditSink := governanceTestRouter(t) + tenant, app, binding := governanceRuntimeConfig() + app.ToolPolicyID = "policy-a" + policy := platform.ToolPolicy{ + TenantID: tenant.TenantID, + AppID: app.AppID, + PolicyID: app.ToolPolicyID, + ToolWhitelist: []string{"read_file", "shell"}, + HighRiskTools: []string{"shell"}, + DangerousToolAction: platform.DangerousToolActionAsk, + } + var captured AgentDependencies + builder, err := NewRuntimeBuilder( + router, + AgentFactoryFunc(func( + _ context.Context, + dependencies AgentDependencies, + ) (agent.Agent, error) { + captured = dependencies + return newToolCallGovernanceAgent(app.AgentName), nil + }), + WithToolPolicyProvider(ToolPolicyProviderFunc(func( + context.Context, + string, + string, + string, + ) (platform.ToolPolicy, error) { + return policy, nil + })), + ) + require.NoError(t, err) + runtime, err := builder.Build(ctx, tenant, app, binding) + require.NoError(t, err) + t.Cleanup(func() { + require.NoError(t, runtime.Runner.Close()) + }) + + require.NotEmpty(t, captured.Plugins) + runnerValue := reflect.Indirect(reflect.ValueOf(runtime.Runner)) + pluginManagerField := runnerValue.FieldByName("pluginManager") + require.True(t, pluginManagerField.IsValid()) + require.False(t, pluginManagerField.IsNil()) + manager, err := plugin.NewManager(captured.Plugins...) + require.NoError(t, err) + require.NotNil(t, manager.ToolCallbacks()) + + approvedCtx := approval.ContextWithAuditContext(ctx, approval.AuditContext{ + TenantID: tenant.TenantID, + AppID: app.AppID, + RequestID: "request-1", + TraceID: "trace-1", + }) + result, err := manager.ToolCallbacks().RunBeforeTool( + approvedCtx, + &tool.BeforeToolArgs{ + ToolName: "shell", + ToolCallID: "call-shell", + Arguments: []byte(`{"command":"restricted"}`), + }, + ) + require.NoError(t, err) + require.NotNil(t, result) + require.NotNil(t, result.Context) + + decision, err := runtime.ToolPermissionPolicy.CheckToolPermission( + result.Context, + &tool.PermissionRequest{ + ToolName: "shell", + ToolCallID: "call-shell", + Arguments: []byte(`{"command":"restricted"}`), + }, + ) + require.NoError(t, err) + assert.Equal(t, tool.PermissionActionAllow, decision.Action) + + decision, err = runtime.ToolPermissionPolicy.CheckToolPermission( + result.Context, + &tool.PermissionRequest{ + ToolName: "shell", + ToolCallID: "call-shell", + Arguments: []byte(`{"command":"mutated-after-approval"}`), + }, + ) + require.NoError(t, err) + assert.Equal(t, tool.PermissionActionAsk, decision.Action) + + decision, err = runtime.ToolPermissionPolicy.CheckToolPermission( + ctx, + &tool.PermissionRequest{ + ToolName: "shell", + ToolCallID: "call-shell-unapproved", + Arguments: []byte(`{"command":"restricted"}`), + }, + ) + require.NoError(t, err) + assert.Equal(t, tool.PermissionActionAsk, decision.Action) + + records := auditSink.Records() + require.Len(t, records, 4) + assert.Equal(t, "approval_requested", records[0].Decision) + assert.Equal(t, "approval_approved", records[1].Decision) + assert.Equal(t, string(tool.PermissionActionAsk), records[2].Decision) + assert.Equal(t, string(tool.PermissionActionAsk), records[3].Decision) + assert.Contains(t, records[2].RedactedDetailRef, "tool_call_id:call-shell") + assert.Contains(t, records[3].RedactedDetailRef, "tool_call_id:call-shell-unapproved") + for _, record := range records[:2] { + assert.Equal(t, tenant.TenantID, record.TenantID) + assert.Equal(t, app.AppID, record.AppID) + assert.Equal(t, "shell", record.ToolName) + assert.Contains(t, record.RedactedDetailRef, "tool_call_id:call-shell") + assert.Contains(t, record.RedactedDetailRef, "args_ref_sha256:") + } +} + func TestRuntimeBuilderRejectsMissingToolPolicyProviderBeforeStorage(t *testing.T) { tenant, app, binding := governanceRuntimeConfig() app.ToolPolicyID = "policy-a" @@ -418,3 +538,84 @@ type governanceProbeTool struct { func (t *governanceProbeTool) Declaration() *tool.Declaration { return &tool.Declaration{Name: t.name} } + +func (t *governanceProbeTool) Call( + context.Context, + json.RawMessage, +) (any, error) { + return "executed:" + t.name, nil +} + +type toolCallGovernanceAgent struct { + *governanceProbeAgent +} + +func newToolCallGovernanceAgent(name string) *toolCallGovernanceAgent { + return &toolCallGovernanceAgent{ + governanceProbeAgent: newGovernanceProbeAgent(name), + } +} + +func (a *toolCallGovernanceAgent) Run( + _ context.Context, + invocation *agent.Invocation, +) (<-chan *event.Event, error) { + out := make(chan *event.Event, 1) + if invocation.Session != nil && len(invocation.Session.Events) > 0 { + lastEvent := invocation.Session.Events[len(invocation.Session.Events)-1] + if lastEvent.Response != nil && len(lastEvent.Response.Choices) > 0 { + last := lastEvent.Response.Choices[0].Message + if last.Role == model.RoleTool && last.ToolID == "call-shell" { + out <- event.NewResponseEvent( + invocation.InvocationID, + a.name, + &model.Response{ + ID: "governance-final-response", + Object: model.ObjectTypeChatCompletion, + Done: true, + Choices: []model.Choice{ + { + Index: 0, + Message: model.Message{ + Role: model.RoleAssistant, + Content: last.Content, + }, + }, + }, + }, + ) + close(out) + return out, nil + } + } + } + out <- event.NewResponseEvent( + invocation.InvocationID, + a.name, + &model.Response{ + ID: "governance-tool-call-response", + Object: model.ObjectTypeChatCompletion, + Done: true, + Choices: []model.Choice{ + { + Index: 0, + Message: model.Message{ + Role: model.RoleAssistant, + ToolCalls: []model.ToolCall{ + { + ID: "call-shell", + Type: "function", + Function: model.FunctionDefinitionParam{ + Name: "shell", + Arguments: []byte(`{"command":"restricted"}`), + }, + }, + }, + }, + }, + }, + }, + ) + close(out) + return out, nil +} diff --git a/plugin/guardrail/approval/approval.go b/plugin/guardrail/approval/approval.go index dbf231b7de..1c369b89fd 100644 --- a/plugin/guardrail/approval/approval.go +++ b/plugin/guardrail/approval/approval.go @@ -173,7 +173,12 @@ func (p *Plugin) beforeTool() tool.BeforeToolCallbackStructured { riskLevel, reason, ) - return nil, nil + if strings.TrimSpace(args.ToolCallID) == "" { + return nil, nil + } + return &tool.BeforeToolResult{ + Context: contextWithApprovedToolCall(ctx, args), + }, nil } denyMessage := fmt.Sprintf( "Automatic approval review denied (risk: %s): %s", diff --git a/plugin/guardrail/approval/approval_test.go b/plugin/guardrail/approval/approval_test.go index 123b61bb10..d229e01d1f 100644 --- a/plugin/guardrail/approval/approval_test.go +++ b/plugin/guardrail/approval/approval_test.go @@ -223,7 +223,13 @@ func TestBeforeTool_RequireApprovalBuildsRequestFromSession(t *testing.T) { Arguments: []byte(`{"command":"pwd"}`), }) require.NoError(t, err) - require.Nil(t, result) + require.NotNil(t, result) + approvedToolCall, ok := ApprovedToolCallFromContext(result.Context) + require.True(t, ok) + require.Equal(t, "call-2", approvedToolCall.ToolCallID) + require.Equal(t, "shell", approvedToolCall.ToolName) + require.NotEmpty(t, approvedToolCall.ArgumentsHash) + require.Equal(t, len([]byte(`{"command":"pwd"}`)), approvedToolCall.ArgumentsBytes) require.NotNil(t, captured) require.Equal(t, "shell", captured.Action.ToolName) require.Equal(t, "Runs shell commands.", captured.Action.ToolDescription) @@ -263,7 +269,13 @@ func TestBeforeTool_ReviewerApprovedLogsInfo(t *testing.T) { Arguments: []byte(`{"command":"pwd"}`), }) require.NoError(t, runErr) - require.Nil(t, result) + require.NotNil(t, result) + approvedToolCall, ok := ApprovedToolCallFromContext(result.Context) + require.True(t, ok) + require.Equal(t, "call-1", approvedToolCall.ToolCallID) + require.Equal(t, "shell", approvedToolCall.ToolName) + require.NotEmpty(t, approvedToolCall.ArgumentsHash) + require.Equal(t, len([]byte(`{"command":"pwd"}`)), approvedToolCall.ArgumentsBytes) require.Equal( t, "Automatic approval review approved (risk: medium): The action is scoped and user-authorized.", @@ -304,7 +316,13 @@ func TestBeforeTool_RequireApprovalWritesApprovedAuditRecords(t *testing.T) { Arguments: []byte(`{"command":"git status --short"}`), }) require.NoError(t, runErr) - require.Nil(t, result) + require.NotNil(t, result) + approvedToolCall, ok := ApprovedToolCallFromContext(result.Context) + require.True(t, ok) + require.Equal(t, "call-1", approvedToolCall.ToolCallID) + require.Equal(t, "shell", approvedToolCall.ToolName) + require.NotEmpty(t, approvedToolCall.ArgumentsHash) + require.Equal(t, len([]byte(`{"command":"git status --short"}`)), approvedToolCall.ArgumentsBytes) records := audit.Records() require.Len(t, records, 2) diff --git a/plugin/guardrail/approval/audit.go b/plugin/guardrail/approval/audit.go index dd5ba76ef5..52114766a3 100644 --- a/plugin/guardrail/approval/audit.go +++ b/plugin/guardrail/approval/audit.go @@ -30,6 +30,7 @@ type AuditContext struct { } type auditContextKey struct{} +type approvedToolCallContextKey struct{} // ContextWithAuditContext attaches trusted platform audit context. func ContextWithAuditContext(ctx context.Context, auditCtx AuditContext) context.Context { @@ -39,6 +40,56 @@ func ContextWithAuditContext(ctx context.Context, auditCtx AuditContext) context return context.WithValue(ctx, auditContextKey{}, auditCtx) } +// ApprovedToolCall binds approval to the exact tool call payload reviewed by +// the approval plugin. +type ApprovedToolCall struct { + ToolCallID string + ToolName string + ArgumentsHash string + ArgumentsBytes int +} + +// contextWithApprovedToolCall marks a reviewed tool call as approved so later +// mandatory permission checks do not ask for the same approval again. +func contextWithApprovedToolCall(ctx context.Context, args *tool.BeforeToolArgs) context.Context { + if ctx == nil { + ctx = context.Background() + } + fingerprint := approvedToolCallFingerprint(args) + if fingerprint.ToolCallID == "" { + return ctx + } + return context.WithValue( + ctx, + approvedToolCallContextKey{}, + fingerprint, + ) +} + +// ApprovedToolCallFromContext returns the approved tool call fingerprint, if any. +func ApprovedToolCallFromContext(ctx context.Context) (ApprovedToolCall, bool) { + if ctx == nil { + return ApprovedToolCall{}, false + } + fingerprint, ok := ctx.Value(approvedToolCallContextKey{}).(ApprovedToolCall) + fingerprint.ToolCallID = strings.TrimSpace(fingerprint.ToolCallID) + fingerprint.ToolName = strings.TrimSpace(fingerprint.ToolName) + return fingerprint, ok && fingerprint.ToolCallID != "" +} + +func approvedToolCallFingerprint(args *tool.BeforeToolArgs) ApprovedToolCall { + if args == nil { + return ApprovedToolCall{} + } + hash, bytes := argumentDigest(args.Arguments) + return ApprovedToolCall{ + ToolCallID: strings.TrimSpace(args.ToolCallID), + ToolName: strings.TrimSpace(args.ToolName), + ArgumentsHash: hash, + ArgumentsBytes: bytes, + } +} + func (p *Plugin) writeApprovalAudit( ctx context.Context, args *tool.BeforeToolArgs, @@ -110,11 +161,19 @@ func approvalAuditContextFrom(ctx context.Context) AuditContext { } func argumentSummaryRef(args []byte) string { - if len(args) == 0 { + hash, bytes := argumentDigest(args) + if hash == "" { return "" } + return "args:sha256:" + hash + " args_bytes:" + strconv.Itoa(bytes) +} + +func argumentDigest(args []byte) (string, int) { + if len(args) == 0 { + return "", 0 + } sum := sha256.Sum256(args) - return "args:sha256:" + hex.EncodeToString(sum[:]) + " args_bytes:" + strconv.Itoa(len(args)) + return hex.EncodeToString(sum[:]), len(args) } func approvalAuditDecisionReason(decision platform.ToolApprovalDecision) string { From 2d3787a25802eb45b57d0ba174e1ae453e39ed3d Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 12:40:54 +0800 Subject: [PATCH 85/95] feat(approval): include tool metadata in approval context --- internal/flow/processor/functioncall.go | 16 ++- internal/flow/processor/functioncall_test.go | 64 ++++++++++ platform/toolpolicy/policy.go | 8 +- platform/toolpolicy/policy_test.go | 55 +++++++++ platform/worker/builder.go | 15 ++- platform/worker/governance_test.go | 111 ++++++++++++++++++ plugin/guardrail/approval/approval.go | 34 +++++- plugin/guardrail/approval/approval_test.go | 64 ++++++++++ plugin/guardrail/approval/audit.go | 2 + plugin/guardrail/approval/option.go | 10 ++ plugin/guardrail/approval/review/message.go | 10 ++ plugin/guardrail/approval/review/review.go | 2 + .../guardrail/approval/review/review_test.go | 14 +++ plugin/guardrail/approval/transcript.go | 1 + tool/callbacks.go | 2 + 15 files changed, 400 insertions(+), 8 deletions(-) diff --git a/internal/flow/processor/functioncall.go b/internal/flow/processor/functioncall.go index 12c92ee200..949e4296f5 100644 --- a/internal/flow/processor/functioncall.go +++ b/internal/flow/processor/functioncall.go @@ -2226,6 +2226,7 @@ func (p *FunctionCallResponseProcessor) runBeforeToolPluginCallbacks( invocation *agent.Invocation, toolCall model.ToolCall, toolDeclaration *tool.Declaration, + toolMetadata tool.ToolMetadata, ) (context.Context, model.ToolCall, any, error) { if invocation == nil || invocation.Plugins == nil { return ctx, toolCall, nil, nil @@ -2241,6 +2242,7 @@ func (p *FunctionCallResponseProcessor) runBeforeToolPluginCallbacks( ToolName: toolCall.Function.Name, Declaration: toolDeclaration, Arguments: toolCall.Function.Arguments, + Metadata: toolMetadata, } result, err := callbacks.RunBeforeTool(ctx, args) if err != nil { @@ -2269,6 +2271,7 @@ func (p *FunctionCallResponseProcessor) runBeforeToolCallbacks( ctx context.Context, toolCall model.ToolCall, toolDeclaration *tool.Declaration, + toolMetadata tool.ToolMetadata, ) (context.Context, model.ToolCall, any, error) { if p.toolCallbacks == nil { return ctx, toolCall, nil, nil @@ -2279,6 +2282,7 @@ func (p *FunctionCallResponseProcessor) runBeforeToolCallbacks( ToolName: toolCall.Function.Name, Declaration: toolDeclaration, Arguments: toolCall.Function.Arguments, + Metadata: toolMetadata, } result, err := p.toolCallbacks.RunBeforeTool(ctx, args) if err != nil { @@ -2429,12 +2433,15 @@ func (p *FunctionCallResponseProcessor) executeToolWithCallbacks( } rememberExecutingToolArgs(ctx, toolCall.Function.Arguments) toolDeclaration := tl.Declaration() + semanticTool := itool.ResolveSemantic(tl) + toolMetadata := tool.MetadataOf(semanticTool) visibilityResult, err := checkMandatoryToolVisibility( ctx, invocation, toolCall, tl, toolDeclaration, + toolMetadata, ) if err != nil { return ctx, nil, toolCall.Function.Arguments, false, false, err @@ -2450,6 +2457,7 @@ func (p *FunctionCallResponseProcessor) executeToolWithCallbacks( invocation, toolCall, toolDeclaration, + toolMetadata, ) if err != nil { return ctx, nil, toolCall.Function.Arguments, false, false, err @@ -2464,6 +2472,7 @@ func (p *FunctionCallResponseProcessor) executeToolWithCallbacks( ctx, toolCall, toolDeclaration, + toolMetadata, ) if err != nil { return ctx, nil, toolCall.Function.Arguments, false, false, err @@ -2480,6 +2489,7 @@ func (p *FunctionCallResponseProcessor) executeToolWithCallbacks( toolCall, tl, toolDeclaration, + toolMetadata, ) if err != nil { return ctx, nil, toolCall.Function.Arguments, false, false, err @@ -2564,6 +2574,7 @@ func checkMandatoryToolVisibility( toolCall model.ToolCall, tl tool.Tool, decl *tool.Declaration, + metadata tool.ToolMetadata, ) (*tool.PermissionResult, error) { if invocation == nil || invocation.RunOptions.MandatoryToolFilter == nil { return nil, nil @@ -2580,7 +2591,7 @@ func checkMandatoryToolVisibility( ToolCallID: toolCall.ID, Declaration: decl, Arguments: toolCall.Function.Arguments, - Metadata: tool.MetadataOf(itool.ResolveSemantic(tl)), + Metadata: metadata, } return normalizeToolPermissionResult( req, @@ -2600,6 +2611,7 @@ func (p *FunctionCallResponseProcessor) checkToolPermission( toolCall model.ToolCall, tl tool.Tool, decl *tool.Declaration, + metadata tool.ToolMetadata, ) (*tool.PermissionResult, error) { semanticTool := itool.ResolveSemantic(tl) req := &tool.PermissionRequest{ @@ -2608,7 +2620,7 @@ func (p *FunctionCallResponseProcessor) checkToolPermission( ToolCallID: toolCall.ID, Declaration: decl, Arguments: toolCall.Function.Arguments, - Metadata: tool.MetadataOf(semanticTool), + Metadata: metadata, } if checker, ok := semanticTool.(tool.PermissionChecker); ok { decision, err := checker.CheckPermission(ctx, req) diff --git a/internal/flow/processor/functioncall_test.go b/internal/flow/processor/functioncall_test.go index 0e1ca08a88..8b5f60a098 100644 --- a/internal/flow/processor/functioncall_test.go +++ b/internal/flow/processor/functioncall_test.go @@ -9401,6 +9401,70 @@ func TestExecuteToolWithCallbacks_ToolPermissionReceivesMetadata( require.JSONEq(t, `{"ok":true}`, string(mustJSON(res))) } +func TestExecuteToolWithCallbacks_BeforeToolReceivesMetadata( + t *testing.T, +) { + const toolName = "web_search" + metadata := tool.ToolMetadata{ + ReadOnly: true, + SearchOrRead: true, + OpenWorld: true, + } + var pluginSawMetadata bool + approvalPlugin := &hookPlugin{ + name: "metadata-before-tool", + reg: func(r *plugin.Registry) { + r.BeforeTool(func(_ context.Context, args *tool.BeforeToolArgs) (*tool.BeforeToolResult, error) { + require.Equal(t, metadata, args.Metadata) + pluginSawMetadata = true + return nil, nil + }) + }, + } + var localSawMetadata bool + callbacks := tool.NewCallbacks() + callbacks.RegisterBeforeTool(func(_ context.Context, args *tool.BeforeToolArgs) (*tool.BeforeToolResult, error) { + require.Equal(t, metadata, args.Metadata) + localSawMetadata = true + return nil, nil + }) + tl := &permissionMockTool{ + mockCallableTool: &mockCallableTool{ + declaration: &tool.Declaration{Name: toolName}, + callFn: func(_ context.Context, _ []byte) (any, error) { + return map[string]any{"ok": true}, nil + }, + }, + metadata: metadata, + decision: tool.AllowPermission(), + } + manager, err := plugin.NewManager(approvalPlugin) + require.NoError(t, err) + inv := &agent.Invocation{ + Plugins: manager, + RunOptions: agent.NewRunOptions(), + } + + _, res, _, _, _, err := NewFunctionCallResponseProcessor(false, callbacks). + executeToolWithCallbacks( + context.Background(), + inv, + model.ToolCall{ + ID: "call-allow", + Function: model.FunctionDefinitionParam{ + Name: toolName, + Arguments: []byte(`{}`), + }, + }, + tl, + nil, + ) + require.NoError(t, err) + require.True(t, pluginSawMetadata) + require.True(t, localSawMetadata) + require.JSONEq(t, `{"ok":true}`, string(mustJSON(res))) +} + func TestExecuteToolCall_StreamableFinalStateOnlyResultAfterToolContextReplacementStillSkipsDefaultMessage(t *testing.T) { ctx := context.Background() callbacks := tool.NewCallbacks() diff --git a/platform/toolpolicy/policy.go b/platform/toolpolicy/policy.go index 8644e19a24..7e3fe6ea42 100644 --- a/platform/toolpolicy/policy.go +++ b/platform/toolpolicy/policy.go @@ -234,7 +234,8 @@ func approvedToolCall(ctx context.Context, req *tool.PermissionRequest) bool { return false } return fingerprint.ArgumentsHash == argumentsHash(req.Arguments) && - fingerprint.ArgumentsBytes == len(req.Arguments) + fingerprint.ArgumentsBytes == len(req.Arguments) && + fingerprint.Metadata == req.Metadata } func argumentsHash(args []byte) string { @@ -255,6 +256,7 @@ func (p *Policy) beforeTool() tool.BeforeToolCallbackStructured { ToolCallID: args.ToolCallID, Declaration: args.Declaration, Arguments: args.Arguments, + Metadata: args.Metadata, } decision, reason, audit := p.decideNameOnly(req, req.ToolName) if audit { @@ -293,6 +295,9 @@ func ApprovalOptions(policy platform.ToolPolicy) ([]approval.Option, error) { defaultPolicy = approval.ToolPolicyDenied } opts := []approval.Option{approval.WithDefaultToolPolicy(defaultPolicy)} + if policy.DangerousToolAction == platform.DangerousToolActionAsk { + opts = append(opts, approval.WithMetadataRiskPolicy(approval.ToolPolicyRequireApproval)) + } whitelist := normalizedList(policy.ToolWhitelist) hasWhitelist := len(whitelist) > 0 for _, name := range whitelist { @@ -342,6 +347,7 @@ func (r *Reviewer) Review(ctx context.Context, req *review.Request) (*review.Dec ToolName: req.Action.ToolName, Declaration: &tool.Declaration{Name: req.Action.ToolName, Description: req.Action.ToolDescription}, Arguments: req.Action.Arguments, + Metadata: req.Action.Metadata, } decision, reason, audit := r.policy.decideReviewer(permissionReq, permissionReq.ToolName) if audit { diff --git a/platform/toolpolicy/policy_test.go b/platform/toolpolicy/policy_test.go index 72f4d1ca3c..aed7adbc32 100644 --- a/platform/toolpolicy/policy_test.go +++ b/platform/toolpolicy/policy_test.go @@ -709,6 +709,39 @@ func TestPolicyRegisterDoesNotTreatUnknownMetadataAsHighRisk(t *testing.T) { } } +func TestPolicyRegisterTreatsMetadataRiskAsHighRisk(t *testing.T) { + audit := platform.NewInMemoryAuditSink() + p := newPolicy(t, platform.ToolPolicy{ + TenantID: "tenant", + AppID: "app", + DangerousToolAction: platform.DangerousToolActionAsk, + }, WithAuditSink(audit)) + manager := plugin.MustNewManager(p) + callbacks := manager.ToolCallbacks() + + result, err := callbacks.RunBeforeTool(context.Background(), &tool.BeforeToolArgs{ + ToolName: "metadata_shell", + Arguments: []byte(`{"command":"pwd"}`), + Metadata: tool.ToolMetadata{ReadOnly: false, OpenWorld: true}, + }) + if err != nil { + t.Fatalf("RunBeforeTool: %v", err) + } + if result == nil || result.CustomResult == nil { + t.Fatalf("expected metadata high-risk approval-required result") + } + permissionResult, ok := result.CustomResult.(tool.PermissionResult) + if !ok { + t.Fatalf("expected tool.PermissionResult, got %T", result.CustomResult) + } + if permissionResult.Status != tool.PermissionResultStatusApprovalRequired { + t.Fatalf("expected approval_required, got %+v", permissionResult) + } + if len(audit.Records()) != 1 || audit.Records()[0].Decision != string(tool.PermissionActionAsk) { + t.Fatalf("expected ask audit record, got %+v", audit.Records()) + } +} + func TestReviewerMapsPolicyDecisionToApprovalDecision(t *testing.T) { reviewer, err := NewReviewer(defaultPolicy(platform.ToolPolicy{ ToolWhitelist: []string{"search"}, @@ -748,6 +781,28 @@ func TestReviewerApprovesAskDecisionForApprovalPluginFlow(t *testing.T) { } } +func TestReviewerApprovesMetadataAskDecisionForApprovalPluginFlow(t *testing.T) { + reviewer, err := NewReviewer(defaultPolicy(platform.ToolPolicy{ + DangerousToolAction: platform.DangerousToolActionAsk, + })) + if err != nil { + t.Fatalf("NewReviewer: %v", err) + } + + decision, err := reviewer.Review(context.Background(), &review.Request{ + Action: review.Action{ + ToolName: "metadata_shell", + Metadata: tool.ToolMetadata{ReadOnly: false, OpenWorld: true}, + }, + }) + if err != nil { + t.Fatalf("Review: %v", err) + } + if !decision.Approved { + t.Fatalf("expected metadata ask decision to be approved inside approval flow, got %+v", decision) + } +} + func TestApprovalOptionsWithReviewerAllowsWhitelistedHighRiskAsk(t *testing.T) { policy := defaultPolicy(platform.ToolPolicy{ ToolWhitelist: []string{"workspace_write"}, diff --git a/platform/worker/builder.go b/platform/worker/builder.go index 3501e021e7..13aaf3d59e 100644 --- a/platform/worker/builder.go +++ b/platform/worker/builder.go @@ -296,16 +296,24 @@ func buildToolGovernancePlugins( } func toolApprovalOptions(policy platform.ToolPolicy) ([]approval.Option, bool) { + defaultPolicy := approval.ToolPolicySkipApproval + if len(normalizedToolNames(policy.ToolWhitelist)) > 0 { + defaultPolicy = approval.ToolPolicyDenied + } opts := []approval.Option{ - approval.WithDefaultToolPolicy(approval.ToolPolicySkipApproval), + approval.WithDefaultToolPolicy(defaultPolicy), } if policy.DangerousToolAction != platform.DangerousToolActionAsk { return opts, false } + opts = append(opts, approval.WithMetadataRiskPolicy(approval.ToolPolicyRequireApproval)) whitelist := normalizedToolNames(policy.ToolWhitelist) denied := normalizedToolNames(policy.ToolDenylist, policy.PlatformDenylist) hasWhitelist := len(whitelist) > 0 - approvalRequired := false + approvalRequired := true + for _, name := range whitelist { + opts = append(opts, approval.WithToolPolicy(name, approval.ToolPolicySkipApproval)) + } for _, name := range normalizedToolNames(policy.HighRiskTools) { if hasWhitelist && !containsToolName(whitelist, name) { continue @@ -316,6 +324,9 @@ func toolApprovalOptions(policy platform.ToolPolicy) ([]approval.Option, bool) { opts = append(opts, approval.WithToolPolicy(name, approval.ToolPolicyRequireApproval)) approvalRequired = true } + for _, name := range denied { + opts = append(opts, approval.WithToolPolicy(name, approval.ToolPolicyDenied)) + } return opts, approvalRequired } diff --git a/platform/worker/governance_test.go b/platform/worker/governance_test.go index ff2bd0e784..1d31991aab 100644 --- a/platform/worker/governance_test.go +++ b/platform/worker/governance_test.go @@ -222,6 +222,117 @@ func TestRuntimeBuilderRunsApprovalPluginBeforeMandatoryPolicy(t *testing.T) { } } +func TestRuntimeBuilderRunsApprovalPluginForMetadataRisk(t *testing.T) { + ctx := context.Background() + router, auditSink := governanceTestRouter(t) + tenant, app, binding := governanceRuntimeConfig() + app.ToolPolicyID = "policy-a" + policy := platform.ToolPolicy{ + TenantID: tenant.TenantID, + AppID: app.AppID, + PolicyID: app.ToolPolicyID, + ToolWhitelist: []string{"read_file", "metadata_shell"}, + DangerousToolAction: platform.DangerousToolActionAsk, + } + var captured AgentDependencies + builder, err := NewRuntimeBuilder( + router, + AgentFactoryFunc(func( + _ context.Context, + dependencies AgentDependencies, + ) (agent.Agent, error) { + captured = dependencies + return newGovernanceProbeAgent(app.AgentName), nil + }), + WithToolPolicyProvider(ToolPolicyProviderFunc(func( + context.Context, + string, + string, + string, + ) (platform.ToolPolicy, error) { + return policy, nil + })), + ) + require.NoError(t, err) + runtime, err := builder.Build(ctx, tenant, app, binding) + require.NoError(t, err) + t.Cleanup(func() { + require.NoError(t, runtime.Runner.Close()) + }) + + manager, err := plugin.NewManager(captured.Plugins...) + require.NoError(t, err) + require.NotNil(t, manager.ToolCallbacks()) + metadata := tool.ToolMetadata{ReadOnly: false, OpenWorld: true} + deniedResult, err := manager.ToolCallbacks().RunBeforeTool( + ctx, + &tool.BeforeToolArgs{ + ToolName: "outside_metadata", + ToolCallID: "call-outside", + Arguments: []byte(`{"command":"restricted"}`), + Metadata: metadata, + }, + ) + require.NoError(t, err) + require.NotNil(t, deniedResult) + assert.Equal( + t, + `tool "outside_metadata" is denied by approval policy`, + deniedResult.CustomResult, + ) + + result, err := manager.ToolCallbacks().RunBeforeTool( + approval.ContextWithAuditContext(ctx, approval.AuditContext{ + TenantID: tenant.TenantID, + AppID: app.AppID, + RequestID: "request-metadata", + TraceID: "trace-metadata", + }), + &tool.BeforeToolArgs{ + ToolName: "metadata_shell", + ToolCallID: "call-metadata", + Arguments: []byte(`{"command":"restricted"}`), + Metadata: metadata, + }, + ) + require.NoError(t, err) + require.NotNil(t, result) + require.NotNil(t, result.Context) + + decision, err := runtime.ToolPermissionPolicy.CheckToolPermission( + result.Context, + &tool.PermissionRequest{ + ToolName: "metadata_shell", + ToolCallID: "call-metadata", + Arguments: []byte(`{"command":"restricted"}`), + Metadata: metadata, + }, + ) + require.NoError(t, err) + assert.Equal(t, tool.PermissionActionAllow, decision.Action) + + decision, err = runtime.ToolPermissionPolicy.CheckToolPermission( + result.Context, + &tool.PermissionRequest{ + ToolName: "metadata_shell", + ToolCallID: "call-metadata", + Arguments: []byte(`{"command":"restricted"}`), + Metadata: tool.ToolMetadata{ReadOnly: true, OpenWorld: true}, + }, + ) + require.NoError(t, err) + assert.Equal(t, tool.PermissionActionAsk, decision.Action) + + records := auditSink.Records() + require.Len(t, records, 3) + assert.Equal(t, "approval_requested", records[0].Decision) + assert.Equal(t, "approval_approved", records[1].Decision) + assert.Equal(t, string(tool.PermissionActionAsk), records[2].Decision) + assert.Equal(t, "metadata_shell", records[0].ToolName) + assert.Equal(t, "metadata_shell", records[1].ToolName) + assert.Equal(t, "metadata_shell", records[2].ToolName) +} + func TestRuntimeBuilderRejectsMissingToolPolicyProviderBeforeStorage(t *testing.T) { tenant, app, binding := governanceRuntimeConfig() app.ToolPolicyID = "policy-a" diff --git a/plugin/guardrail/approval/approval.go b/plugin/guardrail/approval/approval.go index 1c369b89fd..1c82db1c8f 100644 --- a/plugin/guardrail/approval/approval.go +++ b/plugin/guardrail/approval/approval.go @@ -29,6 +29,7 @@ type Plugin struct { reviewer review.Reviewer defaultToolPolicy ToolPolicy toolPolicies map[string]ToolPolicy + metadataPolicy ToolPolicy tokenCounter model.TokenCounter auditSink platform.AuditSink approverUserID string @@ -41,6 +42,9 @@ func New(options ...Option) (*Plugin, error) { if err := validateToolPolicy(opts.defaultToolPolicy); err != nil { return nil, fmt.Errorf("newing approval plugin: default tool policy: %w", err) } + if err := validateToolPolicy(opts.metadataPolicy); err != nil { + return nil, fmt.Errorf("newing approval plugin: metadata policy: %w", err) + } for toolName, policy := range opts.toolPolicies { if toolName == "" { return nil, fmt.Errorf("newing approval plugin: tool policy name is empty") @@ -60,6 +64,7 @@ func New(options ...Option) (*Plugin, error) { reviewer: opts.reviewer, defaultToolPolicy: opts.defaultToolPolicy, toolPolicies: opts.toolPolicies, + metadataPolicy: opts.metadataPolicy, tokenCounter: model.NewSimpleTokenCounter(), auditSink: opts.auditSink, approverUserID: strings.TrimSpace(opts.approverUserID), @@ -85,7 +90,7 @@ func (p *Plugin) beforeTool() tool.BeforeToolCallbackStructured { if args == nil { return nil, nil } - policy := p.resolveToolPolicy(args.ToolName) + policy := p.resolvePolicy(args) switch policy { case ToolPolicyDenied: return &tool.BeforeToolResult{ @@ -214,8 +219,21 @@ func (p *Plugin) beforeTool() tool.BeforeToolCallbackStructured { } } -func (p *Plugin) resolveToolPolicy(toolName string) ToolPolicy { - if policy, ok := p.toolPolicies[toolName]; ok { +func (p *Plugin) resolvePolicy(args *tool.BeforeToolArgs) ToolPolicy { + if args == nil { + return p.defaultToolPolicy + } + policy, explicit := p.toolPolicies[args.ToolName] + if explicit && policy != ToolPolicySkipApproval { + return policy + } + if p.defaultToolPolicy == ToolPolicyDenied && !explicit { + return p.defaultToolPolicy + } + if metadataHighRisk(args.Metadata) { + return p.metadataPolicy + } + if explicit { return policy } return p.defaultToolPolicy @@ -225,6 +243,9 @@ func requiresReviewer(opts *options) bool { if opts.defaultToolPolicy == ToolPolicyRequireApproval { return true } + if opts.metadataPolicy == ToolPolicyRequireApproval { + return true + } for _, policy := range opts.toolPolicies { if policy == ToolPolicyRequireApproval { return true @@ -232,3 +253,10 @@ func requiresReviewer(opts *options) bool { } return false } + +func metadataHighRisk(metadata tool.ToolMetadata) bool { + if metadata == (tool.ToolMetadata{}) { + return false + } + return metadata.Destructive || !metadata.ReadOnly || metadata.OpenWorld +} diff --git a/plugin/guardrail/approval/approval_test.go b/plugin/guardrail/approval/approval_test.go index d229e01d1f..8fe786d9f0 100644 --- a/plugin/guardrail/approval/approval_test.go +++ b/plugin/guardrail/approval/approval_test.go @@ -62,12 +62,31 @@ func TestNew_RequiresReviewerWhenExplicitToolPolicyRequiresApproval(t *testing.T require.Contains(t, err.Error(), "reviewer is nil") } +func TestNew_RequiresReviewerWhenMetadataPolicyRequiresApproval(t *testing.T) { + _, err := New( + WithDefaultToolPolicy(ToolPolicySkipApproval), + WithMetadataRiskPolicy(ToolPolicyRequireApproval), + ) + require.Error(t, err) + require.Contains(t, err.Error(), "reviewer is nil") +} + func TestNew_InvalidToolPolicy(t *testing.T) { _, err := New(WithReviewer(&stubReviewer{}), WithDefaultToolPolicy(ToolPolicy("bad"))) require.Error(t, err) require.Contains(t, err.Error(), "invalid tool policy") } +func TestNew_InvalidMetadataPolicy(t *testing.T) { + _, err := New( + WithReviewer(&stubReviewer{}), + WithMetadataRiskPolicy(ToolPolicy("bad")), + ) + require.Error(t, err) + require.Contains(t, err.Error(), "metadata policy") + require.Contains(t, err.Error(), "invalid tool policy") +} + func TestNew_EmptyToolPolicyName(t *testing.T) { _, err := New(WithReviewer(&stubReviewer{}), WithToolPolicy("", ToolPolicyDenied)) require.Error(t, err) @@ -95,9 +114,11 @@ func TestOptionSettersUpdateOptions(t *testing.T) { WithName("tool-approval")(opts) WithReviewer(reviewer)(opts) WithDefaultToolPolicy(ToolPolicyDenied)(opts) + WithMetadataRiskPolicy(ToolPolicyRequireApproval)(opts) require.Equal(t, "tool-approval", opts.name) require.Equal(t, reviewer, opts.reviewer) require.Equal(t, ToolPolicyDenied, opts.defaultToolPolicy) + require.Equal(t, ToolPolicyRequireApproval, opts.metadataPolicy) } func TestRegister_IgnoresNilReceiverAndNilRegistry(t *testing.T) { @@ -283,6 +304,46 @@ func TestBeforeTool_ReviewerApprovedLogsInfo(t *testing.T) { ) } +func TestBeforeTool_MetadataRiskRequiresApproval(t *testing.T) { + metadata := tool.ToolMetadata{ + ReadOnly: false, + OpenWorld: true, + } + var captured *approvalreview.Request + p, err := New( + WithDefaultToolPolicy(ToolPolicySkipApproval), + WithMetadataRiskPolicy(ToolPolicyRequireApproval), + WithReviewer(&stubReviewer{ + reviewFn: func(ctx context.Context, req *approvalreview.Request) (*approvalreview.Decision, error) { + captured = req + return &approvalreview.Decision{ + Approved: true, + RiskScore: 30, + RiskLevel: "medium", + Reason: "Metadata risk is reviewed.", + }, nil + }, + }), + ) + require.NoError(t, err) + callbacks := registeredToolCallbacks(t, p) + result, runErr := callbacks.RunBeforeTool(context.Background(), &tool.BeforeToolArgs{ + ToolName: "metadata_shell", + ToolCallID: "call-metadata", + Arguments: []byte(`{"command":"pwd"}`), + Metadata: metadata, + }) + require.NoError(t, runErr) + require.NotNil(t, result) + require.NotNil(t, result.Context) + + approvedToolCall, ok := ApprovedToolCallFromContext(result.Context) + require.True(t, ok) + require.Equal(t, metadata, approvedToolCall.Metadata) + require.NotNil(t, captured) + require.Equal(t, metadata, captured.Action.Metadata) +} + func TestBeforeTool_RequireApprovalWritesApprovedAuditRecords(t *testing.T) { now := time.Date(2026, 7, 11, 10, 30, 0, 0, time.UTC) audit := platform.NewInMemoryAuditSink() @@ -563,18 +624,21 @@ func TestBuildTranscript_UserOverflowReturnsOmissionOnly(t *testing.T) { } func TestBuildRequest_WithoutInvocationReturnsActionOnly(t *testing.T) { + metadata := tool.ToolMetadata{ReadOnly: false, OpenWorld: true} p, err := New(WithDefaultToolPolicy(ToolPolicyDenied)) require.NoError(t, err) req, buildErr := p.buildRequest(context.Background(), &tool.BeforeToolArgs{ ToolName: "shell", Declaration: &tool.Declaration{Description: "Runs shell commands."}, Arguments: []byte(`{"command":"pwd"}`), + Metadata: metadata, }) require.NoError(t, buildErr) require.NotNil(t, req) require.Equal(t, "shell", req.Action.ToolName) require.Equal(t, "Runs shell commands.", req.Action.ToolDescription) require.JSONEq(t, `{"command":"pwd"}`, string(req.Action.Arguments)) + require.Equal(t, metadata, req.Action.Metadata) require.Nil(t, req.Transcript) } diff --git a/plugin/guardrail/approval/audit.go b/plugin/guardrail/approval/audit.go index 52114766a3..f4f1dd04e6 100644 --- a/plugin/guardrail/approval/audit.go +++ b/plugin/guardrail/approval/audit.go @@ -47,6 +47,7 @@ type ApprovedToolCall struct { ToolName string ArgumentsHash string ArgumentsBytes int + Metadata tool.ToolMetadata } // contextWithApprovedToolCall marks a reviewed tool call as approved so later @@ -87,6 +88,7 @@ func approvedToolCallFingerprint(args *tool.BeforeToolArgs) ApprovedToolCall { ToolName: strings.TrimSpace(args.ToolName), ArgumentsHash: hash, ArgumentsBytes: bytes, + Metadata: args.Metadata, } } diff --git a/plugin/guardrail/approval/option.go b/plugin/guardrail/approval/option.go index 1ef7bf7b8b..740aa1ae89 100644 --- a/plugin/guardrail/approval/option.go +++ b/plugin/guardrail/approval/option.go @@ -25,6 +25,7 @@ type options struct { reviewer review.Reviewer defaultToolPolicy ToolPolicy toolPolicies map[string]ToolPolicy + metadataPolicy ToolPolicy auditSink platform.AuditSink approverUserID string now func() time.Time @@ -35,6 +36,7 @@ func newOptions(opts ...Option) *options { name: defaultPluginName, defaultToolPolicy: ToolPolicyRequireApproval, toolPolicies: make(map[string]ToolPolicy), + metadataPolicy: ToolPolicySkipApproval, } for _, opt := range opts { if opt != nil { @@ -75,6 +77,14 @@ func WithToolPolicy(name string, policy ToolPolicy) Option { } } +// WithMetadataRiskPolicy sets the policy for calls whose tool metadata is +// high-risk when no explicit tool policy exists. +func WithMetadataRiskPolicy(policy ToolPolicy) Option { + return func(opts *options) { + opts.metadataPolicy = policy + } +} + // WithAuditSink records approval request and decision boundaries to audit. func WithAuditSink(sink platform.AuditSink) Option { return func(opts *options) { diff --git a/plugin/guardrail/approval/review/message.go b/plugin/guardrail/approval/review/message.go index a15087f207..4e9456bd97 100644 --- a/plugin/guardrail/approval/review/message.go +++ b/plugin/guardrail/approval/review/message.go @@ -15,6 +15,7 @@ import ( "text/template" "trpc.group/trpc-go/trpc-agent-go/model" + "trpc.group/trpc-go/trpc-agent-go/tool" ) const defaultSystemPromptTemplateText = `You are the guardian reviewer for tool approval decisions. @@ -58,6 +59,7 @@ type actionPayload struct { ToolName string `json:"tool_name"` ToolDescription string `json:"tool_description,omitempty"` Arguments any `json:"arguments"` + Metadata any `json:"metadata,omitempty"` } type systemPromptTemplateData struct { @@ -125,6 +127,7 @@ func marshalActionPayload(action Action) ([]byte, error) { ToolName: action.ToolName, ToolDescription: action.ToolDescription, Arguments: actionArgumentsForJSON(action.Arguments), + Metadata: actionMetadataForJSON(action.Metadata), } data, err := json.MarshalIndent(payload, "", " ") if err != nil { @@ -133,6 +136,13 @@ func marshalActionPayload(action Action) ([]byte, error) { return data, nil } +func actionMetadataForJSON(metadata tool.ToolMetadata) any { + if metadata == (tool.ToolMetadata{}) { + return nil + } + return metadata +} + func actionArgumentsForJSON(arguments json.RawMessage) any { if len(arguments) == 0 { return json.RawMessage(`{}`) diff --git a/plugin/guardrail/approval/review/review.go b/plugin/guardrail/approval/review/review.go index bf1799b555..c0e45fd990 100644 --- a/plugin/guardrail/approval/review/review.go +++ b/plugin/guardrail/approval/review/review.go @@ -18,6 +18,7 @@ import ( "trpc.group/trpc-go/trpc-agent-go/event" "trpc.group/trpc-go/trpc-agent-go/model" "trpc.group/trpc-go/trpc-agent-go/runner" + "trpc.group/trpc-go/trpc-agent-go/tool" ) // Reviewer evaluates a review request and returns an approval decision. @@ -36,6 +37,7 @@ type Action struct { ToolName string ToolDescription string Arguments json.RawMessage + Metadata tool.ToolMetadata } // TranscriptEntry is a compact transcript line used as approval evidence. diff --git a/plugin/guardrail/approval/review/review_test.go b/plugin/guardrail/approval/review/review_test.go index 3fa766335b..27f781c4b8 100644 --- a/plugin/guardrail/approval/review/review_test.go +++ b/plugin/guardrail/approval/review/review_test.go @@ -21,6 +21,7 @@ import ( "trpc.group/trpc-go/trpc-agent-go/event" "trpc.group/trpc-go/trpc-agent-go/model" "trpc.group/trpc-go/trpc-agent-go/session" + "trpc.group/trpc-go/trpc-agent-go/tool" ) type fakeRunner struct { @@ -496,6 +497,11 @@ func TestRenderUserMessage_UsesStableTemplateLayout(t *testing.T) { ToolName: "shell", ToolDescription: "Runs shell commands.", Arguments: jsonRaw(`{"command":"pwd"}`), + Metadata: tool.ToolMetadata{ + ReadOnly: false, + OpenWorld: true, + SearchOrRead: false, + }, }, Transcript: []TranscriptEntry{ {Role: model.RoleUser, Content: "Show the current directory."}, @@ -520,6 +526,14 @@ Planned action JSON: "tool_description": "Runs shell commands.", "arguments": { "command": "pwd" + }, + "metadata": { + "ReadOnly": false, + "Destructive": false, + "ConcurrencySafe": false, + "SearchOrRead": false, + "OpenWorld": true, + "MaxResultSize": 0 } } >>> APPROVAL REQUEST END`, message) diff --git a/plugin/guardrail/approval/transcript.go b/plugin/guardrail/approval/transcript.go index a2ff186861..1013a93953 100644 --- a/plugin/guardrail/approval/transcript.go +++ b/plugin/guardrail/approval/transcript.go @@ -35,6 +35,7 @@ func (p *Plugin) buildRequest(ctx context.Context, args *tool.BeforeToolArgs) (* ToolName: args.ToolName, ToolDescription: declarationDescription(args.Declaration), Arguments: cloneJSON(args.Arguments), + Metadata: args.Metadata, }, } invocation, ok := agent.InvocationFromContext(ctx) diff --git a/tool/callbacks.go b/tool/callbacks.go index aca53ec78b..bbf474c80f 100644 --- a/tool/callbacks.go +++ b/tool/callbacks.go @@ -64,6 +64,8 @@ type BeforeToolArgs struct { Declaration *Declaration // Arguments is the tool arguments in JSON bytes (can be modified). Arguments []byte + // Metadata describes execution properties published by the tool. + Metadata ToolMetadata // ResumeValue is the value of the resume. ResumeValue any // ResumeMap is the map of resume values. From 7f543cd080339794cff9b6d102759a15ef38a14a Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 13:53:13 +0800 Subject: [PATCH 86/95] feat(approval): return approval denied result --- internal/flow/processor/functioncall.go | 3 +- internal/flow/processor/functioncall_test.go | 66 ++++++++++++++++++++ plugin/guardrail/approval/approval.go | 2 +- plugin/guardrail/approval/approval_test.go | 18 +++++- tool/permission.go | 13 ++++ tool/permission_test.go | 9 +++ 6 files changed, 106 insertions(+), 5 deletions(-) diff --git a/internal/flow/processor/functioncall.go b/internal/flow/processor/functioncall.go index 949e4296f5..f20e5d0a13 100644 --- a/internal/flow/processor/functioncall.go +++ b/internal/flow/processor/functioncall.go @@ -1888,7 +1888,8 @@ func isPermissionResult(result any) bool { func isPermissionResultStatus(status string) bool { switch status { case tool.PermissionResultStatusDenied, - tool.PermissionResultStatusApprovalRequired: + tool.PermissionResultStatusApprovalRequired, + tool.PermissionResultStatusApprovalDenied: return true default: return false diff --git a/internal/flow/processor/functioncall_test.go b/internal/flow/processor/functioncall_test.go index 8b5f60a098..ce9d757095 100644 --- a/internal/flow/processor/functioncall_test.go +++ b/internal/flow/processor/functioncall_test.go @@ -9145,6 +9145,72 @@ func TestExecuteToolCall_ToolPermissionResultSkipsToolResultMessagesCallback( require.JSONEq(t, permissionJSON, choices[0].Message.Content) } +func TestExecuteToolCall_ApprovalDeniedSkipsToolResultMessagesCallback( + t *testing.T, +) { + const ( + toolName = "delete_file" + toolCallID = "call-approval-denied" + denyReason = "Automatic approval review denied (risk: high): write access is disabled" + permissionJSON = `{"status":"approval_denied","tool":"delete_file","reason":"Automatic approval review denied (risk: high): write access is disabled"}` + ) + + var ( + calledTool bool + calledResultMessages bool + ) + callbacks := tool.NewCallbacks() + callbacks.RegisterBeforeTool(func( + _ context.Context, + _ *tool.BeforeToolArgs, + ) (*tool.BeforeToolResult, error) { + return &tool.BeforeToolResult{ + CustomResult: tool.ApprovalDeniedResultFor(toolName, denyReason), + }, nil + }) + callbacks.RegisterToolResultMessages(func( + _ context.Context, + _ *tool.ToolResultMessagesInput, + ) (any, error) { + calledResultMessages = true + return model.Message{ + Role: model.RoleUser, + Content: "overridden", + }, nil + }) + p := NewFunctionCallResponseProcessor(false, callbacks) + tl := &mockCallableTool{ + declaration: &tool.Declaration{Name: toolName}, + callFn: func(_ context.Context, _ []byte) (any, error) { + calledTool = true + return map[string]any{"ok": true}, nil + }, + } + + _, choices, _, _, _, err := p.executeToolCall( + context.Background(), + &agent.Invocation{RunOptions: agent.NewRunOptions()}, + model.ToolCall{ + ID: toolCallID, + Function: model.FunctionDefinitionParam{ + Name: toolName, + Arguments: []byte(`{}`), + }, + }, + map[string]tool.Tool{toolName: tl}, + 0, + nil, + ) + require.NoError(t, err) + require.False(t, calledTool) + require.False(t, calledResultMessages) + require.Len(t, choices, 1) + require.Equal(t, model.RoleTool, choices[0].Message.Role) + require.Equal(t, toolCallID, choices[0].Message.ToolID) + require.Equal(t, toolName, choices[0].Message.ToolName) + require.JSONEq(t, permissionJSON, choices[0].Message.Content) +} + func TestExecuteToolWithCallbacks_MandatoryPermissionDenyCannotBeOverridden( t *testing.T, ) { diff --git a/plugin/guardrail/approval/approval.go b/plugin/guardrail/approval/approval.go index 1c82db1c8f..2ee467e831 100644 --- a/plugin/guardrail/approval/approval.go +++ b/plugin/guardrail/approval/approval.go @@ -209,7 +209,7 @@ func (p *Plugin) beforeTool() tool.BeforeToolCallbackStructured { } log.WarnContext(ctx, denyMessage) return &tool.BeforeToolResult{ - CustomResult: denyMessage, + CustomResult: tool.ApprovalDeniedResultFor(args.ToolName, denyMessage), }, nil default: return &tool.BeforeToolResult{ diff --git a/plugin/guardrail/approval/approval_test.go b/plugin/guardrail/approval/approval_test.go index 8fe786d9f0..55e0f12bf4 100644 --- a/plugin/guardrail/approval/approval_test.go +++ b/plugin/guardrail/approval/approval_test.go @@ -440,7 +440,11 @@ func TestBeforeTool_RequireApprovalWritesRejectedAuditRecords(t *testing.T) { }) require.NoError(t, runErr) require.NotNil(t, result) - require.Equal(t, "Automatic approval review denied (risk: high): Command rm -rf workspace can delete workspace files.", result.CustomResult) + require.Equal(t, tool.PermissionResult{ + Status: tool.PermissionResultStatusApprovalDenied, + Tool: "shell", + Reason: "Automatic approval review denied (risk: high): Command rm -rf workspace can delete workspace files.", + }, result.CustomResult) records := audit.Records() require.Len(t, records, 2) @@ -530,7 +534,11 @@ func TestBeforeTool_EmptyDecisionFieldsDoNotFail(t *testing.T) { }) require.NoError(t, runErr) require.NotNil(t, result) - require.Equal(t, "Automatic approval review denied (risk: ): ", result.CustomResult) + require.Equal(t, tool.PermissionResult{ + Status: tool.PermissionResultStatusApprovalDenied, + Tool: "shell", + Reason: "Automatic approval review denied (risk: ): ", + }, result.CustomResult) } func TestBeforeTool_ReviewerDeniedLogsWarning(t *testing.T) { @@ -563,7 +571,11 @@ func TestBeforeTool_ReviewerDeniedLogsWarning(t *testing.T) { require.NotNil(t, result) require.Equal( t, - "Automatic approval review denied (risk: high): The command is destructive and exceeds safe automatic approval.", + tool.PermissionResult{ + Status: tool.PermissionResultStatusApprovalDenied, + Tool: "shell", + Reason: "Automatic approval review denied (risk: high): The command is destructive and exceeds safe automatic approval.", + }, result.CustomResult, ) require.Equal( diff --git a/tool/permission.go b/tool/permission.go index cefa83622e..b9644046e7 100644 --- a/tool/permission.go +++ b/tool/permission.go @@ -27,6 +27,9 @@ const ( PermissionResultStatusDenied = "denied" // PermissionResultStatusApprovalRequired is returned when a tool call needs approval. PermissionResultStatusApprovalRequired = "approval_required" + // PermissionResultStatusApprovalDenied is returned when an approval reviewer + // rejects a tool call. + PermissionResultStatusApprovalDenied = "approval_denied" ) // PermissionAction is the normalized action returned by permission checks. @@ -139,3 +142,13 @@ func PermissionResultFor(toolName string, decision PermissionDecision) Permissio Reason: decision.Reason, } } + +// ApprovalDeniedResultFor builds the structured tool result returned when an +// approval reviewer explicitly rejects a tool call. +func ApprovalDeniedResultFor(toolName string, reason string) PermissionResult { + return PermissionResult{ + Status: PermissionResultStatusApprovalDenied, + Tool: toolName, + Reason: reason, + } +} diff --git a/tool/permission_test.go b/tool/permission_test.go index b3774b56f9..c5bde4ad2b 100644 --- a/tool/permission_test.go +++ b/tool/permission_test.go @@ -144,3 +144,12 @@ func TestPermissionResultFor(t *testing.T) { t.Fatalf("unexpected ask result: %+v", ask) } } + +func TestApprovalDeniedResultFor(t *testing.T) { + result := ApprovalDeniedResultFor(testToolName, testReason) + if result.Status != PermissionResultStatusApprovalDenied || + result.Tool != testToolName || + result.Reason != testReason { + t.Fatalf("unexpected approval denied result: %+v", result) + } +} From bfe5e933a613054149836f71c964a028e89b320d Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 14:23:36 +0800 Subject: [PATCH 87/95] feat(approval): record approval required metrics --- internal/telemetry/metric_chat_test.go | 1 + internal/telemetry/metric_tool_approval.go | 57 +++++++++ .../telemetry/metric_tool_approval_test.go | 108 ++++++++++++++++++ internal/telemetry/trace.go | 1 + plugin/guardrail/approval/approval.go | 1 + plugin/guardrail/approval/approval_test.go | 95 +++++++++++++++ plugin/guardrail/approval/audit.go | 13 +++ telemetry/metric/metric.go | 21 ++++ telemetry/metric/metric_test.go | 35 ++++++ telemetry/semconv/metrics/metrics.go | 4 + telemetry/semconv/trace/trace.go | 2 + 11 files changed, 338 insertions(+) create mode 100644 internal/telemetry/metric_tool_approval.go create mode 100644 internal/telemetry/metric_tool_approval_test.go diff --git a/internal/telemetry/metric_chat_test.go b/internal/telemetry/metric_chat_test.go index d07943d309..472b226256 100644 --- a/internal/telemetry/metric_chat_test.go +++ b/internal/telemetry/metric_chat_test.go @@ -102,6 +102,7 @@ func TestChatMetricsTracker_TrackResponse_ReasoningDuration_UsesLazyNow(t *testi require.True(t, tracker.isFirstToken, "expected empty chunk to be ignored for TTFT") require.Zero(t, tracker.firstTokenTimeDuration, "expected TTFT to remain unset after empty chunk") + time.Sleep(time.Millisecond) tracker.TrackResponse(&model.Response{ Choices: []model.Choice{ { diff --git a/internal/telemetry/metric_tool_approval.go b/internal/telemetry/metric_tool_approval.go new file mode 100644 index 0000000000..ef708f2a1e --- /dev/null +++ b/internal/telemetry/metric_tool_approval.go @@ -0,0 +1,57 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package telemetry + +import ( + "context" + + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/metric" + + "trpc.group/trpc-go/trpc-agent-go/telemetry/semconv/metrics" + semconvtrace "trpc.group/trpc-go/trpc-agent-go/telemetry/semconv/trace" +) + +var ( + // ToolApprovalMeter is the meter used for recording tool approval metrics. + ToolApprovalMeter = MeterProvider.Meter(metrics.MeterNameToolApproval) + + // ToolApprovalMetricRequiredTotal records tool calls that require explicit approval. + ToolApprovalMetricRequiredTotal metric.Int64Counter +) + +// ToolApprovalAttributes is the attributes for tool approval metrics. +type ToolApprovalAttributes struct { + TenantID string + AppName string + ToolName string +} + +func (a ToolApprovalAttributes) toAttributes() []attribute.KeyValue { + attrs := []attribute.KeyValue{ + attribute.String(semconvtrace.KeyGenAIOperationName, OperationToolApproval), + attribute.String(semconvtrace.KeyGenAISystem, semconvtrace.SystemTRPCGoAgent), + attribute.String(semconvtrace.KeyGenAIToolName, a.ToolName), + } + if a.TenantID != "" { + attrs = append(attrs, attribute.String(semconvtrace.KeyTRPCAgentGoTenantID, a.TenantID)) + } + if a.AppName != "" { + attrs = append(attrs, attribute.String(semconvtrace.KeyTRPCAgentGoAppName, a.AppName)) + } + return attrs +} + +// ReportToolApprovalRequiredMetrics reports that a tool call required explicit approval. +func ReportToolApprovalRequiredMetrics(ctx context.Context, attrs ToolApprovalAttributes) { + if ToolApprovalMetricRequiredTotal == nil { + return + } + ToolApprovalMetricRequiredTotal.Add(ctx, 1, metric.WithAttributes(attrs.toAttributes()...)) +} diff --git a/internal/telemetry/metric_tool_approval_test.go b/internal/telemetry/metric_tool_approval_test.go new file mode 100644 index 0000000000..69f736a77b --- /dev/null +++ b/internal/telemetry/metric_tool_approval_test.go @@ -0,0 +1,108 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package telemetry + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" + "go.opentelemetry.io/otel/attribute" + sdkmetric "go.opentelemetry.io/otel/sdk/metric" + "go.opentelemetry.io/otel/sdk/metric/metricdata" + + "trpc.group/trpc-go/trpc-agent-go/telemetry/semconv/metrics" + semconvtrace "trpc.group/trpc-go/trpc-agent-go/telemetry/semconv/trace" +) + +func TestReportToolApprovalRequiredMetricsNoopWhenCounterNil(t *testing.T) { + originalCounter := ToolApprovalMetricRequiredTotal + t.Cleanup(func() { + ToolApprovalMetricRequiredTotal = originalCounter + }) + + ToolApprovalMetricRequiredTotal = nil + require.NotPanics(t, func() { + ReportToolApprovalRequiredMetrics(context.Background(), ToolApprovalAttributes{ + TenantID: "tenant-1", + AppName: "app-1", + ToolName: "shell", + }) + }) +} + +func TestReportToolApprovalRequiredMetrics(t *testing.T) { + reader := sdkmetric.NewManualReader() + provider := sdkmetric.NewMeterProvider(sdkmetric.WithReader(reader)) + + originalProvider := MeterProvider + originalMeter := ToolApprovalMeter + originalCounter := ToolApprovalMetricRequiredTotal + t.Cleanup(func() { + MeterProvider = originalProvider + ToolApprovalMeter = originalMeter + ToolApprovalMetricRequiredTotal = originalCounter + }) + + MeterProvider = provider + ToolApprovalMeter = provider.Meter(metrics.MeterNameToolApproval) + var err error + ToolApprovalMetricRequiredTotal, err = ToolApprovalMeter.Int64Counter(metrics.MetricToolApprovalRequiredTotal) + require.NoError(t, err) + + ctx := context.Background() + ReportToolApprovalRequiredMetrics(ctx, ToolApprovalAttributes{ + TenantID: "tenant-1", + AppName: "app-1", + ToolName: "shell", + }) + + var rm metricdata.ResourceMetrics + require.NoError(t, reader.Collect(ctx, &rm)) + + points := toolApprovalSumPoints(t, rm, metrics.MetricToolApprovalRequiredTotal) + require.Len(t, points, 1) + require.Equal(t, int64(1), points[0].Value) + requireToolApprovalAttr(t, points[0].Attributes, semconvtrace.KeyGenAIOperationName, OperationToolApproval) + requireToolApprovalAttr(t, points[0].Attributes, semconvtrace.KeyGenAISystem, semconvtrace.SystemTRPCGoAgent) + requireToolApprovalAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoTenantID, "tenant-1") + requireToolApprovalAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoAppName, "app-1") + requireToolApprovalAttr(t, points[0].Attributes, semconvtrace.KeyGenAIToolName, "shell") +} + +func toolApprovalSumPoints( + t *testing.T, + rm metricdata.ResourceMetrics, + metricName string, +) []metricdata.DataPoint[int64] { + t.Helper() + for _, scopeMetric := range rm.ScopeMetrics { + for _, metric := range scopeMetric.Metrics { + if metric.Name != metricName { + continue + } + sum, ok := metric.Data.(metricdata.Sum[int64]) + require.True(t, ok) + return sum.DataPoints + } + } + t.Fatalf("metric %s not found", metricName) + return nil +} + +func requireToolApprovalAttr(t *testing.T, set attribute.Set, key string, value string) { + t.Helper() + for _, kv := range set.ToSlice() { + if string(kv.Key) == key { + require.Equal(t, value, kv.Value.AsString()) + return + } + } + t.Fatalf("attribute %s not found", key) +} diff --git a/internal/telemetry/trace.go b/internal/telemetry/trace.go index c91c8e0be7..d13211d09d 100644 --- a/internal/telemetry/trace.go +++ b/internal/telemetry/trace.go @@ -56,6 +56,7 @@ const ( OperationExecuteTool = "execute_tool" OperationToolCall = "tool.call" + OperationToolApproval = "tool.approval" OperationMemorySearch = "memory.search" OperationMemoryWrite = "memory.write" OperationSummaryCreate = "summary.create" diff --git a/plugin/guardrail/approval/approval.go b/plugin/guardrail/approval/approval.go index 2ee467e831..cda093ffcb 100644 --- a/plugin/guardrail/approval/approval.go +++ b/plugin/guardrail/approval/approval.go @@ -128,6 +128,7 @@ func (p *Plugin) beforeTool() tool.BeforeToolCallbackStructured { CustomResult: fmt.Sprintf("approval audit failed for tool %q: %v", args.ToolName, err), }, nil } + reportApprovalRequiredMetric(ctx, args) decision, err := p.reviewer.Review(ctx, req) if err != nil { log.ErrorfContext( diff --git a/plugin/guardrail/approval/approval_test.go b/plugin/guardrail/approval/approval_test.go index 55e0f12bf4..a79e40cc9d 100644 --- a/plugin/guardrail/approval/approval_test.go +++ b/plugin/guardrail/approval/approval_test.go @@ -17,8 +17,12 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "go.opentelemetry.io/otel/attribute" + sdkmetric "go.opentelemetry.io/otel/sdk/metric" + "go.opentelemetry.io/otel/sdk/metric/metricdata" "trpc.group/trpc-go/trpc-agent-go/agent" "trpc.group/trpc-go/trpc-agent-go/event" + itelemetry "trpc.group/trpc-go/trpc-agent-go/internal/telemetry" approvallog "trpc.group/trpc-go/trpc-agent-go/log" "trpc.group/trpc-go/trpc-agent-go/model" "trpc.group/trpc-go/trpc-agent-go/platform" @@ -26,6 +30,8 @@ import ( approvalreview "trpc.group/trpc-go/trpc-agent-go/plugin/guardrail/approval/review" guardtranscript "trpc.group/trpc-go/trpc-agent-go/plugin/guardrail/internal/transcript" "trpc.group/trpc-go/trpc-agent-go/session" + "trpc.group/trpc-go/trpc-agent-go/telemetry/semconv/metrics" + semconvtrace "trpc.group/trpc-go/trpc-agent-go/telemetry/semconv/trace" "trpc.group/trpc-go/trpc-agent-go/tool" ) @@ -344,6 +350,64 @@ func TestBeforeTool_MetadataRiskRequiresApproval(t *testing.T) { require.Equal(t, metadata, captured.Action.Metadata) } +func TestBeforeTool_RequireApprovalRecordsRequiredMetric(t *testing.T) { + reader := sdkmetric.NewManualReader() + provider := sdkmetric.NewMeterProvider(sdkmetric.WithReader(reader)) + + originalProvider := itelemetry.MeterProvider + originalMeter := itelemetry.ToolApprovalMeter + originalRequired := itelemetry.ToolApprovalMetricRequiredTotal + t.Cleanup(func() { + itelemetry.MeterProvider = originalProvider + itelemetry.ToolApprovalMeter = originalMeter + itelemetry.ToolApprovalMetricRequiredTotal = originalRequired + }) + + itelemetry.MeterProvider = provider + itelemetry.ToolApprovalMeter = provider.Meter(metrics.MeterNameToolApproval) + var counterErr error + itelemetry.ToolApprovalMetricRequiredTotal, counterErr = itelemetry.ToolApprovalMeter.Int64Counter( + metrics.MetricToolApprovalRequiredTotal, + ) + require.NoError(t, counterErr) + + p, err := New(WithReviewer(&stubReviewer{ + reviewFn: func(ctx context.Context, req *approvalreview.Request) (*approvalreview.Decision, error) { + return &approvalreview.Decision{ + Approved: true, + RiskScore: 18, + RiskLevel: "low", + Reason: "Approved command.", + }, nil + }, + })) + require.NoError(t, err) + + callbacks := registeredToolCallbacks(t, p) + ctx := ContextWithAuditContext(context.Background(), AuditContext{ + TenantID: "tenant-metric", + AppID: "app-metric", + }) + result, runErr := callbacks.RunBeforeTool(ctx, &tool.BeforeToolArgs{ + ToolName: "shell", + ToolCallID: "call-metric", + Arguments: []byte(`{"command":"pwd"}`), + }) + require.NoError(t, runErr) + require.NotNil(t, result) + + var rm metricdata.ResourceMetrics + require.NoError(t, reader.Collect(ctx, &rm)) + points := approvalMetricSumPoints(t, rm, metrics.MetricToolApprovalRequiredTotal) + require.Len(t, points, 1) + require.Equal(t, int64(1), points[0].Value) + requireApprovalMetricAttr(t, points[0].Attributes, semconvtrace.KeyGenAIOperationName, itelemetry.OperationToolApproval) + requireApprovalMetricAttr(t, points[0].Attributes, semconvtrace.KeyGenAISystem, semconvtrace.SystemTRPCGoAgent) + requireApprovalMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoTenantID, "tenant-metric") + requireApprovalMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoAppName, "app-metric") + requireApprovalMetricAttr(t, points[0].Attributes, semconvtrace.KeyGenAIToolName, "shell") +} + func TestBeforeTool_RequireApprovalWritesApprovedAuditRecords(t *testing.T) { now := time.Date(2026, 7, 11, 10, 30, 0, 0, time.UTC) audit := platform.NewInMemoryAuditSink() @@ -734,3 +798,34 @@ func stringsRepeat(value string, n int) string { } return string(result) } + +func approvalMetricSumPoints( + t *testing.T, + rm metricdata.ResourceMetrics, + metricName string, +) []metricdata.DataPoint[int64] { + t.Helper() + for _, scopeMetric := range rm.ScopeMetrics { + for _, metric := range scopeMetric.Metrics { + if metric.Name != metricName { + continue + } + sum, ok := metric.Data.(metricdata.Sum[int64]) + require.True(t, ok) + return sum.DataPoints + } + } + t.Fatalf("metric %s not found", metricName) + return nil +} + +func requireApprovalMetricAttr(t *testing.T, set attribute.Set, key string, value string) { + t.Helper() + for _, kv := range set.ToSlice() { + if string(kv.Key) == key { + require.Equal(t, value, kv.Value.AsString()) + return + } + } + t.Fatalf("attribute %s not found", key) +} diff --git a/plugin/guardrail/approval/audit.go b/plugin/guardrail/approval/audit.go index f4f1dd04e6..57967d5da0 100644 --- a/plugin/guardrail/approval/audit.go +++ b/plugin/guardrail/approval/audit.go @@ -17,6 +17,7 @@ import ( "time" "trpc.group/trpc-go/trpc-agent-go/agent" + itelemetry "trpc.group/trpc-go/trpc-agent-go/internal/telemetry" "trpc.group/trpc-go/trpc-agent-go/platform" "trpc.group/trpc-go/trpc-agent-go/tool" ) @@ -122,6 +123,18 @@ func (p *Plugin) writeApprovalAudit( return p.auditSink.WriteAudit(ctx, record) } +func reportApprovalRequiredMetric(ctx context.Context, args *tool.BeforeToolArgs) { + if args == nil { + return + } + auditCtx := approvalAuditContextFrom(ctx) + itelemetry.ReportToolApprovalRequiredMetrics(ctx, itelemetry.ToolApprovalAttributes{ + TenantID: auditCtx.TenantID, + AppName: auditCtx.AppID, + ToolName: strings.TrimSpace(args.ToolName), + }) +} + func (p *Plugin) auditApproverUserID() string { if p == nil { return "" diff --git a/telemetry/metric/metric.go b/telemetry/metric/metric.go index 0bf6e6eaa8..9410d2c123 100644 --- a/telemetry/metric/metric.go +++ b/telemetry/metric/metric.go @@ -115,6 +115,9 @@ func InitMeterProvider(mp metric.MeterProvider) error { return fmt.Errorf("failed to create execute tool metric GenAIClientOperationDuration: %w", err) } + if err := initToolApprovalMetrics(mp); err != nil { + return err + } if err := initInvokeAgentMetrics(mp); err != nil { return err } @@ -197,6 +200,24 @@ func setExecuteToolHistogramBuckets(metricName string, boundaries []float64) err } } +func initToolApprovalMetrics(mp metric.MeterProvider) error { + if mp == nil { + return fmt.Errorf("tool approval meter provider is nil") + } + meterName := metrics.MeterNameToolApproval + itelemetry.ToolApprovalMeter = mp.Meter(meterName) + var err error + itelemetry.ToolApprovalMetricRequiredTotal, err = itelemetry.ToolApprovalMeter.Int64Counter( + metrics.MetricToolApprovalRequiredTotal, + metric.WithDescription("Total number of tool calls requiring explicit approval"), + metric.WithUnit("1"), + ) + if err != nil { + return fmt.Errorf("failed to create %s metric %s: %w", meterName, metrics.MetricToolApprovalRequiredTotal, err) + } + return nil +} + func setInvokeAgentHistogramBuckets(metricName string, boundaries []float64) error { switch metricName { case metrics.MetricTRPCAgentGoClientTimeToFirstToken: diff --git a/telemetry/metric/metric_test.go b/telemetry/metric/metric_test.go index d74cf3f181..1ca844a014 100644 --- a/telemetry/metric/metric_test.go +++ b/telemetry/metric/metric_test.go @@ -294,8 +294,12 @@ func TestInitMeterProvider(t *testing.T) { // Save original meter provider originalMP := itelemetry.MeterProvider + originalToolApprovalMeter := itelemetry.ToolApprovalMeter + originalToolApprovalRequired := itelemetry.ToolApprovalMetricRequiredTotal defer func() { itelemetry.MeterProvider = originalMP + itelemetry.ToolApprovalMeter = originalToolApprovalMeter + itelemetry.ToolApprovalMetricRequiredTotal = originalToolApprovalRequired }() // Create a test meter provider @@ -351,6 +355,12 @@ func TestInitMeterProvider(t *testing.T) { if itelemetry.ExecuteToolMetricGenAIClientOperationDuration == nil { t.Error("ExecuteToolMetricGenAIClientOperationDuration was not created") } + if itelemetry.ToolApprovalMeter == nil { + t.Error("ToolApprovalMeter was not created") + } + if itelemetry.ToolApprovalMetricRequiredTotal == nil { + t.Error("ToolApprovalMetricRequiredTotal was not created") + } if itelemetry.WorkflowMeter == nil { t.Error("WorkflowMeter was not created") } @@ -391,6 +401,31 @@ func TestInitMeterProvider_WorkflowMetricError(t *testing.T) { } } +func TestInitToolApprovalMetrics_ErrorHandling(t *testing.T) { + originalToolApprovalMeter := itelemetry.ToolApprovalMeter + originalToolApprovalRequired := itelemetry.ToolApprovalMetricRequiredTotal + defer func() { + itelemetry.ToolApprovalMeter = originalToolApprovalMeter + itelemetry.ToolApprovalMetricRequiredTotal = originalToolApprovalRequired + }() + + if err := initToolApprovalMetrics(nil); err == nil || !strings.Contains(err.Error(), "tool approval meter provider is nil") { + t.Fatalf("expected nil provider error, got %v", err) + } + + mp := &mockMeterProvider{meter: &mockMeter{ + shouldFail: true, + failOn: metrics.MetricToolApprovalRequiredTotal, + }} + err := initToolApprovalMetrics(mp) + if err == nil { + t.Fatalf("expected tool approval counter creation error") + } + if !strings.Contains(err.Error(), "failed to create trpc_agent_go.internal.tool_approval metric tool_approval_required_total") { + t.Fatalf("unexpected error: %v", err) + } +} + func TestInitWorkflowMetrics_ErrorHandling(t *testing.T) { originalWorkflowMeter := itelemetry.WorkflowMeter originalWorkflowOpDur := itelemetry.WorkflowMetricGenAIClientOperationDuration diff --git a/telemetry/semconv/metrics/metrics.go b/telemetry/semconv/metrics/metrics.go index f35cbf0901..00f2c10864 100644 --- a/telemetry/semconv/metrics/metrics.go +++ b/telemetry/semconv/metrics/metrics.go @@ -59,6 +59,8 @@ const ( // MetricTRPCAgentGoClientRequestCnt represents the request count for client. MetricTRPCAgentGoClientRequestCnt = "trpc_agent_go.client.request_cnt" + // MetricToolApprovalRequiredTotal records tool calls that require explicit approval. + MetricToolApprovalRequiredTotal = "tool_approval_required_total" ////////////////////////// server //////////////////////// @@ -79,4 +81,6 @@ const ( MeterNameWorkflow = "trpc_agent_go.internal.workflow" // MeterNameInvokeAgent is the meter name for invoke agent operations. MeterNameInvokeAgent = "trpc_agent_go.internal.invoke_agent" + // MeterNameToolApproval is the meter name for tool approval operations. + MeterNameToolApproval = "trpc_agent_go.internal.tool_approval" ) diff --git a/telemetry/semconv/trace/trace.go b/telemetry/semconv/trace/trace.go index 3d453bcfda..d31ff1b203 100644 --- a/telemetry/semconv/trace/trace.go +++ b/telemetry/semconv/trace/trace.go @@ -42,6 +42,8 @@ const ( // KeyTRPCAgentGoAppName is the attribute key for application name. KeyTRPCAgentGoAppName = "trpc_go_agent.app.name" + // KeyTRPCAgentGoTenantID is the attribute key for tenant ID. + KeyTRPCAgentGoTenantID = "trpc_go_agent.tenant.id" // KeyTRPCAgentGoUserID is the attribute key for user ID. KeyTRPCAgentGoUserID = "trpc_go_agent.user.id" // KeyTRPCAgentGoClientTimeToFirstToken is the attribute key for time to first token metric. From f3a97fd454a16d9d146a7c93b13aff7e37262db5 Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 14:34:12 +0800 Subject: [PATCH 88/95] feat(audit): record audit write failure metrics --- internal/telemetry/metric_audit.go | 63 ++++++++++++ internal/telemetry/metric_audit_test.go | 112 +++++++++++++++++++++ internal/telemetry/trace.go | 1 + platform/gateway/service.go | 9 +- platform/gateway/service_test.go | 99 ++++++++++++++++++ platform/toolpolicy/policy.go | 7 ++ platform/toolpolicy/policy_test.go | 73 ++++++++++++++ plugin/guardrail/approval/approval_test.go | 75 ++++++++++++++ plugin/guardrail/approval/audit.go | 11 +- telemetry/metric/metric.go | 21 ++++ telemetry/metric/metric_test.go | 35 +++++++ telemetry/semconv/metrics/metrics.go | 4 + telemetry/semconv/trace/trace.go | 2 + 13 files changed, 510 insertions(+), 2 deletions(-) create mode 100644 internal/telemetry/metric_audit.go create mode 100644 internal/telemetry/metric_audit_test.go diff --git a/internal/telemetry/metric_audit.go b/internal/telemetry/metric_audit.go new file mode 100644 index 0000000000..9cb1d2442c --- /dev/null +++ b/internal/telemetry/metric_audit.go @@ -0,0 +1,63 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package telemetry + +import ( + "context" + + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/metric" + + "trpc.group/trpc-go/trpc-agent-go/telemetry/semconv/metrics" + semconvtrace "trpc.group/trpc-go/trpc-agent-go/telemetry/semconv/trace" +) + +var ( + // AuditMeter is the meter used for recording audit metrics. + AuditMeter = MeterProvider.Meter(metrics.MeterNameAudit) + + // AuditMetricWriteFailedTotal records failed audit sink writes. + AuditMetricWriteFailedTotal metric.Int64Counter +) + +// AuditAttributes is the attributes for audit metrics. +type AuditAttributes struct { + TenantID string + AppName string + Decision string + Error error +} + +func (a AuditAttributes) toAttributes() []attribute.KeyValue { + attrs := []attribute.KeyValue{ + attribute.String(semconvtrace.KeyGenAIOperationName, OperationAuditWrite), + attribute.String(semconvtrace.KeyGenAISystem, semconvtrace.SystemTRPCGoAgent), + } + if a.TenantID != "" { + attrs = append(attrs, attribute.String(semconvtrace.KeyTRPCAgentGoTenantID, a.TenantID)) + } + if a.AppName != "" { + attrs = append(attrs, attribute.String(semconvtrace.KeyTRPCAgentGoAppName, a.AppName)) + } + if a.Decision != "" { + attrs = append(attrs, attribute.String(semconvtrace.KeyTRPCAgentGoAuditDecision, a.Decision)) + } + if a.Error != nil { + attrs = append(attrs, attribute.String(semconvtrace.KeyErrorType, ToErrorType(a.Error, semconvtrace.ValueDefaultErrorType))) + } + return attrs +} + +// ReportAuditWriteFailedMetrics reports a failed audit sink write. +func ReportAuditWriteFailedMetrics(ctx context.Context, attrs AuditAttributes) { + if AuditMetricWriteFailedTotal == nil { + return + } + AuditMetricWriteFailedTotal.Add(ctx, 1, metric.WithAttributes(attrs.toAttributes()...)) +} diff --git a/internal/telemetry/metric_audit_test.go b/internal/telemetry/metric_audit_test.go new file mode 100644 index 0000000000..c194049e1d --- /dev/null +++ b/internal/telemetry/metric_audit_test.go @@ -0,0 +1,112 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package telemetry + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/require" + "go.opentelemetry.io/otel/attribute" + sdkmetric "go.opentelemetry.io/otel/sdk/metric" + "go.opentelemetry.io/otel/sdk/metric/metricdata" + + "trpc.group/trpc-go/trpc-agent-go/telemetry/semconv/metrics" + semconvtrace "trpc.group/trpc-go/trpc-agent-go/telemetry/semconv/trace" +) + +func TestReportAuditWriteFailedMetricsNoopWhenCounterNil(t *testing.T) { + originalCounter := AuditMetricWriteFailedTotal + t.Cleanup(func() { + AuditMetricWriteFailedTotal = originalCounter + }) + + AuditMetricWriteFailedTotal = nil + require.NotPanics(t, func() { + ReportAuditWriteFailedMetrics(context.Background(), AuditAttributes{ + TenantID: "tenant-1", + AppName: "app-1", + Decision: "reject", + Error: errors.New("audit unavailable"), + }) + }) +} + +func TestReportAuditWriteFailedMetrics(t *testing.T) { + reader := sdkmetric.NewManualReader() + provider := sdkmetric.NewMeterProvider(sdkmetric.WithReader(reader)) + + originalProvider := MeterProvider + originalMeter := AuditMeter + originalCounter := AuditMetricWriteFailedTotal + t.Cleanup(func() { + MeterProvider = originalProvider + AuditMeter = originalMeter + AuditMetricWriteFailedTotal = originalCounter + }) + + MeterProvider = provider + AuditMeter = provider.Meter(metrics.MeterNameAudit) + var err error + AuditMetricWriteFailedTotal, err = AuditMeter.Int64Counter(metrics.MetricAuditWriteFailedTotal) + require.NoError(t, err) + + ctx := context.Background() + ReportAuditWriteFailedMetrics(ctx, AuditAttributes{ + TenantID: "tenant-1", + AppName: "app-1", + Decision: "reject", + Error: errors.New("audit unavailable"), + }) + + var rm metricdata.ResourceMetrics + require.NoError(t, reader.Collect(ctx, &rm)) + + points := auditSumPoints(t, rm, metrics.MetricAuditWriteFailedTotal) + require.Len(t, points, 1) + require.Equal(t, int64(1), points[0].Value) + requireAuditAttr(t, points[0].Attributes, semconvtrace.KeyGenAIOperationName, OperationAuditWrite) + requireAuditAttr(t, points[0].Attributes, semconvtrace.KeyGenAISystem, semconvtrace.SystemTRPCGoAgent) + requireAuditAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoTenantID, "tenant-1") + requireAuditAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoAppName, "app-1") + requireAuditAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoAuditDecision, "reject") + requireAuditAttr(t, points[0].Attributes, semconvtrace.KeyErrorType, semconvtrace.ValueDefaultErrorType) +} + +func auditSumPoints( + t *testing.T, + rm metricdata.ResourceMetrics, + metricName string, +) []metricdata.DataPoint[int64] { + t.Helper() + for _, scopeMetric := range rm.ScopeMetrics { + for _, metric := range scopeMetric.Metrics { + if metric.Name != metricName { + continue + } + sum, ok := metric.Data.(metricdata.Sum[int64]) + require.True(t, ok) + return sum.DataPoints + } + } + t.Fatalf("metric %s not found", metricName) + return nil +} + +func requireAuditAttr(t *testing.T, set attribute.Set, key string, value string) { + t.Helper() + for _, kv := range set.ToSlice() { + if string(kv.Key) == key { + require.Equal(t, value, kv.Value.AsString()) + return + } + } + t.Fatalf("attribute %s not found", key) +} diff --git a/internal/telemetry/trace.go b/internal/telemetry/trace.go index d13211d09d..f04590543f 100644 --- a/internal/telemetry/trace.go +++ b/internal/telemetry/trace.go @@ -57,6 +57,7 @@ const ( OperationExecuteTool = "execute_tool" OperationToolCall = "tool.call" OperationToolApproval = "tool.approval" + OperationAuditWrite = "audit.write" OperationMemorySearch = "memory.search" OperationMemoryWrite = "memory.write" OperationSummaryCreate = "summary.create" diff --git a/platform/gateway/service.go b/platform/gateway/service.go index 334fc93886..fcb0f96c47 100644 --- a/platform/gateway/service.go +++ b/platform/gateway/service.go @@ -1355,7 +1355,14 @@ func (s *Service) writeAuditTo( } record = redactionFailedAuditRecord(record, err) } - _ = auditSink.WriteAudit(ctx, record) + if err := auditSink.WriteAudit(ctx, record); err != nil { + itelemetry.ReportAuditWriteFailedMetrics(ctx, itelemetry.AuditAttributes{ + TenantID: record.TenantID, + AppName: record.AppID, + Decision: record.Decision, + Error: err, + }) + } } func isAuditRedactionFailure(err error) bool { diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go index 1bd393db46..0bcf3ad949 100644 --- a/platform/gateway/service_test.go +++ b/platform/gateway/service_test.go @@ -19,17 +19,23 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "go.opentelemetry.io/otel/attribute" + sdkmetric "go.opentelemetry.io/otel/sdk/metric" + "go.opentelemetry.io/otel/sdk/metric/metricdata" sdktrace "go.opentelemetry.io/otel/sdk/trace" "go.opentelemetry.io/otel/sdk/trace/tracetest" "trpc.group/trpc-go/trpc-agent-go/agent" "trpc.group/trpc-go/trpc-agent-go/event" + itelemetry "trpc.group/trpc-go/trpc-agent-go/internal/telemetry" "trpc.group/trpc-go/trpc-agent-go/model" "trpc.group/trpc-go/trpc-agent-go/platform" "trpc.group/trpc-go/trpc-agent-go/platform/channeladapter" "trpc.group/trpc-go/trpc-agent-go/plugin" "trpc.group/trpc-go/trpc-agent-go/plugin/guardrail/approval" approvalreview "trpc.group/trpc-go/trpc-agent-go/plugin/guardrail/approval/review" + "trpc.group/trpc-go/trpc-agent-go/telemetry/semconv/metrics" + semconvtrace "trpc.group/trpc-go/trpc-agent-go/telemetry/semconv/trace" telemetrytrace "trpc.group/trpc-go/trpc-agent-go/telemetry/trace" "trpc.group/trpc-go/trpc-agent-go/tool" ) @@ -1592,6 +1598,39 @@ func TestServiceWriteAuditRecordsRedactionFailureFallback(t *testing.T) { assert.NotContains(t, record.RedactedDetailRef, "Authorization") } +func TestServiceWriteAuditRecordsAuditWriteFailureMetric(t *testing.T) { + reader, restore := useAuditMetrics(t) + defer restore() + + svc := NewService( + NewInMemoryRegistry(), + platform.NewInMemoryIdempotencyStore(), + NewInMemoryOutboundStore(), + ) + record := platform.AuditRecord{ + AuditID: "audit-write-failure", + TenantID: "tenant-a", + AppID: "app", + RequestID: "request-1", + MessageID: "msg-1", + Decision: "reject", + DecisionReason: "access denied", + CreatedAt: time.Unix(1500, 0), + } + + svc.writeAuditTo(context.Background(), failingGatewayAuditSink{}, record) + + points := collectGatewayAuditWriteFailedPoints(t, reader) + require.Len(t, points, 1) + require.Equal(t, int64(1), points[0].Value) + requireGatewayAuditMetricAttr(t, points[0].Attributes, semconvtrace.KeyGenAIOperationName, itelemetry.OperationAuditWrite) + requireGatewayAuditMetricAttr(t, points[0].Attributes, semconvtrace.KeyGenAISystem, semconvtrace.SystemTRPCGoAgent) + requireGatewayAuditMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoTenantID, "tenant-a") + requireGatewayAuditMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoAppName, "app") + requireGatewayAuditMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoAuditDecision, "reject") + requireGatewayAuditMetricAttr(t, points[0].Attributes, semconvtrace.KeyErrorType, semconvtrace.ValueDefaultErrorType) +} + func TestServiceHandleInboundRejectsBudgetExceededBeforeIdempotency(t *testing.T) { ctx := context.Background() registry := NewInMemoryRegistry() @@ -2661,3 +2700,63 @@ func spanEventsText(span sdktrace.ReadOnlySpan) string { } return strings.Join(values, "\n") } + +type failingGatewayAuditSink struct{} + +func (failingGatewayAuditSink) WriteAudit(context.Context, platform.AuditRecord) error { + return errors.New("audit unavailable") +} + +func useAuditMetrics(t *testing.T) (*sdkmetric.ManualReader, func()) { + t.Helper() + reader := sdkmetric.NewManualReader() + provider := sdkmetric.NewMeterProvider(sdkmetric.WithReader(reader)) + + originalProvider := itelemetry.MeterProvider + originalMeter := itelemetry.AuditMeter + originalCounter := itelemetry.AuditMetricWriteFailedTotal + + itelemetry.MeterProvider = provider + itelemetry.AuditMeter = provider.Meter(metrics.MeterNameAudit) + var err error + itelemetry.AuditMetricWriteFailedTotal, err = itelemetry.AuditMeter.Int64Counter(metrics.MetricAuditWriteFailedTotal) + require.NoError(t, err) + + return reader, func() { + itelemetry.MeterProvider = originalProvider + itelemetry.AuditMeter = originalMeter + itelemetry.AuditMetricWriteFailedTotal = originalCounter + } +} + +func collectGatewayAuditWriteFailedPoints( + t *testing.T, + reader *sdkmetric.ManualReader, +) []metricdata.DataPoint[int64] { + t.Helper() + var rm metricdata.ResourceMetrics + require.NoError(t, reader.Collect(context.Background(), &rm)) + for _, scopeMetric := range rm.ScopeMetrics { + for _, metric := range scopeMetric.Metrics { + if metric.Name != metrics.MetricAuditWriteFailedTotal { + continue + } + sum, ok := metric.Data.(metricdata.Sum[int64]) + require.True(t, ok) + return sum.DataPoints + } + } + t.Fatalf("metric %s not found", metrics.MetricAuditWriteFailedTotal) + return nil +} + +func requireGatewayAuditMetricAttr(t *testing.T, set attribute.Set, key string, value string) { + t.Helper() + for _, kv := range set.ToSlice() { + if string(kv.Key) == key { + require.Equal(t, value, kv.Value.AsString()) + return + } + } + t.Fatalf("attribute %s not found", key) +} diff --git a/platform/toolpolicy/policy.go b/platform/toolpolicy/policy.go index 7e3fe6ea42..dca962a2b0 100644 --- a/platform/toolpolicy/policy.go +++ b/platform/toolpolicy/policy.go @@ -20,6 +20,7 @@ import ( oteltrace "go.opentelemetry.io/otel/trace" "trpc.group/trpc-go/trpc-agent-go/agent" + itelemetry "trpc.group/trpc-go/trpc-agent-go/internal/telemetry" "trpc.group/trpc-go/trpc-agent-go/platform" "trpc.group/trpc-go/trpc-agent-go/plugin" "trpc.group/trpc-go/trpc-agent-go/plugin/guardrail/approval" @@ -709,6 +710,12 @@ func (p *Policy) writeAudit(ctx context.Context, summary ApprovalSummary) error return fmt.Errorf("tool policy audit record: %w", err) } if err := p.audit.WriteAudit(ctx, record); err != nil { + itelemetry.ReportAuditWriteFailedMetrics(ctx, itelemetry.AuditAttributes{ + TenantID: record.TenantID, + AppName: record.AppID, + Decision: record.Decision, + Error: err, + }) return fmt.Errorf("write tool policy audit: %w", err) } return nil diff --git a/platform/toolpolicy/policy_test.go b/platform/toolpolicy/policy_test.go index aed7adbc32..9f2c968a87 100644 --- a/platform/toolpolicy/policy_test.go +++ b/platform/toolpolicy/policy_test.go @@ -16,14 +16,20 @@ import ( "time" "github.com/stretchr/testify/require" + "go.opentelemetry.io/otel/attribute" + sdkmetric "go.opentelemetry.io/otel/sdk/metric" + "go.opentelemetry.io/otel/sdk/metric/metricdata" oteltrace "go.opentelemetry.io/otel/trace" "trpc.group/trpc-go/trpc-agent-go/agent" + itelemetry "trpc.group/trpc-go/trpc-agent-go/internal/telemetry" "trpc.group/trpc-go/trpc-agent-go/platform" "trpc.group/trpc-go/trpc-agent-go/plugin" "trpc.group/trpc-go/trpc-agent-go/plugin/guardrail/approval" "trpc.group/trpc-go/trpc-agent-go/plugin/guardrail/approval/review" "trpc.group/trpc-go/trpc-agent-go/session" + "trpc.group/trpc-go/trpc-agent-go/telemetry/semconv/metrics" + semconvtrace "trpc.group/trpc-go/trpc-agent-go/telemetry/semconv/trace" "trpc.group/trpc-go/trpc-agent-go/tool" ) @@ -590,9 +596,14 @@ func TestPolicyRejectsMissingRuntimeIdentity(t *testing.T) { } func TestPolicyReturnsAuditSinkErrors(t *testing.T) { + reader, restore := usePolicyAuditMetrics(t) + defer restore() + p := newPolicy( t, platform.ToolPolicy{ + TenantID: "tenant", + AppID: "app", DangerousToolAction: platform.DangerousToolActionAllowWithAudit, HighRiskTools: []string{"http_post"}, }, @@ -606,6 +617,14 @@ func TestPolicyReturnsAuditSinkErrors(t *testing.T) { if err == nil || !strings.Contains(err.Error(), "write tool policy audit") { t.Fatalf("expected audit sink error, got %v", err) } + + points := collectPolicyAuditWriteFailedPoints(t, reader) + require.Len(t, points, 1) + require.Equal(t, int64(1), points[0].Value) + requirePolicyAuditMetricAttr(t, points[0].Attributes, semconvtrace.KeyGenAIOperationName, itelemetry.OperationAuditWrite) + requirePolicyAuditMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoTenantID, "tenant") + requirePolicyAuditMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoAppName, "app") + requirePolicyAuditMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoAuditDecision, string(tool.PermissionActionAllow)) } func TestApprovalOptionsMapPolicy(t *testing.T) { @@ -907,3 +926,57 @@ type failingAuditSink struct{} func (failingAuditSink) WriteAudit(context.Context, platform.AuditRecord) error { return errors.New("audit unavailable") } + +func usePolicyAuditMetrics(t *testing.T) (*sdkmetric.ManualReader, func()) { + t.Helper() + reader := sdkmetric.NewManualReader() + provider := sdkmetric.NewMeterProvider(sdkmetric.WithReader(reader)) + + originalProvider := itelemetry.MeterProvider + originalMeter := itelemetry.AuditMeter + originalCounter := itelemetry.AuditMetricWriteFailedTotal + + itelemetry.MeterProvider = provider + itelemetry.AuditMeter = provider.Meter(metrics.MeterNameAudit) + var err error + itelemetry.AuditMetricWriteFailedTotal, err = itelemetry.AuditMeter.Int64Counter(metrics.MetricAuditWriteFailedTotal) + require.NoError(t, err) + + return reader, func() { + itelemetry.MeterProvider = originalProvider + itelemetry.AuditMeter = originalMeter + itelemetry.AuditMetricWriteFailedTotal = originalCounter + } +} + +func collectPolicyAuditWriteFailedPoints( + t *testing.T, + reader *sdkmetric.ManualReader, +) []metricdata.DataPoint[int64] { + t.Helper() + var rm metricdata.ResourceMetrics + require.NoError(t, reader.Collect(context.Background(), &rm)) + for _, scopeMetric := range rm.ScopeMetrics { + for _, metric := range scopeMetric.Metrics { + if metric.Name != metrics.MetricAuditWriteFailedTotal { + continue + } + sum, ok := metric.Data.(metricdata.Sum[int64]) + require.True(t, ok) + return sum.DataPoints + } + } + t.Fatalf("metric %s not found", metrics.MetricAuditWriteFailedTotal) + return nil +} + +func requirePolicyAuditMetricAttr(t *testing.T, set attribute.Set, key string, value string) { + t.Helper() + for _, kv := range set.ToSlice() { + if string(kv.Key) == key { + require.Equal(t, value, kv.Value.AsString()) + return + } + } + t.Fatalf("attribute %s not found", key) +} diff --git a/plugin/guardrail/approval/approval_test.go b/plugin/guardrail/approval/approval_test.go index a79e40cc9d..90313923ac 100644 --- a/plugin/guardrail/approval/approval_test.go +++ b/plugin/guardrail/approval/approval_test.go @@ -531,6 +531,46 @@ func TestBeforeTool_RequireApprovalWritesRejectedAuditRecords(t *testing.T) { assert.Equal(t, "tool approval rejected", records[1].DecisionReason) } +func TestBeforeTool_RequireApprovalRecordsAuditWriteFailureMetric(t *testing.T) { + reader, restore := useApprovalAuditMetrics(t) + defer restore() + + p, err := New( + WithReviewer(&stubReviewer{ + reviewFn: func(ctx context.Context, req *approvalreview.Request) (*approvalreview.Decision, error) { + return &approvalreview.Decision{Approved: true, RiskLevel: "low", Reason: "ok"}, nil + }, + }), + WithAuditSink(failingApprovalAuditSink{}), + WithApproverUserID("security@example.com"), + ) + require.NoError(t, err) + + callbacks := registeredToolCallbacks(t, p) + ctx := ContextWithAuditContext(context.Background(), AuditContext{ + TenantID: "tenant", + AppID: "app", + RequestID: "request-1", + TraceID: "trace-1", + }) + result, runErr := callbacks.RunBeforeTool(ctx, &tool.BeforeToolArgs{ + ToolName: "shell", + ToolCallID: "call-1", + Arguments: []byte(`{"command":"pwd"}`), + }) + require.NoError(t, runErr) + require.NotNil(t, result) + require.Equal(t, `approval audit failed for tool "shell": audit unavailable`, result.CustomResult) + + points := approvalMetricSumPoints(t, collectApprovalAuditMetrics(t, reader), metrics.MetricAuditWriteFailedTotal) + require.Len(t, points, 1) + require.Equal(t, int64(1), points[0].Value) + requireApprovalMetricAttr(t, points[0].Attributes, semconvtrace.KeyGenAIOperationName, itelemetry.OperationAuditWrite) + requireApprovalMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoTenantID, "tenant") + requireApprovalMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoAppName, "app") + requireApprovalMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoAuditDecision, string(platform.ToolApprovalDecisionRequested)) +} + func TestBeforeTool_ReviewerErrorFailsClosed(t *testing.T) { original := approvallog.ErrorfContext var errorLog string @@ -829,3 +869,38 @@ func requireApprovalMetricAttr(t *testing.T, set attribute.Set, key string, valu } t.Fatalf("attribute %s not found", key) } + +type failingApprovalAuditSink struct{} + +func (failingApprovalAuditSink) WriteAudit(context.Context, platform.AuditRecord) error { + return errors.New("audit unavailable") +} + +func useApprovalAuditMetrics(t *testing.T) (*sdkmetric.ManualReader, func()) { + t.Helper() + reader := sdkmetric.NewManualReader() + provider := sdkmetric.NewMeterProvider(sdkmetric.WithReader(reader)) + + originalProvider := itelemetry.MeterProvider + originalMeter := itelemetry.AuditMeter + originalCounter := itelemetry.AuditMetricWriteFailedTotal + + itelemetry.MeterProvider = provider + itelemetry.AuditMeter = provider.Meter(metrics.MeterNameAudit) + var err error + itelemetry.AuditMetricWriteFailedTotal, err = itelemetry.AuditMeter.Int64Counter(metrics.MetricAuditWriteFailedTotal) + require.NoError(t, err) + + return reader, func() { + itelemetry.MeterProvider = originalProvider + itelemetry.AuditMeter = originalMeter + itelemetry.AuditMetricWriteFailedTotal = originalCounter + } +} + +func collectApprovalAuditMetrics(t *testing.T, reader *sdkmetric.ManualReader) metricdata.ResourceMetrics { + t.Helper() + var rm metricdata.ResourceMetrics + require.NoError(t, reader.Collect(context.Background(), &rm)) + return rm +} diff --git a/plugin/guardrail/approval/audit.go b/plugin/guardrail/approval/audit.go index 57967d5da0..9866ee4aa7 100644 --- a/plugin/guardrail/approval/audit.go +++ b/plugin/guardrail/approval/audit.go @@ -120,7 +120,16 @@ func (p *Plugin) writeApprovalAudit( if err != nil { return err } - return p.auditSink.WriteAudit(ctx, record) + if err := p.auditSink.WriteAudit(ctx, record); err != nil { + itelemetry.ReportAuditWriteFailedMetrics(ctx, itelemetry.AuditAttributes{ + TenantID: record.TenantID, + AppName: record.AppID, + Decision: record.Decision, + Error: err, + }) + return err + } + return nil } func reportApprovalRequiredMetric(ctx context.Context, args *tool.BeforeToolArgs) { diff --git a/telemetry/metric/metric.go b/telemetry/metric/metric.go index 9410d2c123..3ce1e9c3ef 100644 --- a/telemetry/metric/metric.go +++ b/telemetry/metric/metric.go @@ -118,6 +118,9 @@ func InitMeterProvider(mp metric.MeterProvider) error { if err := initToolApprovalMetrics(mp); err != nil { return err } + if err := initAuditMetrics(mp); err != nil { + return err + } if err := initInvokeAgentMetrics(mp); err != nil { return err } @@ -218,6 +221,24 @@ func initToolApprovalMetrics(mp metric.MeterProvider) error { return nil } +func initAuditMetrics(mp metric.MeterProvider) error { + if mp == nil { + return fmt.Errorf("audit meter provider is nil") + } + meterName := metrics.MeterNameAudit + itelemetry.AuditMeter = mp.Meter(meterName) + var err error + itelemetry.AuditMetricWriteFailedTotal, err = itelemetry.AuditMeter.Int64Counter( + metrics.MetricAuditWriteFailedTotal, + metric.WithDescription("Total number of failed audit sink writes"), + metric.WithUnit("1"), + ) + if err != nil { + return fmt.Errorf("failed to create %s metric %s: %w", meterName, metrics.MetricAuditWriteFailedTotal, err) + } + return nil +} + func setInvokeAgentHistogramBuckets(metricName string, boundaries []float64) error { switch metricName { case metrics.MetricTRPCAgentGoClientTimeToFirstToken: diff --git a/telemetry/metric/metric_test.go b/telemetry/metric/metric_test.go index 1ca844a014..732a6f0635 100644 --- a/telemetry/metric/metric_test.go +++ b/telemetry/metric/metric_test.go @@ -296,10 +296,14 @@ func TestInitMeterProvider(t *testing.T) { originalMP := itelemetry.MeterProvider originalToolApprovalMeter := itelemetry.ToolApprovalMeter originalToolApprovalRequired := itelemetry.ToolApprovalMetricRequiredTotal + originalAuditMeter := itelemetry.AuditMeter + originalAuditWriteFailed := itelemetry.AuditMetricWriteFailedTotal defer func() { itelemetry.MeterProvider = originalMP itelemetry.ToolApprovalMeter = originalToolApprovalMeter itelemetry.ToolApprovalMetricRequiredTotal = originalToolApprovalRequired + itelemetry.AuditMeter = originalAuditMeter + itelemetry.AuditMetricWriteFailedTotal = originalAuditWriteFailed }() // Create a test meter provider @@ -361,6 +365,12 @@ func TestInitMeterProvider(t *testing.T) { if itelemetry.ToolApprovalMetricRequiredTotal == nil { t.Error("ToolApprovalMetricRequiredTotal was not created") } + if itelemetry.AuditMeter == nil { + t.Error("AuditMeter was not created") + } + if itelemetry.AuditMetricWriteFailedTotal == nil { + t.Error("AuditMetricWriteFailedTotal was not created") + } if itelemetry.WorkflowMeter == nil { t.Error("WorkflowMeter was not created") } @@ -426,6 +436,31 @@ func TestInitToolApprovalMetrics_ErrorHandling(t *testing.T) { } } +func TestInitAuditMetrics_ErrorHandling(t *testing.T) { + originalAuditMeter := itelemetry.AuditMeter + originalAuditWriteFailed := itelemetry.AuditMetricWriteFailedTotal + defer func() { + itelemetry.AuditMeter = originalAuditMeter + itelemetry.AuditMetricWriteFailedTotal = originalAuditWriteFailed + }() + + if err := initAuditMetrics(nil); err == nil || !strings.Contains(err.Error(), "audit meter provider is nil") { + t.Fatalf("expected nil provider error, got %v", err) + } + + mp := &mockMeterProvider{meter: &mockMeter{ + shouldFail: true, + failOn: metrics.MetricAuditWriteFailedTotal, + }} + err := initAuditMetrics(mp) + if err == nil { + t.Fatalf("expected audit counter creation error") + } + if !strings.Contains(err.Error(), "failed to create trpc_agent_go.internal.audit metric audit_write_failed_total") { + t.Fatalf("unexpected error: %v", err) + } +} + func TestInitWorkflowMetrics_ErrorHandling(t *testing.T) { originalWorkflowMeter := itelemetry.WorkflowMeter originalWorkflowOpDur := itelemetry.WorkflowMetricGenAIClientOperationDuration diff --git a/telemetry/semconv/metrics/metrics.go b/telemetry/semconv/metrics/metrics.go index 00f2c10864..b1d7c4f3a4 100644 --- a/telemetry/semconv/metrics/metrics.go +++ b/telemetry/semconv/metrics/metrics.go @@ -61,6 +61,8 @@ const ( MetricTRPCAgentGoClientRequestCnt = "trpc_agent_go.client.request_cnt" // MetricToolApprovalRequiredTotal records tool calls that require explicit approval. MetricToolApprovalRequiredTotal = "tool_approval_required_total" + // MetricAuditWriteFailedTotal records failed audit sink writes. + MetricAuditWriteFailedTotal = "audit_write_failed_total" ////////////////////////// server //////////////////////// @@ -83,4 +85,6 @@ const ( MeterNameInvokeAgent = "trpc_agent_go.internal.invoke_agent" // MeterNameToolApproval is the meter name for tool approval operations. MeterNameToolApproval = "trpc_agent_go.internal.tool_approval" + // MeterNameAudit is the meter name for audit operations. + MeterNameAudit = "trpc_agent_go.internal.audit" ) diff --git a/telemetry/semconv/trace/trace.go b/telemetry/semconv/trace/trace.go index d31ff1b203..fddda0826f 100644 --- a/telemetry/semconv/trace/trace.go +++ b/telemetry/semconv/trace/trace.go @@ -60,6 +60,8 @@ const ( KeyTRPCAgentGoMemorySearchDeduplicate = "trpc.go.agent.memory.search.deduplicate" // KeyTRPCAgentGoMemoryWriteOperation is the memory write operation type. KeyTRPCAgentGoMemoryWriteOperation = "trpc.go.agent.memory.write.operation" + // KeyTRPCAgentGoAuditDecision is the audit record decision. + KeyTRPCAgentGoAuditDecision = "trpc.go.agent.audit.decision" // KeyGenAIAppName is the attribute key for GenAI application name. KeyGenAIAppName = "gen_ai.app.name" From 75fb2395ce27d23b1dba9b4e9c55f344b9863163 Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 16:22:09 +0800 Subject: [PATCH 89/95] fix runtime builder options calls in approval tests --- platform/worker/governance_test.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/platform/worker/governance_test.go b/platform/worker/governance_test.go index 1d31991aab..bacb19d867 100644 --- a/platform/worker/governance_test.go +++ b/platform/worker/governance_test.go @@ -120,7 +120,7 @@ func TestRuntimeBuilderRunsApprovalPluginBeforeMandatoryPolicy(t *testing.T) { DangerousToolAction: platform.DangerousToolActionAsk, } var captured AgentDependencies - builder, err := NewRuntimeBuilder( + builder, err := NewRuntimeBuilderWithOptions( router, AgentFactoryFunc(func( _ context.Context, @@ -235,7 +235,7 @@ func TestRuntimeBuilderRunsApprovalPluginForMetadataRisk(t *testing.T) { DangerousToolAction: platform.DangerousToolActionAsk, } var captured AgentDependencies - builder, err := NewRuntimeBuilder( + builder, err := NewRuntimeBuilderWithOptions( router, AgentFactoryFunc(func( _ context.Context, From f3b3ab113132e57e681c3d7517599f6bbf2ab360 Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 14:48:16 +0800 Subject: [PATCH 90/95] feat(tool): record permission denied metrics --- internal/flow/processor/functioncall.go | 62 ++++++++++ internal/flow/processor/functioncall_test.go | 87 ++++++++++++++ internal/telemetry/metric_execute_tool.go | 45 ++++++++ internal/telemetry/metric_test.go | 112 +++++++++++++++++++ telemetry/metric/metric.go | 7 ++ telemetry/metric/metric_test.go | 32 ++++++ telemetry/semconv/metrics/metrics.go | 2 + telemetry/semconv/trace/trace.go | 2 + 8 files changed, 349 insertions(+) diff --git a/internal/flow/processor/functioncall.go b/internal/flow/processor/functioncall.go index f20e5d0a13..a8556aed7f 100644 --- a/internal/flow/processor/functioncall.go +++ b/internal/flow/processor/functioncall.go @@ -1770,6 +1770,7 @@ func (p *FunctionCallResponseProcessor) executeToolCall( return ctx, nil, modifiedArgs, true, skipSummarization, nil } } + reportToolPermissionDeniedMetric(ctx, invocation, toolCall, result) if suppressDefaultToolMessage { defaultMsg, err := buildDefaultToolMessage(toolCall.ID, result) if err != nil { @@ -1885,6 +1886,67 @@ func isPermissionResult(result any) bool { } } +func reportToolPermissionDeniedMetric( + ctx context.Context, + invocation *agent.Invocation, + toolCall model.ToolCall, + result any, +) { + status, ok := permissionDeniedMetricStatus(result) + if !ok { + return + } + var ( + sess = &session.Session{} + modelName string + agentName string + ) + if invocation != nil { + if invocation.Session != nil { + sess = invocation.Session + } + if invocation.Model != nil { + modelName = invocation.Model.Info().Name + } + if invocation.AgentName != "" { + agentName = invocation.AgentName + } + } + itelemetry.ReportToolPermissionDeniedMetrics(ctx, itelemetry.ToolPermissionDeniedAttributes{ + RequestModelName: modelName, + ToolName: toolCall.Function.Name, + AppName: sess.AppName, + UserID: sess.UserID, + SessionID: sess.ID, + AgentName: agentName, + Status: status, + }) +} + +func permissionDeniedMetricStatus(result any) (string, bool) { + switch v := result.(type) { + case tool.PermissionResult: + return permissionDeniedMetricStatusValue(v.Status) + case *tool.PermissionResult: + if v == nil { + return "", false + } + return permissionDeniedMetricStatusValue(v.Status) + default: + return "", false + } +} + +func permissionDeniedMetricStatusValue(status string) (string, bool) { + switch status { + case tool.PermissionResultStatusDenied, + tool.PermissionResultStatusApprovalDenied: + return status, true + default: + return "", false + } +} + func isPermissionResultStatus(status string) bool { switch status { case tool.PermissionResultStatusDenied, diff --git a/internal/flow/processor/functioncall_test.go b/internal/flow/processor/functioncall_test.go index ce9d757095..6c5e873641 100644 --- a/internal/flow/processor/functioncall_test.go +++ b/internal/flow/processor/functioncall_test.go @@ -24,17 +24,22 @@ import ( "github.com/google/uuid" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "go.opentelemetry.io/otel/attribute" + sdkmetric "go.opentelemetry.io/otel/sdk/metric" + "go.opentelemetry.io/otel/sdk/metric/metricdata" sdktrace "go.opentelemetry.io/otel/sdk/trace" "go.opentelemetry.io/otel/sdk/trace/tracetest" "trpc.group/trpc-go/trpc-agent-go/agent" "trpc.group/trpc-go/trpc-agent-go/event" "trpc.group/trpc-go/trpc-agent-go/graph" "trpc.group/trpc-go/trpc-agent-go/internal/state/appender" + itelemetry "trpc.group/trpc-go/trpc-agent-go/internal/telemetry" itool "trpc.group/trpc-go/trpc-agent-go/internal/tool" "trpc.group/trpc-go/trpc-agent-go/model" "trpc.group/trpc-go/trpc-agent-go/plugin" "trpc.group/trpc-go/trpc-agent-go/session" skillstate "trpc.group/trpc-go/trpc-agent-go/skill" + "trpc.group/trpc-go/trpc-agent-go/telemetry/semconv/metrics" semconvtrace "trpc.group/trpc-go/trpc-agent-go/telemetry/semconv/trace" "trpc.group/trpc-go/trpc-agent-go/telemetry/trace" "trpc.group/trpc-go/trpc-agent-go/tool" @@ -9089,6 +9094,7 @@ func TestExecuteToolCall_ToolPermissionResultSkipsToolResultMessagesCallback( denyReason = "write access is disabled" permissionJSON = `{"status":"denied","tool":"delete_file","reason":"write access is disabled"}` ) + reader := setupToolPermissionDeniedMetric(t) var ( calledTool bool @@ -9143,6 +9149,7 @@ func TestExecuteToolCall_ToolPermissionResultSkipsToolResultMessagesCallback( require.Equal(t, toolCallID, choices[0].Message.ToolID) require.Equal(t, toolName, choices[0].Message.ToolName) require.JSONEq(t, permissionJSON, choices[0].Message.Content) + requireToolPermissionDeniedMetric(t, reader, tool.PermissionResultStatusDenied) } func TestExecuteToolCall_ApprovalDeniedSkipsToolResultMessagesCallback( @@ -9154,6 +9161,7 @@ func TestExecuteToolCall_ApprovalDeniedSkipsToolResultMessagesCallback( denyReason = "Automatic approval review denied (risk: high): write access is disabled" permissionJSON = `{"status":"approval_denied","tool":"delete_file","reason":"Automatic approval review denied (risk: high): write access is disabled"}` ) + reader := setupToolPermissionDeniedMetric(t) var ( calledTool bool @@ -9209,6 +9217,7 @@ func TestExecuteToolCall_ApprovalDeniedSkipsToolResultMessagesCallback( require.Equal(t, toolCallID, choices[0].Message.ToolID) require.Equal(t, toolName, choices[0].Message.ToolName) require.JSONEq(t, permissionJSON, choices[0].Message.Content) + requireToolPermissionDeniedMetric(t, reader, tool.PermissionResultStatusApprovalDenied) } func TestExecuteToolWithCallbacks_MandatoryPermissionDenyCannotBeOverridden( @@ -9373,6 +9382,7 @@ func TestExecuteToolWithCallbacks_ToolPermissionCheckerAskSkipsRunPolicy( askReason = "shell commands require review" permissionJSON = `{"status":"approval_required","tool":"shell","reason":"shell commands require review"}` ) + reader := setupToolPermissionDeniedMetric(t) var ( calledTool bool @@ -9415,6 +9425,83 @@ func TestExecuteToolWithCallbacks_ToolPermissionCheckerAskSkipsRunPolicy( require.False(t, calledTool) require.False(t, calledRunPolicy) require.JSONEq(t, permissionJSON, string(mustJSON(res))) + requireNoToolPermissionDeniedMetric(t, reader) +} + +func setupToolPermissionDeniedMetric(t *testing.T) *sdkmetric.ManualReader { + t.Helper() + + reader := sdkmetric.NewManualReader() + provider := sdkmetric.NewMeterProvider(sdkmetric.WithReader(reader)) + + originalProvider := itelemetry.MeterProvider + originalMeter := itelemetry.ExecuteToolMeter + originalCounter := itelemetry.ExecuteToolMetricToolPermissionDeniedTotal + t.Cleanup(func() { + itelemetry.MeterProvider = originalProvider + itelemetry.ExecuteToolMeter = originalMeter + itelemetry.ExecuteToolMetricToolPermissionDeniedTotal = originalCounter + }) + + itelemetry.MeterProvider = provider + itelemetry.ExecuteToolMeter = provider.Meter(metrics.MeterNameExecuteTool) + var err error + itelemetry.ExecuteToolMetricToolPermissionDeniedTotal, err = + itelemetry.ExecuteToolMeter.Int64Counter(metrics.MetricToolPermissionDeniedTotal) + require.NoError(t, err) + return reader +} + +func requireToolPermissionDeniedMetric( + t *testing.T, + reader *sdkmetric.ManualReader, + status string, +) { + t.Helper() + + points := collectToolPermissionDeniedMetricPoints(t, reader) + require.Len(t, points, 1) + require.Equal(t, int64(1), points[0].Value) + requireProcessorMetricAttr(t, points[0].Attributes, semconvtrace.KeyGenAIOperationName, itelemetry.OperationExecuteTool) + requireProcessorMetricAttr(t, points[0].Attributes, semconvtrace.KeyGenAIToolName, "delete_file") + requireProcessorMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoToolPermissionStatus, status) +} + +func requireNoToolPermissionDeniedMetric(t *testing.T, reader *sdkmetric.ManualReader) { + t.Helper() + require.Empty(t, collectToolPermissionDeniedMetricPoints(t, reader)) +} + +func collectToolPermissionDeniedMetricPoints( + t *testing.T, + reader *sdkmetric.ManualReader, +) []metricdata.DataPoint[int64] { + t.Helper() + + var rm metricdata.ResourceMetrics + require.NoError(t, reader.Collect(context.Background(), &rm)) + for _, scopeMetric := range rm.ScopeMetrics { + for _, metric := range scopeMetric.Metrics { + if metric.Name != metrics.MetricToolPermissionDeniedTotal { + continue + } + sum, ok := metric.Data.(metricdata.Sum[int64]) + require.True(t, ok) + return sum.DataPoints + } + } + return nil +} + +func requireProcessorMetricAttr(t *testing.T, set attribute.Set, key string, value string) { + t.Helper() + for _, kv := range set.ToSlice() { + if string(kv.Key) == key { + require.Equal(t, value, kv.Value.AsString()) + return + } + } + t.Fatalf("attribute %s not found", key) } func TestExecuteToolWithCallbacks_ToolPermissionReceivesMetadata( diff --git a/internal/telemetry/metric_execute_tool.go b/internal/telemetry/metric_execute_tool.go index 624f47c3ca..6b5fbab6a3 100644 --- a/internal/telemetry/metric_execute_tool.go +++ b/internal/telemetry/metric_execute_tool.go @@ -26,6 +26,8 @@ var ( // ExecuteToolMetricTRPCAgentGoClientRequestCnt records the number of tool execution requests made. ExecuteToolMetricTRPCAgentGoClientRequestCnt metric.Int64Counter + // ExecuteToolMetricToolPermissionDeniedTotal records tool calls denied before execution. + ExecuteToolMetricToolPermissionDeniedTotal metric.Int64Counter // ExecuteToolMetricGenAIClientOperationDuration records the distribution of tool execution durations in seconds. ExecuteToolMetricGenAIClientOperationDuration *histogram.DynamicFloat64Histogram ) @@ -42,6 +44,17 @@ type ExecuteToolAttributes struct { ErrorType string } +// ToolPermissionDeniedAttributes is the attributes for tool permission denial metrics. +type ToolPermissionDeniedAttributes struct { + RequestModelName string + ToolName string + AppName string + AgentName string + UserID string + SessionID string + Status string +} + func (a ExecuteToolAttributes) toAttributes() []attribute.KeyValue { attrs := []attribute.KeyValue{ attribute.String(semconvtrace.KeyGenAIOperationName, OperationExecuteTool), @@ -68,6 +81,30 @@ func (a ExecuteToolAttributes) toAttributes() []attribute.KeyValue { return attrs } +func (a ToolPermissionDeniedAttributes) toAttributes() []attribute.KeyValue { + attrs := []attribute.KeyValue{ + attribute.String(semconvtrace.KeyGenAIOperationName, OperationExecuteTool), + attribute.String(semconvtrace.KeyGenAISystem, a.RequestModelName), + attribute.String(semconvtrace.KeyGenAIToolName, a.ToolName), + } + if a.Status != "" { + attrs = append(attrs, attribute.String(semconvtrace.KeyTRPCAgentGoToolPermissionStatus, a.Status)) + } + if a.AppName != "" { + attrs = append(attrs, attribute.String(semconvtrace.KeyTRPCAgentGoAppName, a.AppName)) + } + if a.UserID != "" { + attrs = append(attrs, attribute.String(semconvtrace.KeyTRPCAgentGoUserID, a.UserID)) + } + if a.SessionID != "" { + attrs = append(attrs, attribute.String(semconvtrace.KeyGenAIConversationID, a.SessionID)) + } + if a.AgentName != "" { + attrs = append(attrs, attribute.String(semconvtrace.KeyGenAIAgentName, a.AgentName)) + } + return attrs +} + // ReportExecuteToolMetrics reports the tool execution metrics. func ReportExecuteToolMetrics(ctx context.Context, attrs ExecuteToolAttributes, duration time.Duration) { as := attrs.toAttributes() @@ -78,3 +115,11 @@ func ReportExecuteToolMetrics(ctx context.Context, attrs ExecuteToolAttributes, ExecuteToolMetricGenAIClientOperationDuration.Record(ctx, duration.Seconds(), metric.WithAttributes(as...)) } } + +// ReportToolPermissionDeniedMetrics reports tool calls denied before execution. +func ReportToolPermissionDeniedMetrics(ctx context.Context, attrs ToolPermissionDeniedAttributes) { + if ExecuteToolMetricToolPermissionDeniedTotal == nil { + return + } + ExecuteToolMetricToolPermissionDeniedTotal.Add(ctx, 1, metric.WithAttributes(attrs.toAttributes()...)) +} diff --git a/internal/telemetry/metric_test.go b/internal/telemetry/metric_test.go index b347985174..e6aa2299c3 100644 --- a/internal/telemetry/metric_test.go +++ b/internal/telemetry/metric_test.go @@ -811,3 +811,115 @@ func TestReportExecuteToolMetrics(t *testing.T) { t.Error("expected metrics to be recorded") } } + +func TestReportToolPermissionDeniedMetricsNoopWhenCounterNil(t *testing.T) { + originalCounter := ExecuteToolMetricToolPermissionDeniedTotal + t.Cleanup(func() { + ExecuteToolMetricToolPermissionDeniedTotal = originalCounter + }) + + ExecuteToolMetricToolPermissionDeniedTotal = nil + if panicked := func() (panicked bool) { + defer func() { + panicked = recover() != nil + }() + ReportToolPermissionDeniedMetrics(context.Background(), ToolPermissionDeniedAttributes{ + RequestModelName: "gpt-4", + ToolName: "shell", + Status: "denied", + }) + return false + }(); panicked { + t.Fatal("ReportToolPermissionDeniedMetrics should not panic when counter is nil") + } +} + +func TestReportToolPermissionDeniedMetrics(t *testing.T) { + reader := sdkmetric.NewManualReader() + provider := sdkmetric.NewMeterProvider(sdkmetric.WithReader(reader)) + + originalProvider := MeterProvider + originalMeter := ExecuteToolMeter + originalCounter := ExecuteToolMetricToolPermissionDeniedTotal + t.Cleanup(func() { + MeterProvider = originalProvider + ExecuteToolMeter = originalMeter + ExecuteToolMetricToolPermissionDeniedTotal = originalCounter + }) + + MeterProvider = provider + ExecuteToolMeter = provider.Meter(metrics.MeterNameExecuteTool) + var err error + ExecuteToolMetricToolPermissionDeniedTotal, err = + ExecuteToolMeter.Int64Counter(metrics.MetricToolPermissionDeniedTotal) + if err != nil { + t.Fatalf("failed to create counter: %v", err) + } + + ctx := context.Background() + ReportToolPermissionDeniedMetrics(ctx, ToolPermissionDeniedAttributes{ + RequestModelName: "gpt-4", + ToolName: "shell", + AppName: "test-app", + UserID: "user-1", + SessionID: "session-1", + AgentName: "agent-1", + Status: "approval_denied", + }) + + var rm metricdata.ResourceMetrics + if err := reader.Collect(ctx, &rm); err != nil { + t.Fatalf("failed to collect metrics: %v", err) + } + + points := executeToolSumPoints(t, rm, metrics.MetricToolPermissionDeniedTotal) + if len(points) != 1 { + t.Fatalf("expected 1 metric point, got %d", len(points)) + } + if points[0].Value != 1 { + t.Fatalf("expected metric value 1, got %d", points[0].Value) + } + requireMetricAttr(t, points[0].Attributes, semconvtrace.KeyGenAIOperationName, OperationExecuteTool) + requireMetricAttr(t, points[0].Attributes, semconvtrace.KeyGenAISystem, "gpt-4") + requireMetricAttr(t, points[0].Attributes, semconvtrace.KeyGenAIToolName, "shell") + requireMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoToolPermissionStatus, "approval_denied") + requireMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoAppName, "test-app") + requireMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoUserID, "user-1") + requireMetricAttr(t, points[0].Attributes, semconvtrace.KeyGenAIConversationID, "session-1") + requireMetricAttr(t, points[0].Attributes, semconvtrace.KeyGenAIAgentName, "agent-1") +} + +func executeToolSumPoints( + t *testing.T, + rm metricdata.ResourceMetrics, + metricName string, +) []metricdata.DataPoint[int64] { + t.Helper() + for _, scopeMetric := range rm.ScopeMetrics { + for _, metric := range scopeMetric.Metrics { + if metric.Name != metricName { + continue + } + sum, ok := metric.Data.(metricdata.Sum[int64]) + if !ok { + t.Fatalf("metric %s has unexpected data type %T", metricName, metric.Data) + } + return sum.DataPoints + } + } + t.Fatalf("metric %s not found", metricName) + return nil +} + +func requireMetricAttr(t *testing.T, set attribute.Set, key string, value string) { + t.Helper() + for _, kv := range set.ToSlice() { + if string(kv.Key) == key { + if kv.Value.AsString() != value { + t.Fatalf("attribute %s: expected %q, got %q", key, value, kv.Value.AsString()) + } + return + } + } + t.Fatalf("attribute %s not found", key) +} diff --git a/telemetry/metric/metric.go b/telemetry/metric/metric.go index 3ce1e9c3ef..e48e530556 100644 --- a/telemetry/metric/metric.go +++ b/telemetry/metric/metric.go @@ -105,6 +105,13 @@ func InitMeterProvider(mp metric.MeterProvider) error { ); err != nil { return fmt.Errorf("failed to create execute tool metric TRPCAgentGoClientRequestCnt: %w", err) } + if itelemetry.ExecuteToolMetricToolPermissionDeniedTotal, err = itelemetry.ExecuteToolMeter.Int64Counter( + metrics.MetricToolPermissionDeniedTotal, + metric.WithDescription("Total number of tool calls denied before execution"), + metric.WithUnit("1"), + ); err != nil { + return fmt.Errorf("failed to create execute tool metric ToolPermissionDeniedTotal: %w", err) + } if itelemetry.ExecuteToolMetricGenAIClientOperationDuration, err = histogram.NewDynamicFloat64Histogram( mp, metrics.MeterNameExecuteTool, diff --git a/telemetry/metric/metric_test.go b/telemetry/metric/metric_test.go index 732a6f0635..3c9cd36e30 100644 --- a/telemetry/metric/metric_test.go +++ b/telemetry/metric/metric_test.go @@ -296,12 +296,14 @@ func TestInitMeterProvider(t *testing.T) { originalMP := itelemetry.MeterProvider originalToolApprovalMeter := itelemetry.ToolApprovalMeter originalToolApprovalRequired := itelemetry.ToolApprovalMetricRequiredTotal + originalToolPermissionDenied := itelemetry.ExecuteToolMetricToolPermissionDeniedTotal originalAuditMeter := itelemetry.AuditMeter originalAuditWriteFailed := itelemetry.AuditMetricWriteFailedTotal defer func() { itelemetry.MeterProvider = originalMP itelemetry.ToolApprovalMeter = originalToolApprovalMeter itelemetry.ToolApprovalMetricRequiredTotal = originalToolApprovalRequired + itelemetry.ExecuteToolMetricToolPermissionDeniedTotal = originalToolPermissionDenied itelemetry.AuditMeter = originalAuditMeter itelemetry.AuditMetricWriteFailedTotal = originalAuditWriteFailed }() @@ -356,6 +358,9 @@ func TestInitMeterProvider(t *testing.T) { if itelemetry.ExecuteToolMetricTRPCAgentGoClientRequestCnt == nil { t.Error("ExecuteToolMetricTRPCAgentGoClientRequestCnt was not created") } + if itelemetry.ExecuteToolMetricToolPermissionDeniedTotal == nil { + t.Error("ExecuteToolMetricToolPermissionDeniedTotal was not created") + } if itelemetry.ExecuteToolMetricGenAIClientOperationDuration == nil { t.Error("ExecuteToolMetricGenAIClientOperationDuration was not created") } @@ -436,6 +441,33 @@ func TestInitToolApprovalMetrics_ErrorHandling(t *testing.T) { } } +func TestInitMeterProvider_ToolPermissionDeniedMetricError(t *testing.T) { + originalMP := itelemetry.MeterProvider + originalExecuteToolMeter := itelemetry.ExecuteToolMeter + originalExecuteToolRequestCnt := itelemetry.ExecuteToolMetricTRPCAgentGoClientRequestCnt + originalToolPermissionDenied := itelemetry.ExecuteToolMetricToolPermissionDeniedTotal + originalExecuteToolDuration := itelemetry.ExecuteToolMetricGenAIClientOperationDuration + t.Cleanup(func() { + itelemetry.MeterProvider = originalMP + itelemetry.ExecuteToolMeter = originalExecuteToolMeter + itelemetry.ExecuteToolMetricTRPCAgentGoClientRequestCnt = originalExecuteToolRequestCnt + itelemetry.ExecuteToolMetricToolPermissionDeniedTotal = originalToolPermissionDenied + itelemetry.ExecuteToolMetricGenAIClientOperationDuration = originalExecuteToolDuration + }) + + mp := &mockMeterProvider{meter: &mockMeter{ + shouldFail: true, + failOn: metrics.MetricToolPermissionDeniedTotal, + }} + err := InitMeterProvider(mp) + if err == nil { + t.Fatalf("expected tool permission denied counter creation error") + } + if !strings.Contains(err.Error(), "failed to create execute tool metric ToolPermissionDeniedTotal") { + t.Fatalf("unexpected error: %v", err) + } +} + func TestInitAuditMetrics_ErrorHandling(t *testing.T) { originalAuditMeter := itelemetry.AuditMeter originalAuditWriteFailed := itelemetry.AuditMetricWriteFailedTotal diff --git a/telemetry/semconv/metrics/metrics.go b/telemetry/semconv/metrics/metrics.go index b1d7c4f3a4..6acebb4b39 100644 --- a/telemetry/semconv/metrics/metrics.go +++ b/telemetry/semconv/metrics/metrics.go @@ -61,6 +61,8 @@ const ( MetricTRPCAgentGoClientRequestCnt = "trpc_agent_go.client.request_cnt" // MetricToolApprovalRequiredTotal records tool calls that require explicit approval. MetricToolApprovalRequiredTotal = "tool_approval_required_total" + // MetricToolPermissionDeniedTotal records tool calls denied by permission checks. + MetricToolPermissionDeniedTotal = "tool_permission_denied_total" // MetricAuditWriteFailedTotal records failed audit sink writes. MetricAuditWriteFailedTotal = "audit_write_failed_total" diff --git a/telemetry/semconv/trace/trace.go b/telemetry/semconv/trace/trace.go index fddda0826f..dead38d66a 100644 --- a/telemetry/semconv/trace/trace.go +++ b/telemetry/semconv/trace/trace.go @@ -62,6 +62,8 @@ const ( KeyTRPCAgentGoMemoryWriteOperation = "trpc.go.agent.memory.write.operation" // KeyTRPCAgentGoAuditDecision is the audit record decision. KeyTRPCAgentGoAuditDecision = "trpc.go.agent.audit.decision" + // KeyTRPCAgentGoToolPermissionStatus is the structured tool permission result status. + KeyTRPCAgentGoToolPermissionStatus = "trpc.go.agent.tool.permission.status" // KeyGenAIAppName is the attribute key for GenAI application name. KeyGenAIAppName = "gen_ai.app.name" From 97a5fa95649964010db113df78585be8349e5607 Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 14:59:03 +0800 Subject: [PATCH 91/95] feat(gateway): record budget denied metrics --- internal/telemetry/metric_gateway.go | 63 ++++++++++++ internal/telemetry/metric_gateway_test.go | 111 ++++++++++++++++++++++ internal/telemetry/trace.go | 1 + platform/gateway/service.go | 9 ++ platform/gateway/service_test.go | 68 +++++++++++++ telemetry/metric/metric.go | 21 ++++ telemetry/metric/metric_test.go | 35 +++++++ telemetry/semconv/metrics/metrics.go | 4 + telemetry/semconv/trace/trace.go | 4 + 9 files changed, 316 insertions(+) create mode 100644 internal/telemetry/metric_gateway.go create mode 100644 internal/telemetry/metric_gateway_test.go diff --git a/internal/telemetry/metric_gateway.go b/internal/telemetry/metric_gateway.go new file mode 100644 index 0000000000..dbfc1ae5f2 --- /dev/null +++ b/internal/telemetry/metric_gateway.go @@ -0,0 +1,63 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package telemetry + +import ( + "context" + + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/metric" + + "trpc.group/trpc-go/trpc-agent-go/telemetry/semconv/metrics" + semconvtrace "trpc.group/trpc-go/trpc-agent-go/telemetry/semconv/trace" +) + +var ( + // GatewayMeter is the meter used for recording gateway metrics. + GatewayMeter = MeterProvider.Meter(metrics.MeterNameGateway) + + // GatewayMetricBudgetDeniedTotal records gateway requests denied by budget checks. + GatewayMetricBudgetDeniedTotal metric.Int64Counter +) + +// GatewayBudgetDeniedAttributes is the attributes for gateway budget denial metrics. +type GatewayBudgetDeniedAttributes struct { + TenantID string + AppName string + Channel string + Reason string +} + +func (a GatewayBudgetDeniedAttributes) toAttributes() []attribute.KeyValue { + attrs := []attribute.KeyValue{ + attribute.String(semconvtrace.KeyGenAIOperationName, OperationGatewayBudget), + attribute.String(semconvtrace.KeyGenAISystem, semconvtrace.SystemTRPCGoAgent), + } + if a.TenantID != "" { + attrs = append(attrs, attribute.String(semconvtrace.KeyTRPCAgentGoTenantID, a.TenantID)) + } + if a.AppName != "" { + attrs = append(attrs, attribute.String(semconvtrace.KeyTRPCAgentGoAppName, a.AppName)) + } + if a.Channel != "" { + attrs = append(attrs, attribute.String(semconvtrace.KeyTRPCAgentGoChannel, a.Channel)) + } + if a.Reason != "" { + attrs = append(attrs, attribute.String(semconvtrace.KeyTRPCAgentGoBudgetDeniedReason, a.Reason)) + } + return attrs +} + +// ReportGatewayBudgetDeniedMetrics reports a gateway request denied by budget checks. +func ReportGatewayBudgetDeniedMetrics(ctx context.Context, attrs GatewayBudgetDeniedAttributes) { + if GatewayMetricBudgetDeniedTotal == nil { + return + } + GatewayMetricBudgetDeniedTotal.Add(ctx, 1, metric.WithAttributes(attrs.toAttributes()...)) +} diff --git a/internal/telemetry/metric_gateway_test.go b/internal/telemetry/metric_gateway_test.go new file mode 100644 index 0000000000..18fc7c9249 --- /dev/null +++ b/internal/telemetry/metric_gateway_test.go @@ -0,0 +1,111 @@ +// +// Tencent is pleased to support the open source community by making trpc-agent-go available. +// +// Copyright (C) 2025 Tencent. All rights reserved. +// +// trpc-agent-go is licensed under the Apache License Version 2.0. +// + +package telemetry + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" + "go.opentelemetry.io/otel/attribute" + sdkmetric "go.opentelemetry.io/otel/sdk/metric" + "go.opentelemetry.io/otel/sdk/metric/metricdata" + + "trpc.group/trpc-go/trpc-agent-go/telemetry/semconv/metrics" + semconvtrace "trpc.group/trpc-go/trpc-agent-go/telemetry/semconv/trace" +) + +func TestReportGatewayBudgetDeniedMetricsNoopWhenCounterNil(t *testing.T) { + originalCounter := GatewayMetricBudgetDeniedTotal + t.Cleanup(func() { + GatewayMetricBudgetDeniedTotal = originalCounter + }) + + GatewayMetricBudgetDeniedTotal = nil + require.NotPanics(t, func() { + ReportGatewayBudgetDeniedMetrics(context.Background(), GatewayBudgetDeniedAttributes{ + TenantID: "tenant-1", + AppName: "app-1", + Channel: "wecom", + Reason: "total_tokens_exceeded", + }) + }) +} + +func TestReportGatewayBudgetDeniedMetrics(t *testing.T) { + reader := sdkmetric.NewManualReader() + provider := sdkmetric.NewMeterProvider(sdkmetric.WithReader(reader)) + + originalProvider := MeterProvider + originalMeter := GatewayMeter + originalCounter := GatewayMetricBudgetDeniedTotal + t.Cleanup(func() { + MeterProvider = originalProvider + GatewayMeter = originalMeter + GatewayMetricBudgetDeniedTotal = originalCounter + }) + + MeterProvider = provider + GatewayMeter = provider.Meter(metrics.MeterNameGateway) + var err error + GatewayMetricBudgetDeniedTotal, err = GatewayMeter.Int64Counter(metrics.MetricGatewayBudgetDeniedTotal) + require.NoError(t, err) + + ctx := context.Background() + ReportGatewayBudgetDeniedMetrics(ctx, GatewayBudgetDeniedAttributes{ + TenantID: "tenant-1", + AppName: "app-1", + Channel: "wecom", + Reason: "total_tokens_exceeded", + }) + + var rm metricdata.ResourceMetrics + require.NoError(t, reader.Collect(ctx, &rm)) + + points := gatewaySumPoints(t, rm, metrics.MetricGatewayBudgetDeniedTotal) + require.Len(t, points, 1) + require.Equal(t, int64(1), points[0].Value) + requireGatewayAttr(t, points[0].Attributes, semconvtrace.KeyGenAIOperationName, OperationGatewayBudget) + requireGatewayAttr(t, points[0].Attributes, semconvtrace.KeyGenAISystem, semconvtrace.SystemTRPCGoAgent) + requireGatewayAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoTenantID, "tenant-1") + requireGatewayAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoAppName, "app-1") + requireGatewayAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoChannel, "wecom") + requireGatewayAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoBudgetDeniedReason, "total_tokens_exceeded") +} + +func gatewaySumPoints( + t *testing.T, + rm metricdata.ResourceMetrics, + metricName string, +) []metricdata.DataPoint[int64] { + t.Helper() + for _, scopeMetric := range rm.ScopeMetrics { + for _, metric := range scopeMetric.Metrics { + if metric.Name != metricName { + continue + } + sum, ok := metric.Data.(metricdata.Sum[int64]) + require.True(t, ok) + return sum.DataPoints + } + } + t.Fatalf("metric %s not found", metricName) + return nil +} + +func requireGatewayAttr(t *testing.T, set attribute.Set, key string, value string) { + t.Helper() + for _, kv := range set.ToSlice() { + if string(kv.Key) == key { + require.Equal(t, value, kv.Value.AsString()) + return + } + } + t.Fatalf("attribute %s not found", key) +} diff --git a/internal/telemetry/trace.go b/internal/telemetry/trace.go index f04590543f..13e6b56475 100644 --- a/internal/telemetry/trace.go +++ b/internal/telemetry/trace.go @@ -58,6 +58,7 @@ const ( OperationToolCall = "tool.call" OperationToolApproval = "tool.approval" OperationAuditWrite = "audit.write" + OperationGatewayBudget = "gateway.budget" OperationMemorySearch = "memory.search" OperationMemoryWrite = "memory.write" OperationSummaryCreate = "summary.create" diff --git a/platform/gateway/service.go b/platform/gateway/service.go index fcb0f96c47..51336d238d 100644 --- a/platform/gateway/service.go +++ b/platform/gateway/service.go @@ -478,6 +478,15 @@ func (s *Service) checkBudget( return nil } budgetSpan.SetAttributes(attribute.String("decision", "deny")) + itelemetry.ReportGatewayBudgetDeniedMetrics( + auditCtx, + itelemetry.GatewayBudgetDeniedAttributes{ + TenantID: runtime.Tenant.TenantID, + AppName: runtime.App.AppID, + Channel: msg.Channel, + Reason: decision.Reason, + }, + ) s.writeBudgetDeniedAudit( auditCtx, auditSink, diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go index 0bcf3ad949..99007fad2b 100644 --- a/platform/gateway/service_test.go +++ b/platform/gateway/service_test.go @@ -1633,6 +1633,9 @@ func TestServiceWriteAuditRecordsAuditWriteFailureMetric(t *testing.T) { func TestServiceHandleInboundRejectsBudgetExceededBeforeIdempotency(t *testing.T) { ctx := context.Background() + reader, restore := useGatewayMetrics(t) + defer restore() + registry := NewInMemoryRegistry() r := &recordingRunner{response: "unused"} runtime := validRuntime("tenant-a", r) @@ -1695,6 +1698,16 @@ func TestServiceHandleInboundRejectsBudgetExceededBeforeIdempotency(t *testing.T assert.NotEmpty(t, estimateRequest.SessionID) assert.NotEmpty(t, estimateRequest.InternalUserID) + points := collectGatewayBudgetDeniedPoints(t, reader) + require.Len(t, points, 1) + require.Equal(t, int64(1), points[0].Value) + requireGatewayMetricAttr(t, points[0].Attributes, semconvtrace.KeyGenAIOperationName, itelemetry.OperationGatewayBudget) + requireGatewayMetricAttr(t, points[0].Attributes, semconvtrace.KeyGenAISystem, semconvtrace.SystemTRPCGoAgent) + requireGatewayMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoTenantID, "tenant-a") + requireGatewayMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoAppName, "app") + requireGatewayMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoChannel, "wecom") + requireGatewayMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoBudgetDeniedReason, "total_tokens_exceeded") + matches, queryErr := audit.Query(platform.AuditQueryFilter{ TenantID: "tenant-a", ToolName: "budget:tenant", @@ -2729,6 +2742,29 @@ func useAuditMetrics(t *testing.T) (*sdkmetric.ManualReader, func()) { } } +func useGatewayMetrics(t *testing.T) (*sdkmetric.ManualReader, func()) { + t.Helper() + reader := sdkmetric.NewManualReader() + provider := sdkmetric.NewMeterProvider(sdkmetric.WithReader(reader)) + + originalProvider := itelemetry.MeterProvider + originalMeter := itelemetry.GatewayMeter + originalCounter := itelemetry.GatewayMetricBudgetDeniedTotal + + itelemetry.MeterProvider = provider + itelemetry.GatewayMeter = provider.Meter(metrics.MeterNameGateway) + var err error + itelemetry.GatewayMetricBudgetDeniedTotal, err = + itelemetry.GatewayMeter.Int64Counter(metrics.MetricGatewayBudgetDeniedTotal) + require.NoError(t, err) + + return reader, func() { + itelemetry.MeterProvider = originalProvider + itelemetry.GatewayMeter = originalMeter + itelemetry.GatewayMetricBudgetDeniedTotal = originalCounter + } +} + func collectGatewayAuditWriteFailedPoints( t *testing.T, reader *sdkmetric.ManualReader, @@ -2750,6 +2786,38 @@ func collectGatewayAuditWriteFailedPoints( return nil } +func collectGatewayBudgetDeniedPoints( + t *testing.T, + reader *sdkmetric.ManualReader, +) []metricdata.DataPoint[int64] { + t.Helper() + var rm metricdata.ResourceMetrics + require.NoError(t, reader.Collect(context.Background(), &rm)) + for _, scopeMetric := range rm.ScopeMetrics { + for _, metric := range scopeMetric.Metrics { + if metric.Name != metrics.MetricGatewayBudgetDeniedTotal { + continue + } + sum, ok := metric.Data.(metricdata.Sum[int64]) + require.True(t, ok) + return sum.DataPoints + } + } + t.Fatalf("metric %s not found", metrics.MetricGatewayBudgetDeniedTotal) + return nil +} + +func requireGatewayMetricAttr(t *testing.T, set attribute.Set, key string, value string) { + t.Helper() + for _, kv := range set.ToSlice() { + if string(kv.Key) == key { + require.Equal(t, value, kv.Value.AsString()) + return + } + } + t.Fatalf("attribute %s not found", key) +} + func requireGatewayAuditMetricAttr(t *testing.T, set attribute.Set, key string, value string) { t.Helper() for _, kv := range set.ToSlice() { diff --git a/telemetry/metric/metric.go b/telemetry/metric/metric.go index e48e530556..bc2ea9dd3d 100644 --- a/telemetry/metric/metric.go +++ b/telemetry/metric/metric.go @@ -128,6 +128,9 @@ func InitMeterProvider(mp metric.MeterProvider) error { if err := initAuditMetrics(mp); err != nil { return err } + if err := initGatewayMetrics(mp); err != nil { + return err + } if err := initInvokeAgentMetrics(mp); err != nil { return err } @@ -246,6 +249,24 @@ func initAuditMetrics(mp metric.MeterProvider) error { return nil } +func initGatewayMetrics(mp metric.MeterProvider) error { + if mp == nil { + return fmt.Errorf("gateway meter provider is nil") + } + meterName := metrics.MeterNameGateway + itelemetry.GatewayMeter = mp.Meter(meterName) + var err error + itelemetry.GatewayMetricBudgetDeniedTotal, err = itelemetry.GatewayMeter.Int64Counter( + metrics.MetricGatewayBudgetDeniedTotal, + metric.WithDescription("Total number of gateway requests denied by budget checks"), + metric.WithUnit("1"), + ) + if err != nil { + return fmt.Errorf("failed to create %s metric %s: %w", meterName, metrics.MetricGatewayBudgetDeniedTotal, err) + } + return nil +} + func setInvokeAgentHistogramBuckets(metricName string, boundaries []float64) error { switch metricName { case metrics.MetricTRPCAgentGoClientTimeToFirstToken: diff --git a/telemetry/metric/metric_test.go b/telemetry/metric/metric_test.go index 3c9cd36e30..ce7e59ceae 100644 --- a/telemetry/metric/metric_test.go +++ b/telemetry/metric/metric_test.go @@ -299,6 +299,8 @@ func TestInitMeterProvider(t *testing.T) { originalToolPermissionDenied := itelemetry.ExecuteToolMetricToolPermissionDeniedTotal originalAuditMeter := itelemetry.AuditMeter originalAuditWriteFailed := itelemetry.AuditMetricWriteFailedTotal + originalGatewayMeter := itelemetry.GatewayMeter + originalGatewayBudgetDenied := itelemetry.GatewayMetricBudgetDeniedTotal defer func() { itelemetry.MeterProvider = originalMP itelemetry.ToolApprovalMeter = originalToolApprovalMeter @@ -306,6 +308,8 @@ func TestInitMeterProvider(t *testing.T) { itelemetry.ExecuteToolMetricToolPermissionDeniedTotal = originalToolPermissionDenied itelemetry.AuditMeter = originalAuditMeter itelemetry.AuditMetricWriteFailedTotal = originalAuditWriteFailed + itelemetry.GatewayMeter = originalGatewayMeter + itelemetry.GatewayMetricBudgetDeniedTotal = originalGatewayBudgetDenied }() // Create a test meter provider @@ -376,6 +380,12 @@ func TestInitMeterProvider(t *testing.T) { if itelemetry.AuditMetricWriteFailedTotal == nil { t.Error("AuditMetricWriteFailedTotal was not created") } + if itelemetry.GatewayMeter == nil { + t.Error("GatewayMeter was not created") + } + if itelemetry.GatewayMetricBudgetDeniedTotal == nil { + t.Error("GatewayMetricBudgetDeniedTotal was not created") + } if itelemetry.WorkflowMeter == nil { t.Error("WorkflowMeter was not created") } @@ -493,6 +503,31 @@ func TestInitAuditMetrics_ErrorHandling(t *testing.T) { } } +func TestInitGatewayMetrics_ErrorHandling(t *testing.T) { + originalGatewayMeter := itelemetry.GatewayMeter + originalGatewayBudgetDenied := itelemetry.GatewayMetricBudgetDeniedTotal + defer func() { + itelemetry.GatewayMeter = originalGatewayMeter + itelemetry.GatewayMetricBudgetDeniedTotal = originalGatewayBudgetDenied + }() + + if err := initGatewayMetrics(nil); err == nil || !strings.Contains(err.Error(), "gateway meter provider is nil") { + t.Fatalf("expected nil provider error, got %v", err) + } + + mp := &mockMeterProvider{meter: &mockMeter{ + shouldFail: true, + failOn: metrics.MetricGatewayBudgetDeniedTotal, + }} + err := initGatewayMetrics(mp) + if err == nil { + t.Fatalf("expected gateway counter creation error") + } + if !strings.Contains(err.Error(), "failed to create trpc_agent_go.internal.gateway metric gateway_budget_denied_total") { + t.Fatalf("unexpected error: %v", err) + } +} + func TestInitWorkflowMetrics_ErrorHandling(t *testing.T) { originalWorkflowMeter := itelemetry.WorkflowMeter originalWorkflowOpDur := itelemetry.WorkflowMetricGenAIClientOperationDuration diff --git a/telemetry/semconv/metrics/metrics.go b/telemetry/semconv/metrics/metrics.go index 6acebb4b39..5ce9965716 100644 --- a/telemetry/semconv/metrics/metrics.go +++ b/telemetry/semconv/metrics/metrics.go @@ -65,6 +65,8 @@ const ( MetricToolPermissionDeniedTotal = "tool_permission_denied_total" // MetricAuditWriteFailedTotal records failed audit sink writes. MetricAuditWriteFailedTotal = "audit_write_failed_total" + // MetricGatewayBudgetDeniedTotal records gateway requests denied by budget checks. + MetricGatewayBudgetDeniedTotal = "gateway_budget_denied_total" ////////////////////////// server //////////////////////// @@ -89,4 +91,6 @@ const ( MeterNameToolApproval = "trpc_agent_go.internal.tool_approval" // MeterNameAudit is the meter name for audit operations. MeterNameAudit = "trpc_agent_go.internal.audit" + // MeterNameGateway is the meter name for gateway operations. + MeterNameGateway = "trpc_agent_go.internal.gateway" ) diff --git a/telemetry/semconv/trace/trace.go b/telemetry/semconv/trace/trace.go index dead38d66a..be1c614e6f 100644 --- a/telemetry/semconv/trace/trace.go +++ b/telemetry/semconv/trace/trace.go @@ -64,6 +64,10 @@ const ( KeyTRPCAgentGoAuditDecision = "trpc.go.agent.audit.decision" // KeyTRPCAgentGoToolPermissionStatus is the structured tool permission result status. KeyTRPCAgentGoToolPermissionStatus = "trpc.go.agent.tool.permission.status" + // KeyTRPCAgentGoChannel is the inbound channel identifier. + KeyTRPCAgentGoChannel = "trpc.go.agent.channel" + // KeyTRPCAgentGoBudgetDeniedReason is the normalized budget denial reason. + KeyTRPCAgentGoBudgetDeniedReason = "trpc.go.agent.budget.denied.reason" // KeyGenAIAppName is the attribute key for GenAI application name. KeyGenAIAppName = "gen_ai.app.name" From fd732fac2f881c2443afc8b5f904a31bb2e8ac62 Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 15:11:12 +0800 Subject: [PATCH 92/95] feat(gateway): record rate limit metrics --- internal/telemetry/metric_gateway.go | 34 ++++++++++++++ internal/telemetry/metric_gateway_test.go | 55 +++++++++++++++++++++++ internal/telemetry/trace.go | 29 ++++++------ platform/gateway/service.go | 8 ++++ platform/gateway/service_test.go | 42 ++++++++++++++++- telemetry/metric/metric.go | 8 ++++ telemetry/metric/metric_test.go | 19 ++++++++ telemetry/semconv/metrics/metrics.go | 2 + 8 files changed, 181 insertions(+), 16 deletions(-) diff --git a/internal/telemetry/metric_gateway.go b/internal/telemetry/metric_gateway.go index dbfc1ae5f2..35c9184f09 100644 --- a/internal/telemetry/metric_gateway.go +++ b/internal/telemetry/metric_gateway.go @@ -24,6 +24,8 @@ var ( // GatewayMetricBudgetDeniedTotal records gateway requests denied by budget checks. GatewayMetricBudgetDeniedTotal metric.Int64Counter + // GatewayMetricRateLimitedTotal records inbound IM messages rejected by gateway rate limits. + GatewayMetricRateLimitedTotal metric.Int64Counter ) // GatewayBudgetDeniedAttributes is the attributes for gateway budget denial metrics. @@ -34,6 +36,13 @@ type GatewayBudgetDeniedAttributes struct { Reason string } +// GatewayRateLimitedAttributes is the attributes for gateway rate limit metrics. +type GatewayRateLimitedAttributes struct { + TenantID string + AppName string + Channel string +} + func (a GatewayBudgetDeniedAttributes) toAttributes() []attribute.KeyValue { attrs := []attribute.KeyValue{ attribute.String(semconvtrace.KeyGenAIOperationName, OperationGatewayBudget), @@ -54,6 +63,23 @@ func (a GatewayBudgetDeniedAttributes) toAttributes() []attribute.KeyValue { return attrs } +func (a GatewayRateLimitedAttributes) toAttributes() []attribute.KeyValue { + attrs := []attribute.KeyValue{ + attribute.String(semconvtrace.KeyGenAIOperationName, OperationGatewayRateLimit), + attribute.String(semconvtrace.KeyGenAISystem, semconvtrace.SystemTRPCGoAgent), + } + if a.TenantID != "" { + attrs = append(attrs, attribute.String(semconvtrace.KeyTRPCAgentGoTenantID, a.TenantID)) + } + if a.AppName != "" { + attrs = append(attrs, attribute.String(semconvtrace.KeyTRPCAgentGoAppName, a.AppName)) + } + if a.Channel != "" { + attrs = append(attrs, attribute.String(semconvtrace.KeyTRPCAgentGoChannel, a.Channel)) + } + return attrs +} + // ReportGatewayBudgetDeniedMetrics reports a gateway request denied by budget checks. func ReportGatewayBudgetDeniedMetrics(ctx context.Context, attrs GatewayBudgetDeniedAttributes) { if GatewayMetricBudgetDeniedTotal == nil { @@ -61,3 +87,11 @@ func ReportGatewayBudgetDeniedMetrics(ctx context.Context, attrs GatewayBudgetDe } GatewayMetricBudgetDeniedTotal.Add(ctx, 1, metric.WithAttributes(attrs.toAttributes()...)) } + +// ReportGatewayRateLimitedMetrics reports an inbound IM message rejected by gateway rate limits. +func ReportGatewayRateLimitedMetrics(ctx context.Context, attrs GatewayRateLimitedAttributes) { + if GatewayMetricRateLimitedTotal == nil { + return + } + GatewayMetricRateLimitedTotal.Add(ctx, 1, metric.WithAttributes(attrs.toAttributes()...)) +} diff --git a/internal/telemetry/metric_gateway_test.go b/internal/telemetry/metric_gateway_test.go index 18fc7c9249..1481fea6ec 100644 --- a/internal/telemetry/metric_gateway_test.go +++ b/internal/telemetry/metric_gateway_test.go @@ -79,6 +79,61 @@ func TestReportGatewayBudgetDeniedMetrics(t *testing.T) { requireGatewayAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoBudgetDeniedReason, "total_tokens_exceeded") } +func TestReportGatewayRateLimitedMetricsNoopWhenCounterNil(t *testing.T) { + originalCounter := GatewayMetricRateLimitedTotal + t.Cleanup(func() { + GatewayMetricRateLimitedTotal = originalCounter + }) + + GatewayMetricRateLimitedTotal = nil + require.NotPanics(t, func() { + ReportGatewayRateLimitedMetrics(context.Background(), GatewayRateLimitedAttributes{ + TenantID: "tenant-1", + AppName: "app-1", + Channel: "wecom", + }) + }) +} + +func TestReportGatewayRateLimitedMetrics(t *testing.T) { + reader := sdkmetric.NewManualReader() + provider := sdkmetric.NewMeterProvider(sdkmetric.WithReader(reader)) + + originalProvider := MeterProvider + originalMeter := GatewayMeter + originalCounter := GatewayMetricRateLimitedTotal + t.Cleanup(func() { + MeterProvider = originalProvider + GatewayMeter = originalMeter + GatewayMetricRateLimitedTotal = originalCounter + }) + + MeterProvider = provider + GatewayMeter = provider.Meter(metrics.MeterNameGateway) + var err error + GatewayMetricRateLimitedTotal, err = GatewayMeter.Int64Counter(metrics.MetricIMRateLimitedTotal) + require.NoError(t, err) + + ctx := context.Background() + ReportGatewayRateLimitedMetrics(ctx, GatewayRateLimitedAttributes{ + TenantID: "tenant-1", + AppName: "app-1", + Channel: "wecom", + }) + + var rm metricdata.ResourceMetrics + require.NoError(t, reader.Collect(ctx, &rm)) + + points := gatewaySumPoints(t, rm, metrics.MetricIMRateLimitedTotal) + require.Len(t, points, 1) + require.Equal(t, int64(1), points[0].Value) + requireGatewayAttr(t, points[0].Attributes, semconvtrace.KeyGenAIOperationName, OperationGatewayRateLimit) + requireGatewayAttr(t, points[0].Attributes, semconvtrace.KeyGenAISystem, semconvtrace.SystemTRPCGoAgent) + requireGatewayAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoTenantID, "tenant-1") + requireGatewayAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoAppName, "app-1") + requireGatewayAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoChannel, "wecom") +} + func gatewaySumPoints( t *testing.T, rm metricdata.ResourceMetrics, diff --git a/internal/telemetry/trace.go b/internal/telemetry/trace.go index 13e6b56475..1a03d921d1 100644 --- a/internal/telemetry/trace.go +++ b/internal/telemetry/trace.go @@ -54,20 +54,21 @@ const ( SpanNamePrefixExecuteTool = "execute_tool" - OperationExecuteTool = "execute_tool" - OperationToolCall = "tool.call" - OperationToolApproval = "tool.approval" - OperationAuditWrite = "audit.write" - OperationGatewayBudget = "gateway.budget" - OperationMemorySearch = "memory.search" - OperationMemoryWrite = "memory.write" - OperationSummaryCreate = "summary.create" - OperationChat = "chat" - OperationGenerateContent = "generate_content" - OperationInvokeAgent = "invoke_agent" - OperationCreateAgent = "create_agent" - OperationEmbeddings = "embeddings" - OperationWorkflow = "workflow" + OperationExecuteTool = "execute_tool" + OperationToolCall = "tool.call" + OperationToolApproval = "tool.approval" + OperationAuditWrite = "audit.write" + OperationGatewayBudget = "gateway.budget" + OperationGatewayRateLimit = "gateway.rate_limit" + OperationMemorySearch = "memory.search" + OperationMemoryWrite = "memory.write" + OperationSummaryCreate = "summary.create" + OperationChat = "chat" + OperationGenerateContent = "generate_content" + OperationInvokeAgent = "invoke_agent" + OperationCreateAgent = "create_agent" + OperationEmbeddings = "embeddings" + OperationWorkflow = "workflow" ) // Memory write operation values. diff --git a/platform/gateway/service.go b/platform/gateway/service.go index 51336d238d..d2f838b9b9 100644 --- a/platform/gateway/service.go +++ b/platform/gateway/service.go @@ -414,6 +414,14 @@ func (s *Service) checkRateLimit( if allowed { return nil } + itelemetry.ReportGatewayRateLimitedMetrics( + ctx, + itelemetry.GatewayRateLimitedAttributes{ + TenantID: runtime.Tenant.TenantID, + AppName: runtime.App.AppID, + Channel: msg.Channel, + }, + ) s.writeRejectAuditWithContextTo(ctx, auditSink, msg, start, ErrRateLimited, auditContext) recordSpanError(routeSpan, ErrRateLimited) return ErrRateLimited diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go index 99007fad2b..55dfad8a51 100644 --- a/platform/gateway/service_test.go +++ b/platform/gateway/service_test.go @@ -1292,6 +1292,9 @@ func TestServiceHandleInboundSkipsMIMEFilterWhenAllowlistUnset(t *testing.T) { func TestServiceHandleInboundRejectsRateLimitedBeforeBudgetAndIdempotency(t *testing.T) { ctx := context.Background() + reader, restore := useGatewayMetrics(t) + defer restore() + registry := NewInMemoryRegistry() r := &recordingRunner{response: "unused"} runtime := validRuntime("tenant-a", r) @@ -1334,6 +1337,15 @@ func TestServiceHandleInboundRejectsRateLimitedBeforeBudgetAndIdempotency(t *tes assert.Equal(t, ErrRateLimited.Error(), audit.Records()[1].DecisionReason) assert.NotEmpty(t, audit.Records()[1].SessionID) assert.NotEmpty(t, audit.Records()[1].InternalUserID) + + points := collectGatewayRateLimitedPoints(t, reader) + require.Len(t, points, 1) + require.Equal(t, int64(1), points[0].Value) + requireGatewayMetricAttr(t, points[0].Attributes, semconvtrace.KeyGenAIOperationName, itelemetry.OperationGatewayRateLimit) + requireGatewayMetricAttr(t, points[0].Attributes, semconvtrace.KeyGenAISystem, semconvtrace.SystemTRPCGoAgent) + requireGatewayMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoTenantID, "tenant-a") + requireGatewayMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoAppName, "app") + requireGatewayMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoChannel, "wecom") } func TestServiceHandleInboundRateLimitRefillsOverTime(t *testing.T) { @@ -2749,7 +2761,8 @@ func useGatewayMetrics(t *testing.T) (*sdkmetric.ManualReader, func()) { originalProvider := itelemetry.MeterProvider originalMeter := itelemetry.GatewayMeter - originalCounter := itelemetry.GatewayMetricBudgetDeniedTotal + originalBudgetCounter := itelemetry.GatewayMetricBudgetDeniedTotal + originalRateLimitCounter := itelemetry.GatewayMetricRateLimitedTotal itelemetry.MeterProvider = provider itelemetry.GatewayMeter = provider.Meter(metrics.MeterNameGateway) @@ -2757,11 +2770,15 @@ func useGatewayMetrics(t *testing.T) (*sdkmetric.ManualReader, func()) { itelemetry.GatewayMetricBudgetDeniedTotal, err = itelemetry.GatewayMeter.Int64Counter(metrics.MetricGatewayBudgetDeniedTotal) require.NoError(t, err) + itelemetry.GatewayMetricRateLimitedTotal, err = + itelemetry.GatewayMeter.Int64Counter(metrics.MetricIMRateLimitedTotal) + require.NoError(t, err) return reader, func() { itelemetry.MeterProvider = originalProvider itelemetry.GatewayMeter = originalMeter - itelemetry.GatewayMetricBudgetDeniedTotal = originalCounter + itelemetry.GatewayMetricBudgetDeniedTotal = originalBudgetCounter + itelemetry.GatewayMetricRateLimitedTotal = originalRateLimitCounter } } @@ -2807,6 +2824,27 @@ func collectGatewayBudgetDeniedPoints( return nil } +func collectGatewayRateLimitedPoints( + t *testing.T, + reader *sdkmetric.ManualReader, +) []metricdata.DataPoint[int64] { + t.Helper() + var rm metricdata.ResourceMetrics + require.NoError(t, reader.Collect(context.Background(), &rm)) + for _, scopeMetric := range rm.ScopeMetrics { + for _, metric := range scopeMetric.Metrics { + if metric.Name != metrics.MetricIMRateLimitedTotal { + continue + } + sum, ok := metric.Data.(metricdata.Sum[int64]) + require.True(t, ok) + return sum.DataPoints + } + } + t.Fatalf("metric %s not found", metrics.MetricIMRateLimitedTotal) + return nil +} + func requireGatewayMetricAttr(t *testing.T, set attribute.Set, key string, value string) { t.Helper() for _, kv := range set.ToSlice() { diff --git a/telemetry/metric/metric.go b/telemetry/metric/metric.go index bc2ea9dd3d..f57afa69c4 100644 --- a/telemetry/metric/metric.go +++ b/telemetry/metric/metric.go @@ -264,6 +264,14 @@ func initGatewayMetrics(mp metric.MeterProvider) error { if err != nil { return fmt.Errorf("failed to create %s metric %s: %w", meterName, metrics.MetricGatewayBudgetDeniedTotal, err) } + itelemetry.GatewayMetricRateLimitedTotal, err = itelemetry.GatewayMeter.Int64Counter( + metrics.MetricIMRateLimitedTotal, + metric.WithDescription("Total number of inbound IM messages rejected by gateway rate limits"), + metric.WithUnit("1"), + ) + if err != nil { + return fmt.Errorf("failed to create %s metric %s: %w", meterName, metrics.MetricIMRateLimitedTotal, err) + } return nil } diff --git a/telemetry/metric/metric_test.go b/telemetry/metric/metric_test.go index ce7e59ceae..3c49d3966f 100644 --- a/telemetry/metric/metric_test.go +++ b/telemetry/metric/metric_test.go @@ -301,6 +301,7 @@ func TestInitMeterProvider(t *testing.T) { originalAuditWriteFailed := itelemetry.AuditMetricWriteFailedTotal originalGatewayMeter := itelemetry.GatewayMeter originalGatewayBudgetDenied := itelemetry.GatewayMetricBudgetDeniedTotal + originalGatewayRateLimited := itelemetry.GatewayMetricRateLimitedTotal defer func() { itelemetry.MeterProvider = originalMP itelemetry.ToolApprovalMeter = originalToolApprovalMeter @@ -310,6 +311,7 @@ func TestInitMeterProvider(t *testing.T) { itelemetry.AuditMetricWriteFailedTotal = originalAuditWriteFailed itelemetry.GatewayMeter = originalGatewayMeter itelemetry.GatewayMetricBudgetDeniedTotal = originalGatewayBudgetDenied + itelemetry.GatewayMetricRateLimitedTotal = originalGatewayRateLimited }() // Create a test meter provider @@ -386,6 +388,9 @@ func TestInitMeterProvider(t *testing.T) { if itelemetry.GatewayMetricBudgetDeniedTotal == nil { t.Error("GatewayMetricBudgetDeniedTotal was not created") } + if itelemetry.GatewayMetricRateLimitedTotal == nil { + t.Error("GatewayMetricRateLimitedTotal was not created") + } if itelemetry.WorkflowMeter == nil { t.Error("WorkflowMeter was not created") } @@ -506,9 +511,11 @@ func TestInitAuditMetrics_ErrorHandling(t *testing.T) { func TestInitGatewayMetrics_ErrorHandling(t *testing.T) { originalGatewayMeter := itelemetry.GatewayMeter originalGatewayBudgetDenied := itelemetry.GatewayMetricBudgetDeniedTotal + originalGatewayRateLimited := itelemetry.GatewayMetricRateLimitedTotal defer func() { itelemetry.GatewayMeter = originalGatewayMeter itelemetry.GatewayMetricBudgetDeniedTotal = originalGatewayBudgetDenied + itelemetry.GatewayMetricRateLimitedTotal = originalGatewayRateLimited }() if err := initGatewayMetrics(nil); err == nil || !strings.Contains(err.Error(), "gateway meter provider is nil") { @@ -526,6 +533,18 @@ func TestInitGatewayMetrics_ErrorHandling(t *testing.T) { if !strings.Contains(err.Error(), "failed to create trpc_agent_go.internal.gateway metric gateway_budget_denied_total") { t.Fatalf("unexpected error: %v", err) } + + mp = &mockMeterProvider{meter: &mockMeter{ + shouldFail: true, + failOn: metrics.MetricIMRateLimitedTotal, + }} + err = initGatewayMetrics(mp) + if err == nil { + t.Fatalf("expected gateway rate limited counter creation error") + } + if !strings.Contains(err.Error(), "failed to create trpc_agent_go.internal.gateway metric im_rate_limited_total") { + t.Fatalf("unexpected error: %v", err) + } } func TestInitWorkflowMetrics_ErrorHandling(t *testing.T) { diff --git a/telemetry/semconv/metrics/metrics.go b/telemetry/semconv/metrics/metrics.go index 5ce9965716..54e5da3733 100644 --- a/telemetry/semconv/metrics/metrics.go +++ b/telemetry/semconv/metrics/metrics.go @@ -67,6 +67,8 @@ const ( MetricAuditWriteFailedTotal = "audit_write_failed_total" // MetricGatewayBudgetDeniedTotal records gateway requests denied by budget checks. MetricGatewayBudgetDeniedTotal = "gateway_budget_denied_total" + // MetricIMRateLimitedTotal records inbound IM messages rejected by gateway rate limits. + MetricIMRateLimitedTotal = "im_rate_limited_total" ////////////////////////// server //////////////////////// From 359a9f488830b201ed5c84642d454bc7dd305a79 Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 15:28:51 +0800 Subject: [PATCH 93/95] feat(gateway): record idempotency hit metrics --- internal/telemetry/metric_gateway.go | 38 ++++++ internal/telemetry/metric_gateway_test.go | 58 +++++++++ internal/telemetry/trace.go | 31 ++--- platform/gateway/service.go | 18 +++ platform/gateway/service_test.go | 143 ++++++++++++++++++++++ telemetry/metric/metric.go | 8 ++ telemetry/metric/metric_test.go | 19 +++ telemetry/semconv/metrics/metrics.go | 2 + telemetry/semconv/trace/trace.go | 2 + 9 files changed, 304 insertions(+), 15 deletions(-) diff --git a/internal/telemetry/metric_gateway.go b/internal/telemetry/metric_gateway.go index 35c9184f09..ab5dbe2287 100644 --- a/internal/telemetry/metric_gateway.go +++ b/internal/telemetry/metric_gateway.go @@ -26,6 +26,8 @@ var ( GatewayMetricBudgetDeniedTotal metric.Int64Counter // GatewayMetricRateLimitedTotal records inbound IM messages rejected by gateway rate limits. GatewayMetricRateLimitedTotal metric.Int64Counter + // GatewayMetricIdempotencyHitTotal records inbound IM messages served by gateway idempotency. + GatewayMetricIdempotencyHitTotal metric.Int64Counter ) // GatewayBudgetDeniedAttributes is the attributes for gateway budget denial metrics. @@ -43,6 +45,14 @@ type GatewayRateLimitedAttributes struct { Channel string } +// GatewayIdempotencyHitAttributes is the attributes for gateway idempotency hit metrics. +type GatewayIdempotencyHitAttributes struct { + TenantID string + AppName string + Channel string + Status string +} + func (a GatewayBudgetDeniedAttributes) toAttributes() []attribute.KeyValue { attrs := []attribute.KeyValue{ attribute.String(semconvtrace.KeyGenAIOperationName, OperationGatewayBudget), @@ -80,6 +90,26 @@ func (a GatewayRateLimitedAttributes) toAttributes() []attribute.KeyValue { return attrs } +func (a GatewayIdempotencyHitAttributes) toAttributes() []attribute.KeyValue { + attrs := []attribute.KeyValue{ + attribute.String(semconvtrace.KeyGenAIOperationName, OperationGatewayIdempotency), + attribute.String(semconvtrace.KeyGenAISystem, semconvtrace.SystemTRPCGoAgent), + } + if a.TenantID != "" { + attrs = append(attrs, attribute.String(semconvtrace.KeyTRPCAgentGoTenantID, a.TenantID)) + } + if a.AppName != "" { + attrs = append(attrs, attribute.String(semconvtrace.KeyTRPCAgentGoAppName, a.AppName)) + } + if a.Channel != "" { + attrs = append(attrs, attribute.String(semconvtrace.KeyTRPCAgentGoChannel, a.Channel)) + } + if a.Status != "" { + attrs = append(attrs, attribute.String(semconvtrace.KeyTRPCAgentGoIdempotencyStatus, a.Status)) + } + return attrs +} + // ReportGatewayBudgetDeniedMetrics reports a gateway request denied by budget checks. func ReportGatewayBudgetDeniedMetrics(ctx context.Context, attrs GatewayBudgetDeniedAttributes) { if GatewayMetricBudgetDeniedTotal == nil { @@ -95,3 +125,11 @@ func ReportGatewayRateLimitedMetrics(ctx context.Context, attrs GatewayRateLimit } GatewayMetricRateLimitedTotal.Add(ctx, 1, metric.WithAttributes(attrs.toAttributes()...)) } + +// ReportGatewayIdempotencyHitMetrics reports an inbound IM message served by gateway idempotency. +func ReportGatewayIdempotencyHitMetrics(ctx context.Context, attrs GatewayIdempotencyHitAttributes) { + if GatewayMetricIdempotencyHitTotal == nil { + return + } + GatewayMetricIdempotencyHitTotal.Add(ctx, 1, metric.WithAttributes(attrs.toAttributes()...)) +} diff --git a/internal/telemetry/metric_gateway_test.go b/internal/telemetry/metric_gateway_test.go index 1481fea6ec..9550f36b7c 100644 --- a/internal/telemetry/metric_gateway_test.go +++ b/internal/telemetry/metric_gateway_test.go @@ -134,6 +134,64 @@ func TestReportGatewayRateLimitedMetrics(t *testing.T) { requireGatewayAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoChannel, "wecom") } +func TestReportGatewayIdempotencyHitMetricsNoopWhenCounterNil(t *testing.T) { + originalCounter := GatewayMetricIdempotencyHitTotal + t.Cleanup(func() { + GatewayMetricIdempotencyHitTotal = originalCounter + }) + + GatewayMetricIdempotencyHitTotal = nil + require.NotPanics(t, func() { + ReportGatewayIdempotencyHitMetrics(context.Background(), GatewayIdempotencyHitAttributes{ + TenantID: "tenant-1", + AppName: "app-1", + Channel: "wecom", + Status: "completed", + }) + }) +} + +func TestReportGatewayIdempotencyHitMetrics(t *testing.T) { + reader := sdkmetric.NewManualReader() + provider := sdkmetric.NewMeterProvider(sdkmetric.WithReader(reader)) + + originalProvider := MeterProvider + originalMeter := GatewayMeter + originalCounter := GatewayMetricIdempotencyHitTotal + t.Cleanup(func() { + MeterProvider = originalProvider + GatewayMeter = originalMeter + GatewayMetricIdempotencyHitTotal = originalCounter + }) + + MeterProvider = provider + GatewayMeter = provider.Meter(metrics.MeterNameGateway) + var err error + GatewayMetricIdempotencyHitTotal, err = GatewayMeter.Int64Counter(metrics.MetricGatewayIdempotencyHitTotal) + require.NoError(t, err) + + ctx := context.Background() + ReportGatewayIdempotencyHitMetrics(ctx, GatewayIdempotencyHitAttributes{ + TenantID: "tenant-1", + AppName: "app-1", + Channel: "wecom", + Status: "completed", + }) + + var rm metricdata.ResourceMetrics + require.NoError(t, reader.Collect(ctx, &rm)) + + points := gatewaySumPoints(t, rm, metrics.MetricGatewayIdempotencyHitTotal) + require.Len(t, points, 1) + require.Equal(t, int64(1), points[0].Value) + requireGatewayAttr(t, points[0].Attributes, semconvtrace.KeyGenAIOperationName, OperationGatewayIdempotency) + requireGatewayAttr(t, points[0].Attributes, semconvtrace.KeyGenAISystem, semconvtrace.SystemTRPCGoAgent) + requireGatewayAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoTenantID, "tenant-1") + requireGatewayAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoAppName, "app-1") + requireGatewayAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoChannel, "wecom") + requireGatewayAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoIdempotencyStatus, "completed") +} + func gatewaySumPoints( t *testing.T, rm metricdata.ResourceMetrics, diff --git a/internal/telemetry/trace.go b/internal/telemetry/trace.go index 1a03d921d1..9d2cd18a91 100644 --- a/internal/telemetry/trace.go +++ b/internal/telemetry/trace.go @@ -54,21 +54,22 @@ const ( SpanNamePrefixExecuteTool = "execute_tool" - OperationExecuteTool = "execute_tool" - OperationToolCall = "tool.call" - OperationToolApproval = "tool.approval" - OperationAuditWrite = "audit.write" - OperationGatewayBudget = "gateway.budget" - OperationGatewayRateLimit = "gateway.rate_limit" - OperationMemorySearch = "memory.search" - OperationMemoryWrite = "memory.write" - OperationSummaryCreate = "summary.create" - OperationChat = "chat" - OperationGenerateContent = "generate_content" - OperationInvokeAgent = "invoke_agent" - OperationCreateAgent = "create_agent" - OperationEmbeddings = "embeddings" - OperationWorkflow = "workflow" + OperationExecuteTool = "execute_tool" + OperationToolCall = "tool.call" + OperationToolApproval = "tool.approval" + OperationAuditWrite = "audit.write" + OperationGatewayBudget = "gateway.budget" + OperationGatewayIdempotency = "gateway.idempotency" + OperationGatewayRateLimit = "gateway.rate_limit" + OperationMemorySearch = "memory.search" + OperationMemoryWrite = "memory.write" + OperationSummaryCreate = "summary.create" + OperationChat = "chat" + OperationGenerateContent = "generate_content" + OperationInvokeAgent = "invoke_agent" + OperationCreateAgent = "create_agent" + OperationEmbeddings = "embeddings" + OperationWorkflow = "workflow" ) // Memory write operation values. diff --git a/platform/gateway/service.go b/platform/gateway/service.go index d2f838b9b9..59a08310c7 100644 --- a/platform/gateway/service.go +++ b/platform/gateway/service.go @@ -616,6 +616,7 @@ func (s *Service) startInboundRun( return inboundRunRecord{}, false, Result{}, err } if ok { + reportGatewayIdempotencyHit(resultCtx, msg, existing) result, err := s.duplicateResult(resultCtx, existing) return inboundRunRecord{}, true, result, err } @@ -685,12 +686,29 @@ func (s *Service) acquireSessionLeaseAndStart( if !started { s.releaseSessionLease(resultCtx, lease) s.releaseUserConcurrency(resultCtx, userLease) + reportGatewayIdempotencyHit(resultCtx, msg, record) result, err := s.duplicateResult(resultCtx, record) return inboundRunRecord{}, true, result, err } return inboundRunRecord{Record: record, SessionLease: lease, UserLease: userLease}, false, Result{}, nil } +func reportGatewayIdempotencyHit( + ctx context.Context, + msg platform.InboundMessage, + record platform.IdempotencyRecord, +) { + itelemetry.ReportGatewayIdempotencyHitMetrics( + ctx, + itelemetry.GatewayIdempotencyHitAttributes{ + TenantID: record.TenantID, + AppName: msg.AppID, + Channel: record.Channel, + Status: string(record.Status), + }, + ) +} + func (s *Service) acquireUserConcurrency( routeCtx context.Context, msg platform.InboundMessage, diff --git a/platform/gateway/service_test.go b/platform/gateway/service_test.go index 55dfad8a51..ff4d2c9516 100644 --- a/platform/gateway/service_test.go +++ b/platform/gateway/service_test.go @@ -582,6 +582,8 @@ func TestServiceHandleInboundOutboxFailureDoesNotCompleteIdempotency(t *testing. func TestServiceHandleInboundDuplicateReplyFailedReusesStoredOutbound(t *testing.T) { ctx := context.Background() + reader, restore := useGatewayMetrics(t) + defer restore() registry := NewInMemoryRegistry() r := &recordingRunner{response: "queued"} registerRuntime(t, registry, "tenant-a", r) @@ -616,10 +618,22 @@ func TestServiceHandleInboundDuplicateReplyFailedReusesStoredOutbound(t *testing assert.Equal(t, resultRef, dup.ResultRef) assert.Equal(t, "queued", dup.Outbound.Content) assert.Len(t, r.calls, 1) + + points := collectGatewayIdempotencyHitPoints(t, reader) + require.Len(t, points, 1) + assert.Equal(t, int64(1), points[0].Value) + requireGatewayMetricAttr(t, points[0].Attributes, semconvtrace.KeyGenAIOperationName, itelemetry.OperationGatewayIdempotency) + requireGatewayMetricAttr(t, points[0].Attributes, semconvtrace.KeyGenAISystem, semconvtrace.SystemTRPCGoAgent) + requireGatewayMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoTenantID, "tenant-a") + requireGatewayMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoAppName, "app") + requireGatewayMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoChannel, "wecom") + requireGatewayMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoIdempotencyStatus, "reply_failed") } func TestServiceHandleInboundDuplicateProcessingDoesNotRun(t *testing.T) { ctx := context.Background() + reader, restore := useGatewayMetrics(t) + defer restore() registry := NewInMemoryRegistry() r := &blockingRunner{started: make(chan struct{})} registerRuntime(t, registry, "tenant-a", r) @@ -642,10 +656,63 @@ func TestServiceHandleInboundDuplicateProcessingDoesNotRun(t *testing.T) { assert.True(t, dup.Duplicate) assert.True(t, dup.Processing) assert.Len(t, r.calls, 1) + + points := collectGatewayIdempotencyHitPoints(t, reader) + require.Len(t, points, 1) + assert.Equal(t, int64(1), points[0].Value) + requireGatewayMetricAttr(t, points[0].Attributes, semconvtrace.KeyGenAIOperationName, itelemetry.OperationGatewayIdempotency) + requireGatewayMetricAttr(t, points[0].Attributes, semconvtrace.KeyGenAISystem, semconvtrace.SystemTRPCGoAgent) + requireGatewayMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoTenantID, "tenant-a") + requireGatewayMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoAppName, "app") + requireGatewayMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoChannel, "wecom") + requireGatewayMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoIdempotencyStatus, "processing") r.finish("done") require.NoError(t, <-errCh) } +func TestServiceHandleInboundIdempotencyStartConflictRecordsMetric(t *testing.T) { + ctx := context.Background() + reader, restore := useGatewayMetrics(t) + defer restore() + registry := NewInMemoryRegistry() + r := &recordingRunner{response: "first"} + registerRuntime(t, registry, "tenant-a", r) + key := platform.IdempotencyKey("tenant-a", "wecom", "acct", "msg-1") + store := &startConflictIdempotencyStore{ + record: platform.IdempotencyRecord{ + TenantID: "tenant-a", + Channel: "wecom", + AccountID: "acct", + PlatformMessageID: "msg-1", + IdempotencyKey: key, + RequestID: "existing-request", + SessionID: "existing-session", + Status: platform.IdempotencyStatusProcessing, + }, + } + svc := NewService(registry, store, NewInMemoryOutboundStore()) + + result, err := svc.HandleInbound(ctx, inbound("tenant-a", "msg-1", "user-1", "hello")) + require.NoError(t, err) + + assert.True(t, result.Duplicate) + assert.True(t, result.Processing) + assert.Equal(t, platform.IdempotencyStatusProcessing, result.Status) + assert.Len(t, r.calls, 0) + assert.Equal(t, 1, store.getCalls) + assert.Equal(t, 1, store.startCalls) + + points := collectGatewayIdempotencyHitPoints(t, reader) + require.Len(t, points, 1) + assert.Equal(t, int64(1), points[0].Value) + requireGatewayMetricAttr(t, points[0].Attributes, semconvtrace.KeyGenAIOperationName, itelemetry.OperationGatewayIdempotency) + requireGatewayMetricAttr(t, points[0].Attributes, semconvtrace.KeyGenAISystem, semconvtrace.SystemTRPCGoAgent) + requireGatewayMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoTenantID, "tenant-a") + requireGatewayMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoAppName, "app") + requireGatewayMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoChannel, "wecom") + requireGatewayMetricAttr(t, points[0].Attributes, semconvtrace.KeyTRPCAgentGoIdempotencyStatus, "processing") +} + func TestServiceHandleInboundSerializesSameSession(t *testing.T) { ctx := context.Background() registry := NewInMemoryRegistry() @@ -2763,6 +2830,7 @@ func useGatewayMetrics(t *testing.T) (*sdkmetric.ManualReader, func()) { originalMeter := itelemetry.GatewayMeter originalBudgetCounter := itelemetry.GatewayMetricBudgetDeniedTotal originalRateLimitCounter := itelemetry.GatewayMetricRateLimitedTotal + originalIdempotencyHitCounter := itelemetry.GatewayMetricIdempotencyHitTotal itelemetry.MeterProvider = provider itelemetry.GatewayMeter = provider.Meter(metrics.MeterNameGateway) @@ -2773,12 +2841,16 @@ func useGatewayMetrics(t *testing.T) (*sdkmetric.ManualReader, func()) { itelemetry.GatewayMetricRateLimitedTotal, err = itelemetry.GatewayMeter.Int64Counter(metrics.MetricIMRateLimitedTotal) require.NoError(t, err) + itelemetry.GatewayMetricIdempotencyHitTotal, err = + itelemetry.GatewayMeter.Int64Counter(metrics.MetricGatewayIdempotencyHitTotal) + require.NoError(t, err) return reader, func() { itelemetry.MeterProvider = originalProvider itelemetry.GatewayMeter = originalMeter itelemetry.GatewayMetricBudgetDeniedTotal = originalBudgetCounter itelemetry.GatewayMetricRateLimitedTotal = originalRateLimitCounter + itelemetry.GatewayMetricIdempotencyHitTotal = originalIdempotencyHitCounter } } @@ -2845,6 +2917,27 @@ func collectGatewayRateLimitedPoints( return nil } +func collectGatewayIdempotencyHitPoints( + t *testing.T, + reader *sdkmetric.ManualReader, +) []metricdata.DataPoint[int64] { + t.Helper() + var rm metricdata.ResourceMetrics + require.NoError(t, reader.Collect(context.Background(), &rm)) + for _, scopeMetric := range rm.ScopeMetrics { + for _, metric := range scopeMetric.Metrics { + if metric.Name != metrics.MetricGatewayIdempotencyHitTotal { + continue + } + sum, ok := metric.Data.(metricdata.Sum[int64]) + require.True(t, ok) + return sum.DataPoints + } + } + t.Fatalf("metric %s not found", metrics.MetricGatewayIdempotencyHitTotal) + return nil +} + func requireGatewayMetricAttr(t *testing.T, set attribute.Set, key string, value string) { t.Helper() for _, kv := range set.ToSlice() { @@ -2866,3 +2959,53 @@ func requireGatewayAuditMetricAttr(t *testing.T, set attribute.Set, key string, } t.Fatalf("attribute %s not found", key) } + +type startConflictIdempotencyStore struct { + record platform.IdempotencyRecord + getCalls int + startCalls int +} + +func (s *startConflictIdempotencyStore) Start( + ctx context.Context, + record platform.IdempotencyRecord, +) (platform.IdempotencyRecord, bool, error) { + if err := ctx.Err(); err != nil { + return platform.IdempotencyRecord{}, false, err + } + s.startCalls++ + return s.record, false, nil +} + +func (s *startConflictIdempotencyStore) Complete( + ctx context.Context, + key string, + resultRef string, +) (platform.IdempotencyRecord, error) { + if err := ctx.Err(); err != nil { + return platform.IdempotencyRecord{}, err + } + return platform.IdempotencyRecord{}, platform.ErrIdempotencyRecordNotFound +} + +func (s *startConflictIdempotencyStore) MarkReplyFailed( + ctx context.Context, + key string, + resultRef string, +) (platform.IdempotencyRecord, error) { + if err := ctx.Err(); err != nil { + return platform.IdempotencyRecord{}, err + } + return platform.IdempotencyRecord{}, platform.ErrIdempotencyRecordNotFound +} + +func (s *startConflictIdempotencyStore) Get( + ctx context.Context, + key string, +) (platform.IdempotencyRecord, bool, error) { + if err := ctx.Err(); err != nil { + return platform.IdempotencyRecord{}, false, err + } + s.getCalls++ + return platform.IdempotencyRecord{}, false, nil +} diff --git a/telemetry/metric/metric.go b/telemetry/metric/metric.go index f57afa69c4..7993136860 100644 --- a/telemetry/metric/metric.go +++ b/telemetry/metric/metric.go @@ -272,6 +272,14 @@ func initGatewayMetrics(mp metric.MeterProvider) error { if err != nil { return fmt.Errorf("failed to create %s metric %s: %w", meterName, metrics.MetricIMRateLimitedTotal, err) } + itelemetry.GatewayMetricIdempotencyHitTotal, err = itelemetry.GatewayMeter.Int64Counter( + metrics.MetricGatewayIdempotencyHitTotal, + metric.WithDescription("Total number of inbound IM messages served by gateway idempotency"), + metric.WithUnit("1"), + ) + if err != nil { + return fmt.Errorf("failed to create %s metric %s: %w", meterName, metrics.MetricGatewayIdempotencyHitTotal, err) + } return nil } diff --git a/telemetry/metric/metric_test.go b/telemetry/metric/metric_test.go index 3c49d3966f..cbfcfccaca 100644 --- a/telemetry/metric/metric_test.go +++ b/telemetry/metric/metric_test.go @@ -302,6 +302,7 @@ func TestInitMeterProvider(t *testing.T) { originalGatewayMeter := itelemetry.GatewayMeter originalGatewayBudgetDenied := itelemetry.GatewayMetricBudgetDeniedTotal originalGatewayRateLimited := itelemetry.GatewayMetricRateLimitedTotal + originalGatewayIdempotencyHit := itelemetry.GatewayMetricIdempotencyHitTotal defer func() { itelemetry.MeterProvider = originalMP itelemetry.ToolApprovalMeter = originalToolApprovalMeter @@ -312,6 +313,7 @@ func TestInitMeterProvider(t *testing.T) { itelemetry.GatewayMeter = originalGatewayMeter itelemetry.GatewayMetricBudgetDeniedTotal = originalGatewayBudgetDenied itelemetry.GatewayMetricRateLimitedTotal = originalGatewayRateLimited + itelemetry.GatewayMetricIdempotencyHitTotal = originalGatewayIdempotencyHit }() // Create a test meter provider @@ -391,6 +393,9 @@ func TestInitMeterProvider(t *testing.T) { if itelemetry.GatewayMetricRateLimitedTotal == nil { t.Error("GatewayMetricRateLimitedTotal was not created") } + if itelemetry.GatewayMetricIdempotencyHitTotal == nil { + t.Error("GatewayMetricIdempotencyHitTotal was not created") + } if itelemetry.WorkflowMeter == nil { t.Error("WorkflowMeter was not created") } @@ -512,10 +517,12 @@ func TestInitGatewayMetrics_ErrorHandling(t *testing.T) { originalGatewayMeter := itelemetry.GatewayMeter originalGatewayBudgetDenied := itelemetry.GatewayMetricBudgetDeniedTotal originalGatewayRateLimited := itelemetry.GatewayMetricRateLimitedTotal + originalGatewayIdempotencyHit := itelemetry.GatewayMetricIdempotencyHitTotal defer func() { itelemetry.GatewayMeter = originalGatewayMeter itelemetry.GatewayMetricBudgetDeniedTotal = originalGatewayBudgetDenied itelemetry.GatewayMetricRateLimitedTotal = originalGatewayRateLimited + itelemetry.GatewayMetricIdempotencyHitTotal = originalGatewayIdempotencyHit }() if err := initGatewayMetrics(nil); err == nil || !strings.Contains(err.Error(), "gateway meter provider is nil") { @@ -545,6 +552,18 @@ func TestInitGatewayMetrics_ErrorHandling(t *testing.T) { if !strings.Contains(err.Error(), "failed to create trpc_agent_go.internal.gateway metric im_rate_limited_total") { t.Fatalf("unexpected error: %v", err) } + + mp = &mockMeterProvider{meter: &mockMeter{ + shouldFail: true, + failOn: metrics.MetricGatewayIdempotencyHitTotal, + }} + err = initGatewayMetrics(mp) + if err == nil { + t.Fatalf("expected gateway idempotency hit counter creation error") + } + if !strings.Contains(err.Error(), "failed to create trpc_agent_go.internal.gateway metric gateway_idempotency_hit_total") { + t.Fatalf("unexpected error: %v", err) + } } func TestInitWorkflowMetrics_ErrorHandling(t *testing.T) { diff --git a/telemetry/semconv/metrics/metrics.go b/telemetry/semconv/metrics/metrics.go index 54e5da3733..6f17fff4b2 100644 --- a/telemetry/semconv/metrics/metrics.go +++ b/telemetry/semconv/metrics/metrics.go @@ -69,6 +69,8 @@ const ( MetricGatewayBudgetDeniedTotal = "gateway_budget_denied_total" // MetricIMRateLimitedTotal records inbound IM messages rejected by gateway rate limits. MetricIMRateLimitedTotal = "im_rate_limited_total" + // MetricGatewayIdempotencyHitTotal records inbound IM messages served by gateway idempotency. + MetricGatewayIdempotencyHitTotal = "gateway_idempotency_hit_total" ////////////////////////// server //////////////////////// diff --git a/telemetry/semconv/trace/trace.go b/telemetry/semconv/trace/trace.go index be1c614e6f..94c28d11bb 100644 --- a/telemetry/semconv/trace/trace.go +++ b/telemetry/semconv/trace/trace.go @@ -68,6 +68,8 @@ const ( KeyTRPCAgentGoChannel = "trpc.go.agent.channel" // KeyTRPCAgentGoBudgetDeniedReason is the normalized budget denial reason. KeyTRPCAgentGoBudgetDeniedReason = "trpc.go.agent.budget.denied.reason" + // KeyTRPCAgentGoIdempotencyStatus is the stored gateway idempotency record status. + KeyTRPCAgentGoIdempotencyStatus = "trpc.go.agent.idempotency.status" // KeyGenAIAppName is the attribute key for GenAI application name. KeyGenAIAppName = "gen_ai.app.name" From a8819d3d023d3bfbff5223f2fb625e5657a74652 Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 15:39:22 +0800 Subject: [PATCH 94/95] fix(approval): redact reviewer decision text --- plugin/guardrail/approval/approval.go | 16 ++++- plugin/guardrail/approval/approval_test.go | 80 ++++++++++++++++++++++ 2 files changed, 94 insertions(+), 2 deletions(-) diff --git a/plugin/guardrail/approval/approval.go b/plugin/guardrail/approval/approval.go index cda093ffcb..db373d4652 100644 --- a/plugin/guardrail/approval/approval.go +++ b/plugin/guardrail/approval/approval.go @@ -153,8 +153,8 @@ func (p *Plugin) beforeTool() tool.BeforeToolCallbackStructured { CustomResult: fmt.Sprintf("approval review failed for tool %q: %v", args.ToolName, err), }, nil } - riskLevel := strings.TrimSpace(decision.RiskLevel) - reason := strings.TrimSpace(decision.Reason) + riskLevel := sanitizeReviewerText(decision.RiskLevel) + reason := sanitizeReviewerText(decision.Reason) if decision.Approved { if err := p.writeApprovalAudit( ctx, @@ -220,6 +220,18 @@ func (p *Plugin) beforeTool() tool.BeforeToolCallbackStructured { } } +func sanitizeReviewerText(value string) string { + value = strings.TrimSpace(value) + if value == "" { + return "" + } + redactor, err := platform.NewRedactor() + if err != nil { + return value + } + return redactor.Redact(value) +} + func (p *Plugin) resolvePolicy(args *tool.BeforeToolArgs) ToolPolicy { if args == nil { return p.defaultToolPolicy diff --git a/plugin/guardrail/approval/approval_test.go b/plugin/guardrail/approval/approval_test.go index 90313923ac..f4e9e3a0c6 100644 --- a/plugin/guardrail/approval/approval_test.go +++ b/plugin/guardrail/approval/approval_test.go @@ -310,6 +310,44 @@ func TestBeforeTool_ReviewerApprovedLogsInfo(t *testing.T) { ) } +func TestBeforeTool_ReviewerApprovedLogRedactsSensitiveReason(t *testing.T) { + original := approvallog.InfofContext + var infoLog string + approvallog.InfofContext = func(ctx context.Context, format string, args ...any) { + infoLog = fmt.Sprintf(format, args...) + } + defer func() { + approvallog.InfofContext = original + }() + p, err := New(WithReviewer(&stubReviewer{ + reviewFn: func(ctx context.Context, req *approvalreview.Request) (*approvalreview.Decision, error) { + return &approvalreview.Decision{ + Approved: true, + RiskScore: 42, + RiskLevel: "medium token=sk-risk-level-secret", + Reason: "Allowed with Authorization: Bearer raw-token and password=plain.", + }, nil + }, + })) + require.NoError(t, err) + callbacks := registeredToolCallbacks(t, p) + result, runErr := callbacks.RunBeforeTool(context.Background(), &tool.BeforeToolArgs{ + ToolName: "shell", + ToolCallID: "call-1", + Arguments: []byte(`{"command":"pwd"}`), + }) + require.NoError(t, runErr) + require.NotNil(t, result) + + require.Contains(t, infoLog, "Automatic approval review approved") + require.NotContains(t, infoLog, "sk-risk-level-secret") + require.NotContains(t, infoLog, "raw-token") + require.NotContains(t, infoLog, "password=plain") + require.Contains(t, infoLog, "token=****") + require.Contains(t, infoLog, "Authorization: ****") + require.Contains(t, infoLog, "password=****") +} + func TestBeforeTool_MetadataRiskRequiresApproval(t *testing.T) { metadata := tool.ToolMetadata{ ReadOnly: false, @@ -689,6 +727,48 @@ func TestBeforeTool_ReviewerDeniedLogsWarning(t *testing.T) { ) } +func TestBeforeTool_ReviewerDeniedRedactsSensitiveReason(t *testing.T) { + original := approvallog.WarnContext + var warning string + approvallog.WarnContext = func(ctx context.Context, args ...any) { + warning = fmt.Sprint(args...) + } + defer func() { + approvallog.WarnContext = original + }() + p, err := New(WithReviewer(&stubReviewer{ + reviewFn: func(ctx context.Context, req *approvalreview.Request) (*approvalreview.Decision, error) { + return &approvalreview.Decision{ + Approved: false, + RiskScore: 92, + RiskLevel: "high api_key=sk-risk-level-secret", + Reason: "Blocked because Authorization: Bearer raw-token and password=plain were present.", + }, nil + }, + })) + require.NoError(t, err) + callbacks := registeredToolCallbacks(t, p) + result, runErr := callbacks.RunBeforeTool(context.Background(), &tool.BeforeToolArgs{ + ToolName: "shell", + ToolCallID: "call-1", + Arguments: []byte(`{"command":"rm -rf /tmp/demo"}`), + }) + require.NoError(t, runErr) + require.NotNil(t, result) + permission, ok := result.CustomResult.(tool.PermissionResult) + require.True(t, ok) + + for _, text := range []string{permission.Reason, warning} { + require.Contains(t, text, "Automatic approval review denied") + require.NotContains(t, text, "sk-risk-level-secret") + require.NotContains(t, text, "raw-token") + require.NotContains(t, text, "password=plain") + require.Contains(t, text, "api_key=****") + require.Contains(t, text, "Authorization: ****") + require.Contains(t, text, "password=****") + } +} + func TestBeforeTool_UnsupportedPolicyReturnsFailureMessage(t *testing.T) { p := &Plugin{ name: "approval", From fb858506b9f8f41bd289cae4560ed92948adeba6 Mon Sep 17 00:00:00 2001 From: XnLemon Date: Sat, 11 Jul 2026 15:43:25 +0800 Subject: [PATCH 95/95] fix(platform): redact spaced secret assignments --- platform/redaction.go | 4 ++-- platform/types_test.go | 31 +++++++++++++++++++++++++++++++ 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/platform/redaction.go b/platform/redaction.go index 46643b4232..d402a26c6b 100644 --- a/platform/redaction.go +++ b/platform/redaction.go @@ -19,8 +19,8 @@ var defaultRedactionPatterns = []*regexp.Regexp{ regexp.MustCompile(`(?i)(Bearer\s+)[A-Za-z0-9._~+/\-]+=*`), regexp.MustCompile(`(?im)(authorization\s*:\s*(?:token|digest)\s+)[^\r\n]+`), regexp.MustCompile(`(?im)(authorization\s*=\s*(?:token|digest)\s+)[^\r\n]+`), - regexp.MustCompile(`(?i)(api[_-]?key|token|secret|password|passwd|cookie)=([^&\s]+)`), - regexp.MustCompile(`(?i)(api[_-]?key|token|secret|password|passwd|cookie):\s*([^,\s]+)`), + regexp.MustCompile(`(?i)(api[_-]?key|token|secret|password|passwd|cookie)\s*=\s*([^&\s]+)`), + regexp.MustCompile(`(?i)(api[_-]?key|token|secret|password|passwd|cookie)\s*:\s*([^,\s]+)`), regexp.MustCompile(`(?i)("(?:api[_-]?key|token|secret|password|passwd|authorization|cookie)"\s*:\s*")([^"]+)(")`), regexp.MustCompile(`(?i)(sk-[A-Za-z0-9._~+/\-]{8,})`), regexp.MustCompile(`(?i)[a-z][a-z0-9+.-]*://[^\s/?#]*@[^\s/?#]+`), diff --git a/platform/types_test.go b/platform/types_test.go index 0f5b686d68..67711d9cda 100644 --- a/platform/types_test.go +++ b/platform/types_test.go @@ -819,6 +819,37 @@ func TestRedactorMasksSecrets(t *testing.T) { } } +func TestRedactorMasksSpacedSecretAssignments(t *testing.T) { + redactor, err := NewRedactor() + if err != nil { + t.Fatalf("NewRedactor: %v", err) + } + input := `api_key = sk-1234567890abcdef password : plain-token token = raw-token secret : sk-secret-value cookie = session-secret` + got := redactor.Redact(input) + for _, leaked := range []string{ + "sk-1234567890abcdef", + "plain-token", + "raw-token", + "sk-secret-value", + "session-secret", + } { + if strings.Contains(got, leaked) { + t.Fatalf("redacted output leaked %q: %q", leaked, got) + } + } + for _, want := range []string{ + "api_key =****", + "password : ****", + "token =****", + "secret : ****", + "cookie =****", + } { + if !strings.Contains(got, want) { + t.Fatalf("expected %q in redacted output, got %q", want, got) + } + } +} + func TestRedactorMasksNonBearerAuthorizationCredentials(t *testing.T) { redactor, err := NewRedactor() if err != nil {