diff --git a/internal/mcp/audit_ctx.go b/internal/mcp/audit_ctx.go new file mode 100644 index 0000000..efed057 --- /dev/null +++ b/internal/mcp/audit_ctx.go @@ -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) +} diff --git a/internal/mcp/audit_ctx_test.go b/internal/mcp/audit_ctx_test.go new file mode 100644 index 0000000..4ef54b4 --- /dev/null +++ b/internal/mcp/audit_ctx_test.go @@ -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"]) +}