You've already forked agentic-coding-workflow
feat(mcp): audit recorder wrapper injecting via_mcp + mcp_token_id from ctx
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,45 @@
|
||||
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)
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/yan1h/agent-coding-workflow/internal/audit"
|
||||
)
|
||||
|
||||
type captureRec struct {
|
||||
mu sync.Mutex
|
||||
entries []audit.Entry
|
||||
}
|
||||
|
||||
func (r *captureRec) Record(_ context.Context, e audit.Entry) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.entries = append(r.entries, e)
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestAuditWrap_InjectsTokenID(t *testing.T) {
|
||||
t.Parallel()
|
||||
cap := &captureRec{}
|
||||
wrapped := AuditWrap(cap)
|
||||
|
||||
uid := uuid.New()
|
||||
ctx := withMcpTokenID(context.Background(), "tok-123")
|
||||
require.NoError(t, wrapped.Record(ctx, audit.Entry{
|
||||
Action: "issue.create",
|
||||
UserID: &uid,
|
||||
Metadata: map[string]any{"slug": "demo"},
|
||||
}))
|
||||
|
||||
require.Len(t, cap.entries, 1)
|
||||
got := cap.entries[0]
|
||||
assert.Equal(t, true, got.Metadata["via_mcp"])
|
||||
assert.Equal(t, "tok-123", got.Metadata["mcp_token_id"])
|
||||
assert.Equal(t, "demo", got.Metadata["slug"])
|
||||
}
|
||||
|
||||
func TestAuditWrap_NoTokenID_Passthrough(t *testing.T) {
|
||||
t.Parallel()
|
||||
cap := &captureRec{}
|
||||
wrapped := AuditWrap(cap)
|
||||
|
||||
require.NoError(t, wrapped.Record(context.Background(), audit.Entry{
|
||||
Action: "user.login",
|
||||
}))
|
||||
|
||||
require.Len(t, cap.entries, 1)
|
||||
got := cap.entries[0]
|
||||
_, hasViaMCP := got.Metadata["via_mcp"]
|
||||
assert.False(t, hasViaMCP)
|
||||
}
|
||||
|
||||
func TestAuditWrap_NilInner_ReturnsNil(t *testing.T) {
|
||||
t.Parallel()
|
||||
assert.Nil(t, AuditWrap(nil))
|
||||
}
|
||||
|
||||
func TestAuditWrap_NilMetadata_Initialized(t *testing.T) {
|
||||
t.Parallel()
|
||||
cap := &captureRec{}
|
||||
wrapped := AuditWrap(cap)
|
||||
|
||||
ctx := withMcpTokenID(context.Background(), "tok-x")
|
||||
require.NoError(t, wrapped.Record(ctx, audit.Entry{Action: "x"}))
|
||||
|
||||
require.Len(t, cap.entries, 1)
|
||||
require.NotNil(t, cap.entries[0].Metadata)
|
||||
assert.Equal(t, "tok-x", cap.entries[0].Metadata["mcp_token_id"])
|
||||
}
|
||||
Reference in New Issue
Block a user