feat(mcp): audit recorder wrapper injecting via_mcp + mcp_token_id from ctx

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-05-08 10:35:44 +08:00
parent d6279c6ee1
commit 5fd4eab7ca
2 changed files with 123 additions and 0 deletions
@@ -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"])
+}