You've already forked agentic-coding-workflow
46 lines
1.1 KiB
Go
46 lines
1.1 KiB
Go
package mcp
|
|
|
|
import (
|
|
"context"
|
|
|
|
"github.com/yan1h/agent-coding-workflow/internal/audit"
|
|
)
|
|
|
|
type mcpTokenIDCtxKey struct{}
|
|
|
|
func withMcpTokenID(ctx context.Context, tokenID string) context.Context {
|
|
return context.WithValue(ctx, mcpTokenIDCtxKey{}, tokenID)
|
|
}
|
|
|
|
func mcpTokenIDFromCtx(ctx context.Context) string {
|
|
if v, ok := ctx.Value(mcpTokenIDCtxKey{}).(string); ok {
|
|
return v
|
|
}
|
|
return ""
|
|
}
|
|
|
|
type wrappedRecorder struct {
|
|
inner audit.Recorder
|
|
}
|
|
|
|
// AuditWrap wraps an audit.Recorder so that every Record call automatically
|
|
// injects "via_mcp" and "mcp_token_id" into the entry's Metadata when the
|
|
// context carries an MCP token ID. Returns nil when inner is nil.
|
|
func AuditWrap(inner audit.Recorder) audit.Recorder {
|
|
if inner == nil {
|
|
return nil
|
|
}
|
|
return &wrappedRecorder{inner: inner}
|
|
}
|
|
|
|
func (r *wrappedRecorder) Record(ctx context.Context, e audit.Entry) error {
|
|
if tid := mcpTokenIDFromCtx(ctx); tid != "" {
|
|
if e.Metadata == nil {
|
|
e.Metadata = map[string]any{}
|
|
}
|
|
e.Metadata["via_mcp"] = true
|
|
e.Metadata["mcp_token_id"] = tid
|
|
}
|
|
return r.inner.Record(ctx, e)
|
|
}
|