From 787ae29b61db12b0f9e35005abe38b583767891b Mon Sep 17 00:00:00 2001 From: Victor Ma Date: Fri, 10 Jul 2026 10:02:54 -0400 Subject: [PATCH] feat: session-scoped correlation ID on feedback analytics events Mint one correlation ID per studio-mcp process (each IDE connection gets a fresh process) and stamp it on every snyk_send_feedback analytics event via NewAnalyticsEventParam, so events emitted by the same session can be joined downstream instead of each call minting its own random ID. --- internal/analytics/analytics.go | 10 ++- internal/analytics/analytics_test.go | 127 +++++++++++++++++++++++++++ internal/mcp/llm_binding.go | 21 +++++ internal/mcp/llm_binding_test.go | 42 +++++++++ internal/mcp/tools.go | 2 +- internal/mcp/tools_test.go | 91 +++++++++++++++++++ 6 files changed, 291 insertions(+), 2 deletions(-) create mode 100644 internal/analytics/analytics_test.go diff --git a/internal/analytics/analytics.go b/internal/analytics/analytics.go index 2343160..923f15d 100644 --- a/internal/analytics/analytics.go +++ b/internal/analytics/analytics.go @@ -47,7 +47,14 @@ type EventParam struct { InteractionUUID string `json:"interactionId"` } -func NewAnalyticsEventParam(interactionType string, err error, path types.FilePath) EventParam { +// NewAnalyticsEventParam builds a new analytics event. +// correlationID, when non-empty, is stamped on the event as its +// InteractionUUID so that events sharing a correlationID (e.g. a +// session-scoped ID minted once per MCP process) can be joined downstream, +// instead of each call minting its own fresh random ID. +// Pass an empty string to preserve the previous behavior of letting +// PayloadForAnalyticsEventParam mint a fresh UUID for this event. +func NewAnalyticsEventParam(interactionType string, err error, path types.FilePath, correlationID string) EventParam { status := string(analytics.Success) if err != nil { status = string(analytics.Failure) @@ -71,6 +78,7 @@ func NewAnalyticsEventParam(interactionType string, err error, path types.FilePa Status: status, TimestampMs: time.Now().UnixMilli(), TargetId: targetId, + InteractionUUID: correlationID, } } diff --git a/internal/analytics/analytics_test.go b/internal/analytics/analytics_test.go new file mode 100644 index 0000000..7566deb --- /dev/null +++ b/internal/analytics/analytics_test.go @@ -0,0 +1,127 @@ +/* + * © 2025 Snyk Limited + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package analytics + +import ( + "encoding/json" + "testing" + + "github.com/golang/mock/gomock" + "github.com/snyk/go-application-framework/pkg/analytics" + "github.com/snyk/go-application-framework/pkg/configuration" + "github.com/snyk/go-application-framework/pkg/mocks" + "github.com/snyk/go-application-framework/pkg/runtimeinfo" + "github.com/snyk/studio-mcp/internal/types" + "github.com/stretchr/testify/require" +) + +// These tests cover the "Session-scoped correlation ID on feedback events" +// requirement (specs/snyk-fix-feedback-verification): studio-mcp mints one +// correlation ID per process and stamps it on every snyk_send_feedback +// analytics event, replacing the previous behavior of minting a fresh random +// ID on every call. + +func TestNewAnalyticsEventParam_CorrelationID(t *testing.T) { + t.Run("stamps the provided correlation ID as InteractionUUID", func(t *testing.T) { + event := NewAnalyticsEventParam("Send feedback", nil, types.FilePath("/tmp"), "session-correlation-id") + require.Equal(t, "session-correlation-id", event.InteractionUUID) + }) + + t.Run("two calls with the same correlation ID (same session) share it", func(t *testing.T) { + correlationID := "shared-session-id" + first := NewAnalyticsEventParam("Send feedback", nil, types.FilePath("/tmp/a"), correlationID) + second := NewAnalyticsEventParam("Send feedback", nil, types.FilePath("/tmp/b"), correlationID) + + require.Equal(t, correlationID, first.InteractionUUID) + require.Equal(t, first.InteractionUUID, second.InteractionUUID) + }) + + t.Run("different correlation IDs (different sessions/processes) differ", func(t *testing.T) { + first := NewAnalyticsEventParam("Send feedback", nil, types.FilePath("/tmp"), "session-a") + second := NewAnalyticsEventParam("Send feedback", nil, types.FilePath("/tmp"), "session-b") + + require.NotEqual(t, first.InteractionUUID, second.InteractionUUID) + }) + + t.Run("empty correlation ID leaves InteractionUUID unset", func(t *testing.T) { + event := NewAnalyticsEventParam("Send feedback", nil, types.FilePath("/tmp"), "") + require.Empty(t, event.InteractionUUID) + }) +} + +// setupMockEngine builds a minimal mock workflow.Engine sufficient for +// PayloadForAnalyticsEventParam, which only reads GetConfiguration()/GetRuntimeInfo(). +func setupMockEngine(t *testing.T) *mocks.MockEngine { + t.Helper() + ctrl := gomock.NewController(t) + engine := mocks.NewMockEngine(ctrl) + engineConfig := configuration.NewWithOpts(configuration.WithAutomaticEnv()) + engine.EXPECT().GetConfiguration().Return(engineConfig).AnyTimes() + engine.EXPECT().GetRuntimeInfo().Return(runtimeinfo.New(runtimeinfo.WithName("test"), runtimeinfo.WithVersion("1.0.0"))).AnyTimes() + return engine +} + +// interactionURN renders the instrumentation collector's payload to JSON and +// returns it as a string, so tests can assert on the embedded +// "urn:snyk:interaction:" value without reaching into go-application-framework's +// internal API types. +func interactionURNPayload(t *testing.T, engine *mocks.MockEngine, event EventParam) string { + t.Helper() + ic := PayloadForAnalyticsEventParam(engine, "", event) + obj, err := analytics.GetV2InstrumentationObject(ic) + require.NoError(t, err) + b, err := json.Marshal(obj) + require.NoError(t, err) + return string(b) +} + +func TestPayloadForAnalyticsEventParam_CorrelationID(t *testing.T) { + t.Run("preserves a pre-set session correlation ID instead of minting a fresh one", func(t *testing.T) { + engine := setupMockEngine(t) + correlationID := "session-correlation-id" + event := NewAnalyticsEventParam("Send feedback", nil, types.FilePath("/tmp"), correlationID) + + payload := interactionURNPayload(t, engine, event) + + require.Contains(t, payload, "urn:snyk:interaction:"+correlationID) + }) + + t.Run("two events built from the same session correlation ID resolve to the same interaction URN", func(t *testing.T) { + engine := setupMockEngine(t) + correlationID := "shared-session-id" + + eventOne := NewAnalyticsEventParam("Send feedback", nil, types.FilePath("/tmp/a"), correlationID) + eventTwo := NewAnalyticsEventParam("Send feedback", nil, types.FilePath("/tmp/b"), correlationID) + + payloadOne := interactionURNPayload(t, engine, eventOne) + payloadTwo := interactionURNPayload(t, engine, eventTwo) + + expected := "urn:snyk:interaction:" + correlationID + require.Contains(t, payloadOne, expected) + require.Contains(t, payloadTwo, expected) + }) + + t.Run("mints a fresh interaction ID when InteractionUUID is unset", func(t *testing.T) { + engine := setupMockEngine(t) + event := NewAnalyticsEventParam("Send feedback", nil, types.FilePath("/tmp"), "") + require.Empty(t, event.InteractionUUID, "sanity check: no correlation ID was provided") + + payload := interactionURNPayload(t, engine, event) + + require.Contains(t, payload, "urn:snyk:interaction:") + }) +} diff --git a/internal/mcp/llm_binding.go b/internal/mcp/llm_binding.go index 26a4fa4..b2812ef 100644 --- a/internal/mcp/llm_binding.go +++ b/internal/mcp/llm_binding.go @@ -62,6 +62,13 @@ type McpLLMBinding struct { started bool cliPath string openBrowserFunc types.OpenBrowserFunc + + // correlationID is a session-scoped identifier minted once per process in + // Start() and stamped on every snyk_send_feedback analytics event, so + // events from the same MCP session can be correlated. + // It is set once, before the server starts accepting tool calls, and is + // thereafter only read, so it needs no dedicated lock. + correlationID string } func NewMcpLLMBinding(opts ...Option) *McpLLMBinding { @@ -78,8 +85,22 @@ func NewMcpLLMBinding(opts ...Option) *McpLLMBinding { return mcpServerImpl } +// mintCorrelationID mints the session-scoped correlation ID once per process, +// if it hasn't been minted already. +// Every shipped MCP config launches studio-mcp as a fresh process per IDE +// connection, so minting once here (before any tool call can occur) is +// sufficient to correlate every snyk_send_feedback event emitted during this +// session's lifetime. +func (m *McpLLMBinding) mintCorrelationID() { + if m.correlationID == "" { + m.correlationID = uuid.New().String() + } +} + // Start starts the MCP server. It blocks until the server is stopped via Shutdown. func (m *McpLLMBinding) Start(invocationContext workflow.InvocationContext) error { + m.mintCorrelationID() + runTimeInfo := invocationContext.GetRuntimeInfo() version := "" if runTimeInfo != nil { diff --git a/internal/mcp/llm_binding_test.go b/internal/mcp/llm_binding_test.go index 6105aad..c256c48 100644 --- a/internal/mcp/llm_binding_test.go +++ b/internal/mcp/llm_binding_test.go @@ -402,3 +402,45 @@ func TestStart(t *testing.T) { }) }) } + +// TestMintCorrelationID covers the "Session-scoped correlation ID on feedback +// events" requirement (specs/snyk-fix-feedback-verification): one correlation +// ID is minted per process and reused for the rest of that process's +// lifetime, while a different process (a different McpLLMBinding here) gets +// its own distinct ID. +func TestMintCorrelationID(t *testing.T) { + t.Run("mints a non-empty ID", func(t *testing.T) { + binding := NewMcpLLMBinding() + require.Empty(t, binding.correlationID, "correlationID must be unset before minting") + + binding.mintCorrelationID() + + require.NotEmpty(t, binding.correlationID) + }) + + t.Run("two feedback calls in the same process (same binding) share one correlation ID", func(t *testing.T) { + binding := NewMcpLLMBinding() + + binding.mintCorrelationID() + first := binding.correlationID + + // Simulates a second call within the same session/process reading the + // same session-scoped ID; minting must be idempotent. + binding.mintCorrelationID() + second := binding.correlationID + + require.Equal(t, first, second) + }) + + t.Run("two different sessions (different bindings) get different correlation IDs", func(t *testing.T) { + bindingA := NewMcpLLMBinding() + bindingB := NewMcpLLMBinding() + + bindingA.mintCorrelationID() + bindingB.mintCorrelationID() + + require.NotEmpty(t, bindingA.correlationID) + require.NotEmpty(t, bindingB.correlationID) + require.NotEqual(t, bindingA.correlationID, bindingB.correlationID) + }) +} diff --git a/internal/mcp/tools.go b/internal/mcp/tools.go index cee6dd5..42d056f 100644 --- a/internal/mcp/tools.go +++ b/internal/mcp/tools.go @@ -516,7 +516,7 @@ func (m *McpLLMBinding) snykSendFeedback(invocationCtx workflow.InvocationContex clientInfo := ClientInfoFromContext(ctx) m.updateGafConfigWithIntegrationEnvironment(invocationCtx, clientInfo.Name, clientInfo.Version) - event := analytics.NewAnalyticsEventParam("Send feedback", nil, types.FilePath(path)) + event := analytics.NewAnalyticsEventParam("Send feedback", nil, types.FilePath(path), m.correlationID) event.Extension = buildSendFeedbackExtension(&logger, int(preventedCount), int(remediatedCount), preventedIDs) go analytics.SendAnalytics(invocationCtx.GetEngine(), "", event, nil) diff --git a/internal/mcp/tools_test.go b/internal/mcp/tools_test.go index 2f4d501..9fdf591 100644 --- a/internal/mcp/tools_test.go +++ b/internal/mcp/tools_test.go @@ -29,6 +29,7 @@ import ( "reflect" "runtime" "strings" + "sync" "testing" "github.com/golang/mock/gomock" @@ -125,6 +126,60 @@ func (f *testFixture) mockCliOutput(output string) { createMockSnykCli(f.t, f.snykCliPath, output) } +// analyticsCapture records the raw payload bytes handed to the mock engine's +// InvokeWithInputAndConfig by each background analytics.SendAnalytics call, and +// lets tests block until a known number of dispatches have completed before +// asserting on them. +type analyticsCapture struct { + wg sync.WaitGroup + mu sync.Mutex + payloads [][]byte +} + +func (c *analyticsCapture) add(n int) { + c.wg.Add(n) +} + +func (c *analyticsCapture) wait() { + c.wg.Wait() +} + +func (c *analyticsCapture) all() [][]byte { + c.mu.Lock() + defer c.mu.Unlock() + return append([][]byte{}, c.payloads...) +} + +// allowAnalyticsDispatch wires the fixture's mock engine so that +// snykSendFeedback's `go analytics.SendAnalytics(...)` background call can run +// to completion (GetLogger/GetRuntimeInfo/InvokeWithInputAndConfig are +// otherwise unmocked and would fail the test) and captures the payload each +// dispatch sends. +// Callers must call capture.add(n) before triggering n dispatches, then +// capture.wait() before inspecting capture.all(). +func (f *testFixture) allowAnalyticsDispatch() *analyticsCapture { + f.t.Helper() + capture := &analyticsCapture{} + + nopLogger := zerolog.Nop() + f.mockEngine.EXPECT().GetLogger().Return(&nopLogger).AnyTimes() + f.mockEngine.EXPECT().GetRuntimeInfo().Return(runtimeinfo.New(runtimeinfo.WithName("test"), runtimeinfo.WithVersion("1.0.0"))).AnyTimes() + f.mockEngine.EXPECT().InvokeWithInputAndConfig(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( + func(_ workflow.Identifier, input []workflow.Data, _ configuration.Configuration) ([]workflow.Data, error) { + defer capture.wg.Done() + if len(input) > 0 { + if b, ok := input[0].GetPayload().([]byte); ok { + capture.mu.Lock() + capture.payloads = append(capture.payloads, b) + capture.mu.Unlock() + } + } + return nil, nil + }).AnyTimes() + + return capture +} + func getToolWithName(t *testing.T, tools *SnykMcpTools, toolName string) *SnykMcpToolsDefinition { t.Helper() for _, tool := range tools.Tools { @@ -1799,6 +1854,42 @@ func TestSnykSendFeedbackHandler_Validation(t *testing.T) { }) } +// TestSnykSendFeedbackHandler_CorrelationID covers the "Two feedback calls in +// the same session share a correlation ID" spec scenario end-to-end through +// the actual snykSendFeedback handler, not just the analytics helpers it +// calls into. +func TestSnykSendFeedbackHandler_CorrelationID(t *testing.T) { + fixture := setupTestFixture(t) + toolDef := getToolWithName(t, fixture.tools, ToolName.SendFeedback) + require.NotNil(t, toolDef) + + fixture.binding.mintCorrelationID() + correlationID := fixture.binding.correlationID + require.NotEmpty(t, correlationID) + + capture := fixture.allowAnalyticsDispatch() + handler := fixture.binding.snykSendFeedback(fixture.invocationContext, *toolDef) + + capture.add(2) + for i := 0; i < 2; i++ { + req := mcp.CallToolRequest{Params: mcp.CallToolParams{Arguments: map[string]interface{}{ + "preventedIssuesCount": float64(1), + "fixedExistingIssuesCount": float64(0), + "path": "/tmp", + }}} + result, err := handler(t.Context(), req) + require.NoError(t, err) + require.NotNil(t, result) + } + capture.wait() + + payloads := capture.all() + require.Len(t, payloads, 2) + expectedURN := "urn:snyk:interaction:" + correlationID + require.Contains(t, string(payloads[0]), expectedURN) + require.Contains(t, string(payloads[1]), expectedURN) +} + func TestCoerceStringSlice(t *testing.T) { t.Run("NilInput", func(t *testing.T) { require.Nil(t, coerceStringSlice(nil))