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..feccb2a1cc 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( @@ -4358,54 +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) - - 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, - 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) @@ -4432,14 +4443,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 @@ -4455,30 +4468,178 @@ 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 } +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, + 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, @@ -4561,8 +4722,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 +4781,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..0561134a37 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, @@ -1981,14 +1991,23 @@ func (f *Flow) getFilteredTools( hasUserToolTracking, userToolNames, ) - allTools, userToolNames, hasUserToolTracking, externalToolNames := + allTools, userToolNames, _, externalToolNames := toolsurface.AppendRunOptionTools( allTools, userToolNames, 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) @@ -2004,8 +2023,16 @@ func (f *Flow) getFilteredTools( userToolNames, externalToolNames, ) - hasUserToolTracking = userToolNames != nil } + allTools, userToolNames, externalToolNames = + toolsurface.ApplyMandatoryToolFilter( + ctx, + allTools, + userToolNames, + externalToolNames, + invocation.RunOptions, + ) + filteredHasUserToolTracking := userToolNames != nil // If no filter is specified, return all tools for this invocation. if invocation.RunOptions.ToolFilter == nil { @@ -2014,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 @@ -2027,7 +2054,7 @@ func (f *Flow) getFilteredTools( ctx, allTools, userToolNames, - hasUserToolTracking, + filteredHasUserToolTracking, invocation.RunOptions, ) @@ -2035,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), ) @@ -2181,6 +2208,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 +2247,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 +2261,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/content.go b/internal/flow/processor/content.go index 0db00738ea..86a31372ce 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,8 +3109,23 @@ func (p *ContentRequestProcessor) getAdaptivePreloadMemoryMessage( Deduplicate: true, HybridSearch: true, } - memories, err := reader.SearchMemories( - ctx, + searchCtx, span, startedSpan := itrace.StartSpan(ctx, inv, itelemetry.NewMemorySearchSpanName()) + 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), 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/flow/processor/functioncall.go b/internal/flow/processor/functioncall.go index 95b9837856..a8556aed7f 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() @@ -1768,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 { @@ -1883,10 +1886,72 @@ 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, - tool.PermissionResultStatusApprovalRequired: + tool.PermissionResultStatusApprovalRequired, + tool.PermissionResultStatusApprovalDenied: return true default: return false @@ -2224,6 +2289,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 @@ -2239,6 +2305,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 { @@ -2267,6 +2334,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 @@ -2277,6 +2345,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 { @@ -2427,11 +2496,31 @@ 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 + } + 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, toolCall, toolDeclaration, + toolMetadata, ) if err != nil { return ctx, nil, toolCall.Function.Arguments, false, false, err @@ -2446,6 +2535,7 @@ func (p *FunctionCallResponseProcessor) executeToolWithCallbacks( ctx, toolCall, toolDeclaration, + toolMetadata, ) if err != nil { return ctx, nil, toolCall.Function.Arguments, false, false, err @@ -2462,6 +2552,7 @@ func (p *FunctionCallResponseProcessor) executeToolWithCallbacks( toolCall, tl, toolDeclaration, + toolMetadata, ) if err != nil { return ctx, nil, toolCall.Function.Arguments, false, false, err @@ -2540,12 +2631,50 @@ 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, + metadata tool.ToolMetadata, +) (*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: metadata, + } + 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, toolCall model.ToolCall, tl tool.Tool, decl *tool.Declaration, + metadata tool.ToolMetadata, ) (*tool.PermissionResult, error) { semanticTool := itool.ResolveSemantic(tl) req := &tool.PermissionRequest{ @@ -2554,7 +2683,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) @@ -2563,10 +2692,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 81301d4132..6c5e873641 100644 --- a/internal/flow/processor/functioncall_test.go +++ b/internal/flow/processor/functioncall_test.go @@ -24,17 +24,23 @@ 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" agenttool "trpc.group/trpc-go/trpc-agent-go/tool/agent" @@ -92,6 +98,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 @@ -107,6 +129,7 @@ type permissionMockTool struct { *mockCallableTool metadata tool.ToolMetadata decision tool.PermissionDecision + permissionCalled bool stateDelta map[string][]byte stateDeltaCalled bool skipSummarize bool @@ -132,6 +155,7 @@ func (m *permissionMockTool) CheckPermission( _ context.Context, _ *tool.PermissionRequest, ) (tool.PermissionDecision, error) { + m.permissionCalled = true return m.decision, nil } @@ -226,6 +250,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 +436,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) @@ -8787,6 +9013,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, ) { @@ -8796,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 @@ -8850,6 +9149,136 @@ 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( + 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"}` + ) + reader := setupToolPermissionDeniedMetric(t) + + 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) + requireToolPermissionDeniedMetric(t, reader, tool.PermissionResultStatusApprovalDenied) +} + +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( @@ -8953,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 @@ -8995,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( @@ -9047,6 +9554,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/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/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_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_gateway.go b/internal/telemetry/metric_gateway.go new file mode 100644 index 0000000000..ab5dbe2287 --- /dev/null +++ b/internal/telemetry/metric_gateway.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 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 + // 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. +type GatewayBudgetDeniedAttributes struct { + TenantID string + AppName string + Channel string + Reason string +} + +// GatewayRateLimitedAttributes is the attributes for gateway rate limit metrics. +type GatewayRateLimitedAttributes struct { + TenantID string + AppName string + 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), + 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 +} + +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 +} + +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 { + return + } + 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()...)) +} + +// 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 new file mode 100644 index 0000000000..9550f36b7c --- /dev/null +++ b/internal/telemetry/metric_gateway_test.go @@ -0,0 +1,224 @@ +// +// 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 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 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, + 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/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/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 e65abef0d1..9d2cd18a91 100644 --- a/internal/telemetry/trace.go +++ b/internal/telemetry/trace.go @@ -12,8 +12,14 @@ package telemetry import ( + "crypto/hmac" + "crypto/sha256" + "encoding/hex" "encoding/json" + "errors" "fmt" + "os" + "strings" "time" "go.opentelemetry.io/otel/attribute" @@ -31,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 @@ -44,13 +54,30 @@ const ( SpanNamePrefixExecuteTool = "execute_tool" - OperationExecuteTool = "execute_tool" - 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. +const ( + MemoryWriteOperationAdd = "add" + MemoryWriteOperationUpdate = "update" + MemoryWriteOperationDelete = "delete" + MemoryWriteOperationClear = "clear" ) // NewChatSpanName creates a new chat span name. @@ -63,6 +90,104 @@ func NewExecuteToolSpanName(toolName string) string { return OperationExecuteTool + " " + toolName } +// NewToolCallSpanName creates the stable platform tool-call span contract name. +func NewToolCallSpanName() string { + return OperationToolCall +} + +// NewMemorySearchSpanName creates the stable platform memory-search span contract name. +func NewMemorySearchSpanName() string { + return OperationMemorySearch +} + +// NewMemoryWriteSpanName creates the stable platform memory-write span contract name. +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() { + return + } + 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) + } +} + +// 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) + } +} + +// 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 + } + errorType := ToErrorType(err, fallback) + span.SetAttributes(attribute.String(semconvtrace.KeyErrorType, errorType)) + span.SetStatus(codes.Error, errorType) + 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 @@ -201,7 +326,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 +339,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 +379,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 +406,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..61f0ed499c 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,113 @@ func TestNewExecuteToolSpanName(t *testing.T) { } } +func TestNewToolCallSpanName(t *testing.T) { + require.Equal(t, OperationToolCall, NewToolCallSpanName()) +} + +func TestNewMemorySearchSpanName(t *testing.T) { + require.Equal(t, OperationMemorySearch, NewMemorySearchSpanName()) +} + +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) + 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 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 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 @@ -598,6 +741,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 +752,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 +881,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/internal/toolsurface/toolsurface.go b/internal/toolsurface/toolsurface.go index cb5e815001..518c759400 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 @@ -48,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 @@ -124,6 +129,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 +202,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 @@ -206,13 +276,28 @@ func EffectiveWithExternal( ctx = context.Background() } allTools, userToolNames, hasUserToolTracking := ResolveBase(ctx, invocation) - allTools, userToolNames, hasUserToolTracking, externalNames := + allTools, userToolNames, _, externalNames := AppendRunOptionTools( allTools, userToolNames, hasUserToolTracking, invocation.RunOptions, ) + allTools, userToolNames, externalNames, _ = + ApplyInvocationToolActivation( + ctx, + invocation, + allTools, + userToolNames, + externalNames, + ) + allTools, userToolNames, externalNames = ApplyMandatoryToolFilter( + ctx, + allTools, + userToolNames, + externalNames, + invocation.RunOptions, + ) if invocation.RunOptions.ToolFilter == nil { return allTools, userToolNames, externalNames } @@ -220,7 +305,7 @@ func EffectiveWithExternal( ctx, allTools, userToolNames, - hasUserToolTracking, + userToolNames != nil, invocation.RunOptions, ), userToolNames, externalNames } @@ -260,6 +345,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 +355,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/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/memory/tool/tool.go b/memory/tool/tool.go index e6f4d760b6..b5edada7f1 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" @@ -81,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) } @@ -141,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) } @@ -186,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) } @@ -225,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) } @@ -272,11 +326,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..162022c517 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,88 @@ 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 { + 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() == spanName { + return span + } + } + t.Fatalf("span %q not recorded; ended spans=%v", spanName, 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 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() @@ -191,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() @@ -266,6 +377,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() @@ -343,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() @@ -482,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() @@ -1111,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() @@ -1151,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/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 +} diff --git a/platform/audit.go b/platform/audit.go new file mode 100644 index 0000000000..9bca743481 --- /dev/null +++ b/platform/audit.go @@ -0,0 +1,41 @@ +// +// 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" +) + +// AuditSink stores audit records. +type AuditSink interface { + // WriteAudit writes one audit record. + WriteAudit(ctx context.Context, record AuditRecord) error +} + +// 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(options ...InMemorySinkOption) *InMemoryAuditSink { + return &InMemoryAuditSink{ + records: newInMemoryRecords[AuditRecord](options...), + } +} + +// WriteAudit writes one audit record. +func (s *InMemoryAuditSink) WriteAudit(ctx context.Context, record AuditRecord) error { + return s.records.append(ctx, record, AuditRecord.Validate) +} + +// Records returns a snapshot of written audit records. +func (s *InMemoryAuditSink) Records() []AuditRecord { + return s.records.snapshot() +} 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/audit_query.go b/platform/audit_query.go new file mode 100644 index 0000000000..1bcfd863c3 --- /dev/null +++ b/platform/audit_query.go @@ -0,0 +1,169 @@ +// +// 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 + 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. +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 + } + 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{"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{"redacted_detail_ref", f.RedactedDetailRef}, + safeTextField{"redaction_version", f.RedactionVersion}, + ); 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.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.RedactedDetailRef = strings.TrimSpace(f.RedactedDetailRef) + f.RedactionVersion = strings.TrimSpace(f.RedactionVersion) + return f, nil +} + +func (f AuditQueryFilter) matchesScope(record AuditRecord) bool { + return strings.TrimSpace(record.TenantID) == f.TenantID +} + +func (f AuditQueryFilter) matches(record AuditRecord) bool { + 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.RedactedDetailRef, record.RedactedDetailRef}, + {f.RedactionVersion, record.RedactionVersion}, + {f.TraceID, record.TraceID}, + } { + if !matchOptional(field.want, field.got) { + return false + } + } + return true +} + +func (f AuditQueryFilter) matchesCreatedAt(createdAt time.Time) bool { + if !f.CreatedFrom.IsZero() && createdAt.Before(f.CreatedFrom) { + return false + } + if !f.CreatedTo.IsZero() && 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..b848d5c7c0 --- /dev/null +++ b/platform/audit_query_test.go @@ -0,0 +1,287 @@ +// +// 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 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 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()) + 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) { + 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", + }, + "redacted_detail_ref": { + RedactedDetailRef: "outcome:deny 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) + } + }) + } +} + +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 +} diff --git a/platform/audit_record_test.go b/platform/audit_record_test.go new file mode 100644 index 0000000000..661c971a46 --- /dev/null +++ b/platform/audit_record_test.go @@ -0,0 +1,150 @@ +// +// 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 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" + 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/backend_migration_status.go b/platform/backend_migration_status.go new file mode 100644 index 0000000000..7cc6353bbd --- /dev/null +++ b/platform/backend_migration_status.go @@ -0,0 +1,463 @@ +// +// 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" + // BackendMigrationResourceSummary covers session summary storage migrations. + BackendMigrationResourceSummary BackendMigrationResource = "summary" + // 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 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 + } + 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") + } + return nil +} + +func (r BackendMigrationStatusReport) validateBackendMigrationState() error { + 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") + } + return nil +} + +func (r BackendMigrationStatusReport) validateBackendMigrationSafeText() error { + 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, + BackendMigrationResourceSummary, + 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..984862c0f5 --- /dev/null +++ b/platform/backend_migration_status_test.go @@ -0,0 +1,383 @@ +// +// 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, + BackendMigrationResourceSummary, + 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(), + } +} diff --git a/platform/budget.go b/platform/budget.go new file mode 100644 index 0000000000..2c58b0e778 --- /dev/null +++ b/platform/budget.go @@ -0,0 +1,340 @@ +// +// 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 +} + +// 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 { + 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) +} + +// 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 { + 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) +} + +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 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 { + 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") + } + return nil +} + +func (s BudgetUsageSnapshot) quota() TenantQuota { + return TenantQuota{ + MaxPromptTokens: s.MaxPromptTokens, + MaxCompletionTokens: s.MaxCompletionTokens, + MaxTotalTokens: s.MaxTotalTokens, + MaxCost: s.MaxCost, + } +} + +func (s BudgetUsageSnapshot) estimate() (UsageEstimate, error) { + estimate := UsageEstimate{ + PromptTokens: s.PromptTokensUsed, + CompletionTokens: s.CompletionTokensUsed, + TotalTokens: s.TotalTokensUsed, + Cost: s.CostUsed, + } + effectiveTotalTokens, err := estimate.effectiveTotalTokens() + if err != nil { + return UsageEstimate{}, err + } + if s.TotalTokensUsed != effectiveTotalTokens { + 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 + } + 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") + } + 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_audit.go b/platform/budget_audit.go new file mode 100644 index 0000000000..efee53372f --- /dev/null +++ b/platform/budget_audit.go @@ -0,0 +1,393 @@ +// +// 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 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 + } + 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") + } + return nil +} + +func (s BudgetDecisionSummary) validateEstimateAndQuota() (BudgetDecision, error) { + estimate := UsageEstimate{ + PromptTokens: s.EstimatedPrompt, + CompletionTokens: s.EstimatedCompletion, + TotalTokens: s.EstimatedTotalTokens, + Cost: s.EstimatedCost, + } + if err := validateUsageEstimate(estimate); err != nil { + return BudgetDecision{}, err + } + canonicalTotal, err := estimate.effectiveTotalTokens() + if err != nil { + return BudgetDecision{}, err + } + if canonicalTotal != s.EstimatedTotalTokens { + return BudgetDecision{}, fmt.Errorf("estimated_total_tokens must match effective total tokens") + } + quota := s.quota() + if err := quota.Validate(); err != nil { + return BudgetDecision{}, err + } + expected, err := quota.Check(estimate) + if err != nil { + return BudgetDecision{}, err + } + return expected, nil +} + +func (s BudgetDecisionSummary) validateOutcome(expected BudgetDecision) error { + 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) + } + 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) + } +} diff --git a/platform/budget_test.go b/platform/budget_test.go new file mode 100644 index 0000000000..1bbb1491dc --- /dev/null +++ b/platform/budget_test.go @@ -0,0 +1,356 @@ +// +// 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 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 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 + 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") + } +} 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) + } +} 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..9bad4a9dd3 --- /dev/null +++ b/platform/channeladapter/errors.go @@ -0,0 +1,32 @@ +// +// 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") + // 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 new file mode 100644 index 0000000000..1885bb9535 --- /dev/null +++ b/platform/channeladapter/outbox.go @@ -0,0 +1,473 @@ +// +// 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) +} + +// 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 + 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 +} + +// 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, + 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 +} + +// 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, + 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 "" + } + redactor, redactorErr := platform.NewRedactor() + if redactorErr != nil { + return err.Error() + } + return redactor.Redact(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..db5d3b7788 --- /dev/null +++ b/platform/channeladapter/outbox_test.go @@ -0,0 +1,672 @@ +// +// 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" + "strings" + "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 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, "postgres://****@example/db") || + !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() + 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() + 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") + } +} diff --git a/platform/config_cache_invalidation.go b/platform/config_cache_invalidation.go new file mode 100644 index 0000000000..06c4e81aeb --- /dev/null +++ b/platform/config_cache_invalidation.go @@ -0,0 +1,197 @@ +// +// 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.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") + } + 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..c63475573c --- /dev/null +++ b/platform/config_cache_invalidation_test.go @@ -0,0 +1,236 @@ +// +// 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.Version = " v2 " + whitespace.Checksum = " sha256:previous " + whitespaceNext := next + 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) + } + + 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 || + !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(), + } +} diff --git a/platform/config_diff.go b/platform/config_diff.go new file mode 100644 index 0000000000..2b8b37744c --- /dev/null +++ b/platform/config_diff.go @@ -0,0 +1,211 @@ +// +// 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" + "errors" + "fmt" + "io" + "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 + } + 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 +} + +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..3a7de30a67 --- /dev/null +++ b/platform/config_diff_test.go @@ -0,0 +1,167 @@ +// +// 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 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"]}` + 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) +} 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 +} diff --git a/platform/config_operation_summary.go b/platform/config_operation_summary.go new file mode 100644 index 0000000000..89f9e8a4a6 --- /dev/null +++ b/platform/config_operation_summary.go @@ -0,0 +1,414 @@ +// +// 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 must already carry AppConfigVersionStatusRollback. + PreviousActive AppConfigVersion + // NextActive must carry AppConfigVersionStatusActive. + 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 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 + } + 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") + } + return nil +} + +func (s AppConfigOperationSummary) validateConfigOperationState() error { + 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") + } + return nil +} + +func (s AppConfigOperationSummary) validateConfigOperationLinks() error { + 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") + } + return nil +} + +func (s AppConfigOperationSummary) validateConfigOperationSafeText() error { + 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 { + 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 { + 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 err := validateConfigOperationGrayCandidate(s.GrayStatus); err != nil { + return err + } + 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 validateConfigOperationGrayStatusText(status ConfigGrayStatusSummary) error { + 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 { + 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() { + 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 + } + } + if err := validateAuditRedactedFields( + safeTextField{"operation_id", i.OperationID}, + safeTextField{"trace_id", i.TraceID}, + ); 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(), + } +} 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/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/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..8ea65eaac9 --- /dev/null +++ b/platform/gateway/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 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") + // 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. + 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. + 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") + // 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. + 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/lease.go b/platform/gateway/lease.go new file mode 100644 index 0000000000..4348b7590f --- /dev/null +++ b/platform/gateway/lease.go @@ -0,0 +1,93 @@ +// +// 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 { + 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]int64 + counter map[SessionLeaseKey]int64 +} + +// NewInMemorySessionLeaseStore creates an empty process-local session lease store. +func NewInMemorySessionLeaseStore() *InMemorySessionLeaseStore { + return &InMemorySessionLeaseStore{ + leases: make(map[SessionLeaseKey]int64), + counter: make(map[SessionLeaseKey]int64), + } +} + +// 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 + } + token := s.counter[key] + 1 + s.counter[key] = token + s.leases[key] = token + return &inMemorySessionLease{ + store: s, + key: key, + fencingToken: token, + }, true, nil +} + +type inMemorySessionLease struct { + 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 { + 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/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/registry.go b/platform/gateway/registry.go new file mode 100644 index 0000000000..cc80a21ece --- /dev/null +++ b/platform/gateway/registry.go @@ -0,0 +1,162 @@ +// +// 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" + "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. +type Runtime struct { + 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. + ToolPermissionPolicy tool.PermissionPolicy +} + +// 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 err := validateRuntimeModelProfile(r); err != nil { + return err + } + 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 { + 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 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 && + 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, + channel, + accountID, + 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..632324751f --- /dev/null +++ b/platform/gateway/registry_test.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 gateway + +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 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 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, + 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 new file mode 100644 index 0000000000..59a08310c7 --- /dev/null +++ b/platform/gateway/service.go @@ -0,0 +1,1719 @@ +// +// 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" + "fmt" + "reflect" + "strings" + "sync" + "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" + 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/platform/toolpolicy" + "trpc.group/trpc-go/trpc-agent-go/plugin/guardrail/approval" + telemetrytrace "trpc.group/trpc-go/trpc-agent-go/telemetry/trace" +) + +// Service handles normalized inbound platform messages. +type Service struct { + registry Registry + idempotencyStore platform.IdempotencyStore + outboundStore OutboundStore + leaseStore SessionLeaseStore + userGate UserConcurrencyGate + auditSink platform.AuditSink + messageEventSink platform.MessageEventSink + usageSink platform.UsageSink + budgetEstimator BudgetEstimator + rateLimiter RateLimiter + 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) +} + +// 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 + 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) { + s.auditSink = sink + } +} + +// WithMessageEventSink sets the message event sink used by the service. +func WithMessageEventSink(sink platform.MessageEventSink) Option { + return func(s *Service) { + s.messageEventSink = sink + } +} + +// 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) { + if now != nil { + s.now = now + } + } +} + +// WithSessionLeaseStore sets the lease store used to serialize same-session runs. +func WithSessionLeaseStore(store SessionLeaseStore) Option { + return func(s *Service) { + 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 + } + } +} + +// WithBudgetEstimator enables pre-run tenant budget checks. +func WithBudgetEstimator(estimator BudgetEstimator) Option { + return func(s *Service) { + s.budgetEstimator = estimator + } +} + +// 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, + idempotencyStore platform.IdempotencyStore, + outboundStore OutboundStore, + opts ...Option, +) *Service { + svc := &Service{ + registry: registry, + idempotencyStore: idempotencyStore, + outboundStore: outboundStore, + leaseStore: NewInMemorySessionLeaseStore(), + userGate: NewInMemoryUserConcurrencyGate(), + rateLimiter: NewInMemoryRateLimiter(), + 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() + 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 + } + requestID := requestIDFor(msg) + sessionID, err := platform.SessionIDForInbound(msg) + if err != nil { + recordSpanError(callbackSpan, err) + return Result{}, err + } + 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 + } + auditSink := s.auditSinkForRuntime(runtime) + text, err := s.validateInboundContent(ctx, routeSpan, runtime, msg, start, auditSink, auditContext) + if err != nil { + return Result{}, err + } + if err := s.checkRateLimit(ctx, routeSpan, runtime, auditSink, msg, start, auditContext); err != nil { + return Result{}, err + } + if err := s.checkBudget( + routeCtx, + ctx, + routeSpan, + runtime, + auditSink, + msg, + text, + sessionID, + requestID, + internalUserID, + start, + auditContext, + ); err != nil { + return Result{}, err + } + key := platform.IdempotencyKey( + msg.TenantID, + msg.Channel, + msg.ChannelAccountID, + msg.PlatformMessageID, + ) + record, handled, result, err := s.startInboundRun( + routeCtx, + ctx, + msg, + sessionID, + requestID, + internalUserID, + key, + runtime.Binding.ChannelLimits.MaxConcurrentPerUser, + ) + if err != nil { + return Result{}, err + } + if handled { + return result, nil + } + defer s.releaseSessionLease(ctx, record.SessionLease) + defer s.releaseUserConcurrency(ctx, record.UserLease) + + return s.runAndReply( + routeCtx, + ctx, + runtime, + auditSink, + msg, + inboundRunInput{ + Text: text, + SessionID: sessionID, + InternalUserID: internalUserID, + RequestID: requestID, + Key: key, + FencingToken: record.SessionLease.FencingToken(), + Start: start, + }, + ) +} + +type inboundRunRecord struct { + Record platform.IdempotencyRecord + SessionLease SessionLease + UserLease UserConcurrencyLease +} + +type inboundRunInput struct { + Text string + SessionID string + InternalUserID string + RequestID string + Key string + FencingToken int64 + Start time.Time +} + +type runnerOutput struct { + Content string + Usage *model.Usage +} + +func (s *Service) lookupRuntime( + routeCtx context.Context, + auditCtx context.Context, + routeSpan oteltrace.Span, + msg platform.InboundMessage, + start time.Time, + auditContext rejectAuditContext, +) (Runtime, error) { + runtime, ok, err := s.registry.Lookup(routeCtx, msg) + if err != nil { + recordSpanError(routeSpan, err) + return Runtime{}, err + } + if !ok { + err := ErrRuntimeNotFound + s.writeRejectAuditWithContext(auditCtx, msg, start, err, auditContext) + recordSpanError(routeSpan, err) + return Runtime{}, err + } + if err := validateRuntimeForMessage(runtime, msg); err != nil { + s.writeRejectAuditWithContext(auditCtx, msg, start, err, auditContext) + recordSpanError(routeSpan, err) + return Runtime{}, err + } + if err := authorizeBinding(runtime.Binding, msg); err != nil { + s.writeRejectAuditWithContextTo( + auditCtx, + s.auditSinkForRuntime(runtime), + msg, + start, + err, + auditContext, + ) + recordSpanError(routeSpan, err) + return Runtime{}, err + } + return runtime, nil +} + +func (s *Service) checkRateLimit( + ctx context.Context, + routeSpan oteltrace.Span, + runtime Runtime, + auditSink platform.AuditSink, + msg platform.InboundMessage, + start time.Time, + auditContext rejectAuditContext, +) 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.writeRejectAuditWithContextTo(ctx, auditSink, msg, start, err, auditContext) + recordSpanError(routeSpan, err) + return err + } + 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 +} + +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, + auditContext rejectAuditContext, +) 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.writeRejectAuditWithContextTo(auditCtx, auditSink, msg, start, err, auditContext) + 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.writeRejectAuditWithContextTo(auditCtx, auditSink, msg, start, err, auditContext) + recordSpanError(routeSpan, err) + recordSpanError(budgetSpan, err) + return err + } + decision, err := quota.Check(estimate) + if err != nil { + s.writeRejectAuditWithContextTo(auditCtx, auditSink, msg, start, err, auditContext) + 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")) + itelemetry.ReportGatewayBudgetDeniedMetrics( + auditCtx, + itelemetry.GatewayBudgetDeniedAttributes{ + TenantID: runtime.Tenant.TenantID, + AppName: runtime.App.AppID, + Channel: msg.Channel, + Reason: decision.Reason, + }, + ) + s.writeBudgetDeniedAudit( + auditCtx, + auditSink, + runtime, + requestID, + decision, + estimate, + quota, + msg, + auditContext, + 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, + msg platform.InboundMessage, + auditContext rejectAuditContext, + 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 + } + 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 + record.AgentName = runtime.App.AgentName + record.ModelName = usageModelName(runtime) + if err := record.Validate(); err != nil { + s.writeRejectAuditWithContextTo(ctx, auditSink, msg, start, err, auditContext) + return + } + s.writeAuditTo(ctx, auditSink, record) +} + +func validateRuntimeForMessage(runtime Runtime, msg platform.InboundMessage) error { + if err := runtime.Validate(); err != nil { + return err + } + if !runtime.matchesInbound(msg) { + return ErrRuntimeMismatch + } + return nil +} + +func (s *Service) validateInboundContent( + ctx context.Context, + routeSpan oteltrace.Span, + runtime Runtime, + msg platform.InboundMessage, + start time.Time, + auditSink platform.AuditSink, + auditContext rejectAuditContext, +) (string, error) { + if err := validateFileLimits(msg, runtime.Binding.ChannelLimits); err != nil { + s.writeRejectAuditWithContextTo(ctx, auditSink, msg, start, err, auditContext) + recordSpanError(routeSpan, err) + return "", err + } + text, err := inboundText(msg) + if err != nil { + s.writeRejectAuditWithContextTo(ctx, auditSink, msg, start, err, auditContext) + recordSpanError(routeSpan, err) + return "", err + } + if err := validateTextLimit(text, runtime.Binding.ChannelLimits); err != nil { + s.writeRejectAuditWithContextTo(ctx, auditSink, msg, start, err, auditContext) + 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, + userConcurrencyLimit int, +) (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) + return inboundRunRecord{}, false, Result{}, err + } + if ok { + reportGatewayIdempotencyHit(resultCtx, msg, existing) + 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, + idempotencyCtx, + idempotencySpan, + msg, + sessionID, + requestID, + internalUserID, + key, + userLease, + ) +} + +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, + userLease UserConcurrencyLease, +) (inboundRunRecord, bool, Result, error) { + lease, handled, result, err := s.acquireSessionLease( + routeCtx, + msg, + sessionID, + requestID, + internalUserID, + ) + if err != nil || handled { + s.releaseUserConcurrency(resultCtx, userLease) + return inboundRunRecord{}, handled, result, err + } + record, started, err := s.idempotencyStore.Start(idempotencyCtx, platform.IdempotencyRecord{ + TenantID: msg.TenantID, + Channel: msg.Channel, + AccountID: msg.ChannelAccountID, + PlatformMessageID: msg.PlatformMessageID, + IdempotencyKey: key, + RequestID: requestID, + SessionID: sessionID, + }) + 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) + 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, + 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( + 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 + } + 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) { + 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) +} + +func (s *Service) runAndReply( + routeCtx context.Context, + auditCtx context.Context, + runtime Runtime, + auditSink platform.AuditSink, + msg platform.InboundMessage, + input inboundRunInput, +) (Result, error) { + output, err := s.runGatewayRunner( + routeCtx, + auditCtx, + runtime, + auditSink, + msg, + input, + ) + if err != nil { + return Result{}, err + } + result, err := s.writeReply( + routeCtx, + auditCtx, + runtime, + auditSink, + msg, + input, + output.Content, + ) + if err != nil { + return Result{}, err + } + s.writeUsageRecord(auditCtx, runtime, msg, input, output.Usage) + return result, nil +} + +func (s *Service) runGatewayRunner( + routeCtx context.Context, + auditCtx context.Context, + runtime Runtime, + auditSink platform.AuditSink, + msg platform.InboundMessage, + input inboundRunInput, +) (runnerOutput, error) { + 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, + }) + 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)) + } + 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), + runOptions..., + ) + if err != nil { + s.writeAuditTo(auditCtx, auditSink, auditFromMessage(msg, input.SessionID, input.InternalUserID, "runner_error", err.Error(), input.Start, err)) + recordSpanError(runnerSpan, err) + return runnerOutput{}, err + } + 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 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 + } + 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: 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, + runtime Runtime, + auditSink platform.AuditSink, + msg platform.InboundMessage, + input inboundRunInput, + content string, +) (Result, error) { + reply := newReplyPlan(input.Key, 0) + outbound := platform.OutboundMessage{ + TenantID: msg.TenantID, + BindingID: msg.BindingID, + Channel: msg.Channel, + SessionID: input.SessionID, + ReplyToPlatformMessageID: msg.PlatformMessageID, + Kind: platform.OutboundMessageKindText, + Content: content, + 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, reply.ResultRef, outbound); err != nil { + s.writeAuditTo(auditCtx, auditSink, auditFromMessage(msg, input.SessionID, input.InternalUserID, "outbound_error", err.Error(), input.Start, err)) + recordSpanError(replySpan, err) + return Result{}, err + } + if err := s.outboundStore.Enqueue( + replyCtx, + outbound, + channeladapter.RetryPolicyForBinding(runtime.Binding), + ); err != nil { + if _, markErr := s.idempotencyStore.MarkReplyFailed(replyCtx, input.Key, reply.ResultRef); markErr != nil { + recordSpanError(replySpan, markErr) + return Result{}, markErr + } + s.writeAuditTo(auditCtx, auditSink, auditFromMessage(msg, input.SessionID, input.InternalUserID, "outbound_error", err.Error(), input.Start, err)) + recordSpanError(replySpan, err) + return Result{}, err + } + 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, reply.InboundSequence, input.Start)) + s.writeMessageEvent(auditCtx, messageEventFromAssistant(msg, input.SessionID, reply.ResultRef, input.RequestID, reply.AssistantSequence, s.now())) + s.writeAuditTo(auditCtx, auditSink, auditFromMessage(msg, input.SessionID, input.InternalUserID, "completed", "", input.Start, nil)) + return Result{ + RequestID: input.RequestID, + SessionID: input.SessionID, + 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, + start time.Time, + err error, +) { + 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, + msg platform.InboundMessage, + start time.Time, + err error, +) { + s.writeAuditTo( + ctx, + auditSink, + auditFromMessage(msg, "", "", "reject", err.Error(), start, err), + ) +} + +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") + } + if s.idempotencyStore == nil { + return fmt.Errorf("gateway idempotency store is required") + } + 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") + } + if s.leaseStore == nil { + return fmt.Errorf("gateway session lease 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.ResultRef == "" || + (record.Status != platform.IdempotencyStatusCompleted && + record.Status != platform.IdempotencyStatusReplyFailed) { + 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 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 + } + 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 validateTextLimit(text string, limits platform.ChannelLimits) error { + if limits.MaxTextLength <= 0 { + return nil + } + if len([]rune(text)) > limits.MaxTextLength { + return ErrTextTooLong + } + return nil +} + +func validateFileLimits(msg platform.InboundMessage, limits platform.ChannelLimits) error { + if limits.FileMaxBytes <= 0 && len(limits.AllowedMIMETypes) == 0 { + return nil + } + for _, part := range msg.ContentParts { + if !contentPartHasFile(part) { + continue + } + 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, + 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 { + 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 runnerOutput{}, ctx.Err() + case next, ok := <-ch: + if !ok { + goto done + } + evt = next + } + if evt == nil || evt.Response == nil { + continue + } + if evt.Response.Usage != nil { + usage = evt.Response.Usage + } + if evt.IsTerminalError() { + return runnerOutput{}, 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) + } + } + } +done: + if strings.TrimSpace(final) != "" { + return runnerOutput{Content: strings.TrimSpace(final), Usage: usage}, nil + } + content := strings.TrimSpace(strings.Join(parts, "")) + if content == "" { + return runnerOutput{}, ErrRunnerResponseEmpty + } + return runnerOutput{Content: content, Usage: usage}, 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{ + AuditID: platform.AuditID(msg.TenantID, msg.AppID, msg.Channel, msg.BindingID, msg.PlatformMessageID, sessionID, decision), + 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), + TraceID: requestIDFor(msg), + Decision: decision, + DecisionReason: redactAuditReason(reason), + LatencyMS: time.Since(start).Milliseconds(), + CreatedAt: time.Now(), + } + if err != nil { + record.ErrorType = fmt.Sprintf("%T", err) + } + 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 "" + } + redactor, err := platform.NewRedactor() + if err != nil { + return reason + } + return redactor.Redact(reason) +} + +func (s *Service) writeAudit(ctx context.Context, record platform.AuditRecord) { + s.writeAuditTo(ctx, s.auditSink, record) +} + +func (s *Service) writeAuditTo( + ctx context.Context, + auditSink platform.AuditSink, + record platform.AuditRecord, +) { + if isNilInterfaceValue(auditSink) { + return + } + if err := record.Validate(); err != nil { + if !isAuditRedactionFailure(err) { + return + } + record = redactionFailedAuditRecord(record, err) + } + 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 { + 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 + } + return s.auditSink +} + +func isNilInterfaceValue(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 + } +} + +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, + 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 { + return itelemetry.TraceSafeHash(scope, value) +} + +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 { + if err == nil { + return "" + } + 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"}, + {ErrTextTooLong, "text_too_long"}, + {ErrFileTooLarge, "file_too_large"}, + {ErrMIMETypeNotAllowed, "mime_type_not_allowed"}, + {ErrRateLimited, "rate_limited"}, + {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) +} + +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") +} + +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 new file mode 100644 index 0000000000..ff4d2c9516 --- /dev/null +++ b/platform/gateway/service_test.go @@ -0,0 +1,3011 @@ +// +// 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" + "slices" + "strings" + "sync" + "testing" + "time" + + "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" +) + +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 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 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() + 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() + messageEvents := platform.NewInMemoryMessageEventSink() + svc := NewService( + registry, + idempotency, + NewOutboxBackedOutboundStore(outbox), + WithAuditSink(audit), + WithMessageEventSink(messageEvents), + ) + 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) + 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") + 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") + 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) + assert.Equal(t, platform.IdempotencyStatusCompleted, tenantBWeComResult.Status) + 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) + 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) { + 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.NotEmpty(t, runtimeAudit.Records()[0].SessionID) + assert.NotEmpty(t, runtimeAudit.Records()[0].InternalUserID) + assert.Empty(t, fallbackAudit.Records()) +} + +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) + 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{ + 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), + WithMessageEventSink(messageEvents), + ) + + _, 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.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) { + ctx := context.Background() + reader, restore := useGatewayMetrics(t) + defer restore() + 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) + + 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) + 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) + + 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() + 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) + 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, + 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 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() + 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 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 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() + 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() + 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.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) { + 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.NotEmpty(t, audit.Records()[0].SessionID) + assert.NotEmpty(t, audit.Records()[0].InternalUserID) + 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 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 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() + reader, restore := useGatewayMetrics(t) + defer restore() + + 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) + 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) { + 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() + 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 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 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() + reader, restore := useGatewayMetrics(t) + defer restore() + + 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() + 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.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) + 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) + + 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", + 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) { + 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() + 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() + runnerErr := errors.New("runner failed") + r := &recordingRunner{runErr: runnerErr} + registerRuntime(t, registry, "tenant-a", r) + store := platform.NewInMemoryIdempotencyStore() + messageEvents := platform.NewInMemoryMessageEventSink() + svc := NewService( + registry, + store, + NewInMemoryOutboundStore(), + WithMessageEventSink(messageEvents), + ) + 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) + assert.Empty(t, messageEvents.Events()) +} + +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() + r := &recordingRunner{chunks: []string{"he", "llo"}} + 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"} + + 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.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) + 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 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" + 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( + 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, "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.InDelta(t, 0.000025/18, record.ModelUnitPrice, 0.000000000001) + assert.InDelta(t, 0.000025, record.ModelCost, 0.000000000001) + assert.Zero(t, record.ToolCost) + assert.InDelta(t, 0.000025, record.TotalCost, 0.000000000001) +} + +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" + + 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(context.Background(), ch) + + require.NoError(t, err) + assert.Equal(t, "done", content) +} + +func TestCollectAssistantTextPrefersFinalFullMessage(t *testing.T) { + ch := make(chan *event.Event, 3) + ch <- chunkEvent("he", true) + ch <- chunkEvent("llo", true) + ch <- responseEvent("hello", true) + close(ch) + + content, err := collectAssistantText(context.Background(), ch) + + require.NoError(t, err) + 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)) + 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 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, + 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), + } +} + +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, + 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 + fencingToken int64 + runOptions agent.RunOptions +} + +type recordingRunner struct { + response string + chunks []string + usage *model.Usage + runErr error + 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, + sessionID string, + message model.Message, + runOpts ...agent.RunOption, +) (<-chan *event.Event, error) { + if r.runErr != nil { + 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, + fencingToken: fencingToken, + runOptions: runOptions, + }) + out := make(chan *event.Event, 2) + go func() { + defer close(out) + if len(r.chunks) > 0 { + for i, chunk := range r.chunks { + 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 + } + evt := responseEvent(r.response, true) + if r.usage != nil { + evt.Response.Usage = r.usage + } + out <- evt + }() + return out, nil +} + +func (r *recordingRunner) Close() error { + return nil +} + +type blockingRunner struct { + mu sync.Mutex + started chan struct{} + startedOnce sync.Once + done chan string + calls []runnerCall +} + +type cancelingRunner struct { + cancel func() + runErr error + 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 +} + +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) + } + runOptions := runOptionsFromOptions(runOpts...) + r.calls = append(r.calls, runnerCall{ + userID: userID, + sessionID: sessionID, + message: message, + requestID: runOptions.RequestID, + runOptions: runOptions, + }) + r.startedOnce.Do(func() { + 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 (r *cancelingRunner) 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, + }) + r.cancel() + return nil, r.runErr +} + +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{} + 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) + runOptions := runOptionsFromOptions(runOpts...) + r.calls = append(r.calls, runnerCall{ + userID: userID, + sessionID: sessionID, + message: message, + requestID: runOptions.RequestID, + runOptions: runOptions, + }) + 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 + token int64 +} + +func (l *recordingLease) FencingToken() int64 { + if l.token == 0 { + return 1 + } + return l.token +} + +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", + "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 { + 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 +} + +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") +} + +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 useGatewayMetrics(t *testing.T) (*sdkmetric.ManualReader, func()) { + t.Helper() + reader := sdkmetric.NewManualReader() + provider := sdkmetric.NewMeterProvider(sdkmetric.WithReader(reader)) + + originalProvider := itelemetry.MeterProvider + originalMeter := itelemetry.GatewayMeter + originalBudgetCounter := itelemetry.GatewayMetricBudgetDeniedTotal + originalRateLimitCounter := itelemetry.GatewayMetricRateLimitedTotal + originalIdempotencyHitCounter := itelemetry.GatewayMetricIdempotencyHitTotal + + itelemetry.MeterProvider = provider + itelemetry.GatewayMeter = provider.Meter(metrics.MeterNameGateway) + var err error + itelemetry.GatewayMetricBudgetDeniedTotal, err = + itelemetry.GatewayMeter.Int64Counter(metrics.MetricGatewayBudgetDeniedTotal) + require.NoError(t, err) + 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 + } +} + +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 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 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 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() { + 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() { + if string(kv.Key) == key { + require.Equal(t, value, kv.Value.AsString()) + return + } + } + 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/platform/gateway/store.go b/platform/gateway/store.go new file mode 100644 index 0000000000..8965f2cb3c --- /dev/null +++ b/platform/gateway/store.go @@ -0,0 +1,91 @@ +// +// 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/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) +} + +// InMemoryOutboundStore is a concurrency-safe outbound store for tests and demos. +type InMemoryOutboundStore struct { + mu sync.Mutex + messages map[string]platform.OutboundMessage + outbox channeladapter.OutboxStore +} + +// NewInMemoryOutboundStore creates an in-memory outbound store. +func NewInMemoryOutboundStore() *InMemoryOutboundStore { + return &InMemoryOutboundStore{ + messages: make(map[string]platform.OutboundMessage), + } +} + +// 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, + 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 +} + +// 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, + 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 +} diff --git a/platform/gray.go b/platform/gray.go new file mode 100644 index 0000000000..cf7cdaa1e6 --- /dev/null +++ b/platform/gray.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 ( + "fmt" + "hash/fnv" + "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) + 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 +} + +// 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_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 +} diff --git a/platform/gray_test.go b/platform/gray_test.go new file mode 100644 index 0000000000..425ee314da --- /dev/null +++ b/platform/gray_test.go @@ -0,0 +1,197 @@ +// +// 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") + } +} + +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 +} diff --git a/platform/idempotency.go b/platform/idempotency.go new file mode 100644 index 0000000000..5dfee8bdd9 --- /dev/null +++ b/platform/idempotency.go @@ -0,0 +1,162 @@ +// +// 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" + "fmt" + "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, 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 { + 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, IdempotencyStatusProcessing, 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, + from IdempotencyStatus, + 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 + } + 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, + 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 err := validateRoutingIdentifier("tenant_id", record.TenantID, ErrTenantIDRequired); err != nil { + return "", err + } + if err := validateRoutingIdentifier("channel", record.Channel, ErrChannelRequired); err != nil { + return "", err + } + if err := validateRoutingIdentifier("account_id", record.AccountID, ErrAccountIDRequired); err != nil { + return "", err + } + if err := validateRoutingIdentifier("platform_message_id", record.PlatformMessageID, ErrPlatformMessageIDRequired); err != nil { + return "", err + } + return IdempotencyKey( + record.TenantID, + record.Channel, + record.AccountID, + record.PlatformMessageID, + ), nil +} diff --git a/platform/identity.go b/platform/identity.go new file mode 100644 index 0000000000..89edb68a61 --- /dev/null +++ b/platform/identity.go @@ -0,0 +1,206 @@ +// +// 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" + "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 +} + +// InternalUserID returns a stable tenant-scoped user identifier. +func InternalUserID(tenantID, channel, externalUserID string) string { + 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 stableID( + "user_hash", + stableIDPart{"tenant", tenantID}, + stableIDPart{"channel", channel}, + stableIDPart{"user", 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 stableID( + "idem", + stableIDPart{"tenant", tenantID}, + stableIDPart{"channel", channel}, + stableIDPart{"account", accountID}, + stableIDPart{"message", 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 + } + 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/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 err := validateRoutingIdentifier("tenant_id", tenantID, ErrTenantIDRequired); err != nil { + return "", err + } + 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 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, + 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 err := validateRoutingIdentifier("external_user_id", externalUserID, ErrExternalUserIDRequired); err != nil { + return nil, err + } + return []stableIDPart{ + {"conversation_type", string(ConversationTypeDM)}, + {"user", externalUserID}, + }, nil + case ConversationTypeGroup: + 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 err := validateRoutingIdentifier("external_group_id", externalGroupID, ErrExternalGroupIDRequired); err != nil { + return nil, err + } + if err := validateRoutingIdentifier("thread_id", threadID, ErrThreadIDRequired); err != nil { + return nil, err + } + return []stableIDPart{ + {"conversation_type", string(ConversationTypeThread)}, + {"group", externalGroupID}, + {"thread", threadID}, + }, nil + case "": + return nil, ErrConversationTypeRequired + default: + return nil, ErrInvalidConversationType + } +} + +func stableID(prefix string, parts ...stableIDPart) string { + hash := sha256.New() + writeStablePart(hash, "prefix", prefix) + for _, part := range parts { + writeStablePart(hash, part.name, 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) +} + +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/inmemory_sink.go b/platform/inmemory_sink.go new file mode 100644 index 0000000000..07991942bf --- /dev/null +++ b/platform/inmemory_sink.go @@ -0,0 +1,82 @@ +// +// 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" + "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 + 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 { + 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() + 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 +} + +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/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 +} diff --git a/platform/message_event.go b/platform/message_event.go new file mode 100644 index 0000000000..38e8d0016f --- /dev/null +++ b/platform/message_event.go @@ -0,0 +1,41 @@ +// +// 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" +) + +// 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 bounded message event sink for tests and demos. +type InMemoryMessageEventSink struct { + events inMemoryRecords[MessageEvent] +} + +// NewInMemoryMessageEventSink creates an in-memory message event sink. +func NewInMemoryMessageEventSink(options ...InMemorySinkOption) *InMemoryMessageEventSink { + return &InMemoryMessageEventSink{ + events: newInMemoryRecords[MessageEvent](options...), + } +} + +// WriteMessageEvent writes one message event. +func (s *InMemoryMessageEventSink) WriteMessageEvent(ctx context.Context, event MessageEvent) error { + return s.events.append(ctx, event, MessageEvent.Validate) +} + +// Events returns a snapshot of written message events. +func (s *InMemoryMessageEventSink) Events() []MessageEvent { + return s.events.snapshot() +} diff --git a/platform/message_event_test.go b/platform/message_event_test.go new file mode 100644 index 0000000000..dd1fa88565 --- /dev/null +++ b/platform/message_event_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 ( + "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 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 + 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/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..afd1752a96 --- /dev/null +++ b/platform/migration_test.go @@ -0,0 +1,106 @@ +// +// 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 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", + ProfileID: "profile", + Namespace: "tenant/tenant", + MigrationMode: "dual-read", + } + if err := profile.Validate(); err == nil { + t.Fatalf("expected invalid migration mode error") + } +} + +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") + } + for _, mode := range []StorageMigrationMode{ + StorageMigrationModeDualWrite, + StorageMigrationModeShadowRead, + StorageMigrationModeCutover, + StorageMigrationModeRollback, + } { + if !IsActiveStorageMigrationMode(mode) { + t.Fatalf("%q should be active migration", mode) + } + } +} 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/operational_action_audit.go b/platform/operational_action_audit.go new file mode 100644 index 0000000000..14d72ed9a5 --- /dev/null +++ b/platform/operational_action_audit.go @@ -0,0 +1,214 @@ +// +// 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") + } + 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 +} + +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..bba119dec9 --- /dev/null +++ b/platform/operational_action_audit_test.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 ( + "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 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() + + 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(), + } +} diff --git a/platform/redaction.go b/platform/redaction.go new file mode 100644 index 0000000000..d402a26c6b --- /dev/null +++ b/platform/redaction.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 ( + "regexp" + "strings" +) + +var defaultRedactionPatterns = []*regexp.Regexp{ + 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*(?: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*=\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/?#]+`), + 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:") { + 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] + "****" + } + if strings.Contains(match, "://") && strings.Contains(match, "@") { + if redacted, ok := redactURLUserinfo(match); ok { + return redacted + } + } + 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:] +} + +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/secret_rotation_status.go b/platform/secret_rotation_status.go new file mode 100644 index 0000000000..3f6c274f1f --- /dev/null +++ b/platform/secret_rotation_status.go @@ -0,0 +1,291 @@ +// +// 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") + } + 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 validateSecretRotationStatusGate(r) +} + +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") + } + 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, + 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..87e5c5115b --- /dev/null +++ b/platform/secret_rotation_status_test.go @@ -0,0 +1,272 @@ +// +// 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 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 { + 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(), + } +} 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/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..45dbb2a4ea --- /dev/null +++ b/platform/storagerouter/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 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") + // 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") + // 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 new file mode 100644 index 0000000000..e317520066 --- /dev/null +++ b/platform/storagerouter/router.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 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 + 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) + 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. +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 ErrBackendIDRequired + } + 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 + } + // Defensive for future persistent backends that may not key by tenant_id. + if profile.TenantID != tenantID { + return platform.StorageProfile{}, ErrTenantMismatch + } + 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, + 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 +} + +// 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, + 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" + resourceSummary resourceKind = "summary" + 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 + } + // Defensive for future persistent backends that may not key by tenant_id. + 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 resourceSummary: + return strings.TrimSpace(profile.SummaryBackend) + 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 "" + } +} + +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 new file mode 100644 index 0000000000..289954203d --- /dev/null +++ b/platform/storagerouter/router_test.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 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" + "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" +) + +var _ Router = (*InMemoryRouter)(nil) + +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() + 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", + Session: sessionSvc, + Memory: memorySvc, + Artifact: artifactSvc, + 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") + 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, summarySvc, gotSummary) + assert.Same(t, memorySvc, gotMemory) + assert.Same(t, artifactSvc, gotArtifact) + assert.Same(t, knowledgeSvc, gotKnowledge) + 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() + 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 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") + 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, + SummaryBackend: 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 +} + +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 new file mode 100644 index 0000000000..a1c369c467 --- /dev/null +++ b/platform/storagerouter/status.go @@ -0,0 +1,199 @@ +// +// 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" + "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 + +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: resourceSummary, resource: platform.BackendMigrationResourceSummary}, + {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 + } + // 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) + 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 resourceSummary: + return backend.Summary != 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 := statusBackendIDRedactor() + 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 +} + +func statusBackendIDRedactor() (*platform.Redactor, error) { + statusRedactorOnce.Do(func() { + statusRedactor, statusRedactorErr = platform.NewRedactor() + }) + return statusRedactor, statusRedactorErr +} diff --git a/platform/storagerouter/status_test.go b/platform/storagerouter/status_test.go new file mode 100644 index 0000000000..17b112eb78 --- /dev/null +++ b/platform/storagerouter/status_test.go @@ -0,0 +1,165 @@ +// +// 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(), + Summary: 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, 6, summary.ReadyCount) + assert.Equal(t, 0, summary.MissingCount) + require.Len(t, summary.Resources, 6) + 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, 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) + 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) +} 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), + } +} 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/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 new file mode 100644 index 0000000000..dca962a2b0 --- /dev/null +++ b/platform/toolpolicy/policy.go @@ -0,0 +1,815 @@ +// +// 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" + "regexp" + "strconv" + "strings" + "time" + + 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/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 +} + +// 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 + 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{} + +// 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) + +// 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 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 { + 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 + } + 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) + } + 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 +} + +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 == "" { + 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 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 { + return tool.PermissionDecision{}, err + } + if err := p.writeAudit(ctx, summary); err != nil { + return tool.PermissionDecision{}, err + } + } + 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) && + fingerprint.Metadata == req.Metadata +} + +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 { + return nil, nil + } + req := &tool.PermissionRequest{ + ToolName: args.ToolName, + ToolCallID: args.ToolCallID, + Declaration: args.Declaration, + Arguments: args.Arguments, + Metadata: args.Metadata, + } + decision, reason, audit := p.decideNameOnly(req, req.ToolName) + if audit { + summary, err := p.ApprovalSummary(req, decision, reason) + if err != nil { + return nil, err + } + if err := p.writeAudit(ctx, summary); 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)} + if policy.DangerousToolAction == platform.DangerousToolActionAsk { + opts = append(opts, approval.WithMetadataRiskPolicy(approval.ToolPolicyRequireApproval)) + } + 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, + Metadata: req.Action.Metadata, + } + decision, reason, audit := r.policy.decideReviewer(permissionReq, permissionReq.ToolName) + if audit { + summary, err := r.policy.ApprovalSummary(permissionReq, decision, reason) + if err != nil { + return nil, err + } + if err := r.policy.writeAudit(ctx, summary); err != nil { + return nil, err + } + } + 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) { + 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 + } + 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 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, "": + 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) highRiskForDecision( + req *tool.PermissionRequest, + name string, + opts decisionOptions, +) bool { + if (opts.includeNameRisk || opts.includeMetadataRisk) && + contains(normalizedList(p.policy.HighRiskTools), name) { + return true + } + if !opts.includeMetadataRisk || req == nil || req.Metadata == (tool.ToolMetadata{}) { + return false + } + return req.Metadata.Destructive || !req.Metadata.ReadOnly || req.Metadata.OpenWorld +} + +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 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 + } + 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 +} + +// 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, + 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) + 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, + 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 + } + return summary, nil +} + +// 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.validateToolBudgetRemaining(); 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") + } + 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") + } + 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 + } + return nil +} + +func (s ApprovalSummary) validateArguments() error { + 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") + } + 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: + 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) + } + return nil +} + +func (p *Policy) writeAudit(ctx context.Context, summary ApprovalSummary) error { + if p.audit == nil { + return nil + } + detailRef := summary.DetailRef() + record := platform.AuditRecord{ + 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, + } + 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 { + 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 +} + +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{ + "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.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") + } + 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 "", 0 + } + sum := sha256.Sum256(args) + 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 { + 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 new file mode 100644 index 0000000000..9f2c968a87 --- /dev/null +++ b/platform/toolpolicy/policy_test.go @@ -0,0 +1,982 @@ +// +// 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" + "errors" + "strings" + "testing" + "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" +) + +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.AuditID == "" || + 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.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 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( + 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 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{ + TenantID: "tenant", + AppID: "app", + PolicyID: "policy", + 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) + } + + 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") { + t.Fatalf("expected allow approval invariant, got %v", err) + } +} + +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( + t, + platform.ToolPolicy{ + TenantID: "tenant", + AppID: "app", + 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 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) + } +} + +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(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) { + reader, restore := usePolicyAuditMetrics(t) + defer restore() + + p := newPolicy( + t, + platform.ToolPolicy{ + TenantID: "tenant", + AppID: "app", + 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) + } + + 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) { + 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") + } + 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) + } + 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 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"}, + })) + 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(defaultPolicy(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 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"}, + 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 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) + p, err := New(policy, opts...) + if err != nil { + t.Fatalf("New: %v", err) + } + 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 { + 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 +} + +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/platform/types.go b/platform/types.go new file mode 100644 index 0000000000..f1ceb2bb46 --- /dev/null +++ b/platform/types.go @@ -0,0 +1,496 @@ +// +// 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" +) + +// 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 + +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" +) + +// 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 + +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 +} + +// 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 + 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 + 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. +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 +} + +// 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 + 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 +} + +// 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/types_test.go b/platform/types_test.go new file mode 100644 index 0000000000..67711d9cda --- /dev/null +++ b/platform/types_test.go @@ -0,0 +1,991 @@ +// +// 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", + BindingID: "telegram-bot-1", + 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.HasPrefix(sessionA, "ses_") || strings.Contains(sessionA, ":") { + t.Fatalf("session id should be opaque and delimiter-safe, got %q", sessionA) + } +} + +func TestSessionIDForInboundSupportsGroupAndThread(t *testing.T) { + 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", "binding", "wecom", "bot", 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", "binding", "telegram", "bot", ConversationTypeThread, "user", "chat", "topic/7") + if err != nil { + t.Fatalf("thread session: %v", err) + } + 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 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", "binding", "chat:dm:user", "bot", ConversationTypeDM, "leaf", "", "") + if err != nil { + t.Fatalf("session A: %v", err) + } + sessionB, err := SessionID("tenant", "app", "binding", "chat", "bot", 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) + } + } +} + +func TestValidateInboundRequiresGroupForGroupConversation(t *testing.T) { + msg := InboundMessage{ + TenantID: "tenant", + AppID: "app", + BindingID: "binding", + Channel: "telegram", + ChannelAccountID: "bot", + PlatformMessageID: "msg", + ExternalUserID: "user", + ConversationType: ConversationTypeGroup, + MessageType: MessageTypeText, + } + if err := msg.Validate(); !errors.Is(err, ErrExternalGroupIDRequired) { + t.Fatalf("expected ErrExternalGroupIDRequired, got %v", err) + } +} + +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 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: "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 { + 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", + 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") + 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 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 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{ + 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 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") + } + 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() + 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", + 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 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", + 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 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") + 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 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 { + 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://****@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{ + 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") + } +} + +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) + } +} diff --git a/platform/usage.go b/platform/usage.go new file mode 100644 index 0000000000..83c0f229ec --- /dev/null +++ b/platform/usage.go @@ -0,0 +1,41 @@ +// +// 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" +) + +// 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 bounded usage sink for tests and demos. +type InMemoryUsageSink struct { + records inMemoryRecords[UsageRecord] +} + +// NewInMemoryUsageSink creates an in-memory usage sink. +func NewInMemoryUsageSink(options ...InMemorySinkOption) *InMemoryUsageSink { + return &InMemoryUsageSink{ + records: newInMemoryRecords[UsageRecord](options...), + } +} + +// WriteUsage writes one usage record. +func (s *InMemoryUsageSink) WriteUsage(ctx context.Context, record UsageRecord) error { + return s.records.append(ctx, record, UsageRecord.Validate) +} + +// Records returns a snapshot of written usage records. +func (s *InMemoryUsageSink) Records() []UsageRecord { + return s.records.snapshot() +} diff --git a/platform/usage_record_test.go b/platform/usage_record_test.go new file mode 100644 index 0000000000..619db96dad --- /dev/null +++ b/platform/usage_record_test.go @@ -0,0 +1,199 @@ +// +// 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 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 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 + 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 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 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() + 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", + AppID: "app", + UserIDHash: UserIDHash("tenant", "telegram", "external"), + SessionID: "session", + RequestID: "request", + ModelName: "gpt-test", + ToolName: "knowledge_search", + TraceID: "trace", + } +} 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 +} diff --git a/platform/validation.go b/platform/validation.go new file mode 100644 index 0000000000..2bce624687 --- /dev/null +++ b/platform/validation.go @@ -0,0 +1,566 @@ +// +// 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" + "net/url" + "strconv" + "strings" + "unicode" +) + +var rawSecretPrefixes = []string{ + "sk-", + "xoxb-", + "xoxp-", + "ya29.", + "ghp_", + "github_pat_", + "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 == "" { + 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 err := validateRoutingIdentifier("tenant_id", t.TenantID, ErrTenantIDRequired); err != nil { + return err + } + 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 err := validateRoutingIdentifier("tenant_id", a.TenantID, ErrTenantIDRequired); err != nil { + return err + } + 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") + } + switch a.Status { + case "", AppStatusActive, AppStatusSuspended, AppStatusDeleted: + return nil + default: + return fmt.Errorf("invalid app status %q", a.Status) + } +} + +// Validate checks that an app config version is safe to store and route. +func (v AppConfigVersion) Validate() error { + if err := validateRoutingIdentifier("tenant_id", v.TenantID, ErrTenantIDRequired); err != nil { + return err + } + if err := validateRoutingIdentifier("app_id", v.AppID, ErrAppIDRequired); err != nil { + return err + } + 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 err := validateRoutingIdentifier("tenant_id", p.TenantID, ErrTenantIDRequired); err != nil { + return err + } + 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 + } + if err := validateSecretReference("api_key_ref", p.APIKeyRef); err != nil { + return err + } + if _, err := ParseModelCostPolicy(p); err != nil { + return err + } + return nil +} + +// Validate checks that a binding has safe routing and secret references. +func (b ChannelBinding) Validate() error { + if err := validateRoutingIdentifier("tenant_id", b.TenantID, ErrTenantIDRequired); err != nil { + return err + } + if err := validateRoutingIdentifier("app_id", b.AppID, ErrAppIDRequired); err != nil { + return err + } + if err := validateRoutingIdentifier("binding_id", b.BindingID, ErrBindingIDRequired); err != nil { + return err + } + if err := validateRoutingIdentifier("channel", b.Channel, ErrChannelRequired); err != nil { + return err + } + if err := validateRoutingIdentifier("account_id", b.AccountID, ErrAccountIDRequired); err != nil { + return err + } + 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 err := validateRoutingIdentifier("tenant_id", m.TenantID, ErrTenantIDRequired); err != nil { + return err + } + 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 err := validateRoutingIdentifier("channel", m.Channel, ErrChannelRequired); err != nil { + return err + } + if err := validateRoutingIdentifier("channel_account_id", m.ChannelAccountID, ErrAccountIDRequired); err != nil { + return err + } + if err := validateRoutingIdentifier("platform_message_id", m.PlatformMessageID, ErrPlatformMessageIDRequired); err != nil { + return err + } + if err := validateInboundMessageType(m); err != nil { + return err + } + if m.MessageType == MessageTypeEvent { + return nil + } + if err := validateRoutingIdentifier("external_user_id", m.ExternalUserID, ErrExternalUserIDRequired); err != nil { + return err + } + switch m.ConversationType { + case ConversationTypeDM: + return nil + case ConversationTypeGroup: + if err := validateRoutingIdentifier("external_group_id", m.ExternalGroupID, ErrExternalGroupIDRequired); err != nil { + return err + } + return nil + case ConversationTypeThread: + if err := validateRoutingIdentifier("external_group_id", m.ExternalGroupID, ErrExternalGroupIDRequired); err != nil { + return err + } + if err := validateRoutingIdentifier("thread_id", m.ThreadID, ErrThreadIDRequired); err != nil { + return err + } + return nil + case "": + return ErrConversationTypeRequired + default: + return ErrInvalidConversationType + } +} + +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 err := validateRoutingIdentifier("tenant_id", p.TenantID, ErrTenantIDRequired); err != nil { + return err + } + 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 + } + if _, err := NormalizeStorageMigrationMode(p.MigrationMode); err != nil { + return err + } + 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 { + return err + } + 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") + } + 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 +} + +// 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") + } + for field, value := range map[string]string{ + "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 + } + } + 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 +} + +// Validate checks that a message event has required identity and safe trace metadata. +func (e MessageEvent) Validate() error { + if err := validateRoutingIdentifier("tenant_id", e.TenantID, ErrTenantIDRequired); err != nil { + return err + } + if err := validateRoutingIdentifier("app_id", e.AppID, ErrAppIDRequired); err != nil { + return err + } + 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 err := validateRoutingIdentifier("tenant_id", r.TenantID, ErrTenantIDRequired); err != nil { + return err + } + if err := validateRoutingIdentifier("app_id", r.AppID, ErrAppIDRequired); err != nil { + return err + } + 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 == "" { + 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 == "" { + return nil + } + if strings.Contains(value, "=") || + hasInlineURLCredential(value) || + looksLikeRawSecret(value) { + return fmt.Errorf("%s: %w", field, ErrInlineSecretRejected) + } + return nil +} + +func hasInlineURLCredential(value string) bool { + parsed, err := url.Parse(value) + if err != nil || parsed.Scheme == "" || parsed.Host == "" || parsed.User == nil { + return false + } + if parsed.User.Username() != "" { + return true + } + _, hasPassword := parsed.User.Password() + return hasPassword +} + +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 +} diff --git a/platform/worker/builder.go b/platform/worker/builder.go new file mode 100644 index 0000000000..13aaf3d59e --- /dev/null +++ b/platform/worker/builder.go @@ -0,0 +1,381 @@ +// +// 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/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" +) + +// 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 + // 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 + // Plugins contains runner-scoped plugins assembled by worker governance. + Plugins []plugin.Plugin +} + +// 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 + 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, +) (*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) { + return nil, ErrStorageRouterRequired + } + if isNilDependency(factory) { + return nil, ErrAgentFactoryRequired + } + builder := &RuntimeBuilder{ + router: router, + factory: factory, + } + for _, opt := range opts { + if opt != nil { + opt(builder) + } + } + return builder, 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 + } + 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 { + 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) + } + permissionPolicy, err := compileToolPolicy(resolvedPolicy, auditSink) + if err != nil { + return gateway.Runtime{}, err + } + var toolFilter tool.FilterFunc + if permissionPolicy != nil { + toolFilter = permissionPolicy.ToolFilter() + } + plugins, err := buildToolGovernancePlugins(resolvedPolicy, auditSink) + if err != nil { + return gateway.Runtime{}, err + } + + dependencies := AgentDependencies{ + Tenant: tenant, + App: app, + Binding: binding, + Storage: storage, + Session: sessionService, + Memory: memoryService, + Artifact: artifactService, + Knowledge: knowledgeService, + Audit: auditSink, + ToolPolicy: resolvedPolicy, + ToolFilter: toolFilter, + ToolPermissionPolicy: permissionPolicy, + Plugins: plugins, + } + 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 + } + + 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, + Binding: binding, + Runner: runner.NewRunner( + storage.Scope().ScopedAppName(app.AppID), + ag, + runnerOptions..., + ), + Audit: auditSink, + ToolFilter: toolFilter, + ToolPermissionPolicy: permissionPolicy, + } + 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 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) { + defaultPolicy := approval.ToolPolicySkipApproval + if len(normalizedToolNames(policy.ToolWhitelist)) > 0 { + defaultPolicy = approval.ToolPolicyDenied + } + opts := []approval.Option{ + 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 := 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 + } + if containsToolName(denied, name) { + continue + } + 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 +} + +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 + } + 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..ca2316c59e --- /dev/null +++ b/platform/worker/errors.go @@ -0,0 +1,34 @@ +// +// 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") + // 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..bacb19d867 --- /dev/null +++ b/platform/worker/governance_test.go @@ -0,0 +1,732 @@ +// +// 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" + "encoding/json" + "reflect" + "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" + "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" +) + +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 := NewRuntimeBuilderWithOptions( + 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 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 := NewRuntimeBuilderWithOptions( + 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 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 := NewRuntimeBuilderWithOptions( + 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" + router := &countingRouter{Router: storagerouter.NewInMemoryRouter()} + builder, err := NewRuntimeBuilderWithOptions( + 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 := NewRuntimeBuilderWithOptions( + 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 := NewRuntimeBuilderWithOptions( + 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} +} + +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 4cfca29ac9..db373d4652 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" @@ -27,7 +29,11 @@ type Plugin struct { reviewer review.Reviewer defaultToolPolicy ToolPolicy toolPolicies map[string]ToolPolicy + metadataPolicy ToolPolicy tokenCounter model.TokenCounter + auditSink platform.AuditSink + approverUserID string + now func() time.Time } // New creates a new approval plugin. @@ -36,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") @@ -47,12 +56,19 @@ 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, + metadataPolicy: opts.metadataPolicy, tokenCounter: model.NewSimpleTokenCounter(), + auditSink: opts.auditSink, + approverUserID: strings.TrimSpace(opts.approverUserID), + now: opts.now, }, nil } @@ -74,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{ @@ -95,6 +111,24 @@ 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 + } + reportApprovalRequiredMetric(ctx, args) decision, err := p.reviewer.Review(ctx, req) if err != nil { log.ErrorfContext( @@ -119,25 +153,64 @@ 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, + 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", 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", 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, + CustomResult: tool.ApprovalDeniedResultFor(args.ToolName, denyMessage), }, nil default: return &tool.BeforeToolResult{ @@ -147,8 +220,33 @@ func (p *Plugin) beforeTool() tool.BeforeToolCallbackStructured { } } -func (p *Plugin) resolveToolPolicy(toolName string) ToolPolicy { - if policy, ok := p.toolPolicies[toolName]; ok { +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 + } + 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 @@ -158,6 +256,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 @@ -165,3 +266,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 db25fe2937..f4e9e3a0c6 100644 --- a/plugin/guardrail/approval/approval_test.go +++ b/plugin/guardrail/approval/approval_test.go @@ -13,17 +13,25 @@ import ( "errors" "fmt" "testing" + "time" "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" "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" "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" ) @@ -60,18 +68,46 @@ 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) 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) @@ -84,9 +120,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) { @@ -212,7 +250,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) @@ -252,7 +296,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.", @@ -260,6 +310,305 @@ 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, + 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_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() + 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.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) + 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, 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) + 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_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 @@ -327,7 +676,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) { @@ -360,7 +713,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( @@ -370,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", @@ -421,18 +820,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) } @@ -516,3 +918,69 @@ 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) +} + +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 new file mode 100644 index 0000000000..9866ee4aa7 --- /dev/null +++ b/plugin/guardrail/approval/audit.go @@ -0,0 +1,214 @@ +// +// 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" + 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" +) + +// AuditContext carries trusted platform context for tool approval audit. +type AuditContext struct { + TenantID string + AppID string + RequestID string + TraceID string +} + +type auditContextKey struct{} +type approvedToolCallContextKey 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) +} + +// 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 + Metadata tool.ToolMetadata +} + +// 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, + Metadata: args.Metadata, + } +} + +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 + } + 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) { + 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 "" + } + 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 { + 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 hex.EncodeToString(sum[:]), 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..740aa1ae89 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,10 @@ type options struct { reviewer review.Reviewer defaultToolPolicy ToolPolicy toolPolicies map[string]ToolPolicy + metadataPolicy ToolPolicy + auditSink platform.AuditSink + approverUserID string + now func() time.Time } func newOptions(opts ...Option) *options { @@ -27,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 { @@ -66,3 +76,33 @@ func WithToolPolicy(name string, policy ToolPolicy) Option { opts.toolPolicies[name] = policy } } + +// 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) { + 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 + } + } +} 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/runner/diagnostics.go b/runner/diagnostics.go index d3e31637e4..1c1c53abd5 100644 --- a/runner/diagnostics.go +++ b/runner/diagnostics.go @@ -10,6 +10,8 @@ package runner import ( "context" + "errors" + "fmt" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/codes" @@ -18,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" @@ -95,8 +98,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 +114,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 +130,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 +157,17 @@ func runnerSessionAttrs(key session.Key, sess *session.Session) []attribute.KeyV return attrs } +func runnerTraceSafeHash(scope string, value string) string { + return itelemetry.TraceSafeHash(scope, value) +} + +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() 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 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..91ff44af8b 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" @@ -33,12 +34,13 @@ 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) - defer span.End() - 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 !isummary.NewSummaryDispatchPolicy( s.opts.summaryFilterAllowlist, s.opts.shouldCascadeFullSessionSummary(), diff --git a/telemetry/metric/metric.go b/telemetry/metric/metric.go index 0bf6e6eaa8..7993136860 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, @@ -115,6 +122,15 @@ 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 := initAuditMetrics(mp); err != nil { + return err + } + if err := initGatewayMetrics(mp); err != nil { + return err + } if err := initInvokeAgentMetrics(mp); err != nil { return err } @@ -197,6 +213,76 @@ 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 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 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) + } + 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) + } + 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 +} + 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..cbfcfccaca 100644 --- a/telemetry/metric/metric_test.go +++ b/telemetry/metric/metric_test.go @@ -294,8 +294,26 @@ func TestInitMeterProvider(t *testing.T) { // Save original meter provider originalMP := itelemetry.MeterProvider + originalToolApprovalMeter := itelemetry.ToolApprovalMeter + originalToolApprovalRequired := itelemetry.ToolApprovalMetricRequiredTotal + originalToolPermissionDenied := itelemetry.ExecuteToolMetricToolPermissionDeniedTotal + originalAuditMeter := itelemetry.AuditMeter + originalAuditWriteFailed := itelemetry.AuditMetricWriteFailedTotal + originalGatewayMeter := itelemetry.GatewayMeter + originalGatewayBudgetDenied := itelemetry.GatewayMetricBudgetDeniedTotal + originalGatewayRateLimited := itelemetry.GatewayMetricRateLimitedTotal + originalGatewayIdempotencyHit := itelemetry.GatewayMetricIdempotencyHitTotal defer func() { itelemetry.MeterProvider = originalMP + itelemetry.ToolApprovalMeter = originalToolApprovalMeter + itelemetry.ToolApprovalMetricRequiredTotal = originalToolApprovalRequired + itelemetry.ExecuteToolMetricToolPermissionDeniedTotal = originalToolPermissionDenied + itelemetry.AuditMeter = originalAuditMeter + itelemetry.AuditMetricWriteFailedTotal = originalAuditWriteFailed + itelemetry.GatewayMeter = originalGatewayMeter + itelemetry.GatewayMetricBudgetDeniedTotal = originalGatewayBudgetDenied + itelemetry.GatewayMetricRateLimitedTotal = originalGatewayRateLimited + itelemetry.GatewayMetricIdempotencyHitTotal = originalGatewayIdempotencyHit }() // Create a test meter provider @@ -348,9 +366,36 @@ 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") } + if itelemetry.ToolApprovalMeter == nil { + t.Error("ToolApprovalMeter was not created") + } + 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.GatewayMeter == nil { + t.Error("GatewayMeter was not created") + } + if itelemetry.GatewayMetricBudgetDeniedTotal == nil { + t.Error("GatewayMetricBudgetDeniedTotal was not created") + } + 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") } @@ -391,6 +436,136 @@ 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 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 + 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 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") { + 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) + } + + 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) + } + + 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) { originalWorkflowMeter := itelemetry.WorkflowMeter originalWorkflowOpDur := itelemetry.WorkflowMetricGenAIClientOperationDuration diff --git a/telemetry/semconv/metrics/metrics.go b/telemetry/semconv/metrics/metrics.go index f35cbf0901..6f17fff4b2 100644 --- a/telemetry/semconv/metrics/metrics.go +++ b/telemetry/semconv/metrics/metrics.go @@ -59,6 +59,18 @@ 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" + // MetricToolPermissionDeniedTotal records tool calls denied by permission checks. + 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" + // 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 //////////////////////// @@ -79,4 +91,10 @@ 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" + // 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 bec8f780f5..94c28d11bb 100644 --- a/telemetry/semconv/trace/trace.go +++ b/telemetry/semconv/trace/trace.go @@ -42,10 +42,34 @@ 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. 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" + // 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" + // 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" + // 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" @@ -137,6 +161,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" 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/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. 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 } 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) + } +}