@posthog/mcp 0.10.6 → 0.10.8

package/src/extensions/conversation-id.ts

@@ -11,7 +11,6 @@ import {
 } from './analytics-parameters'
 import { DEFAULT_CONVERSATION_ID_DESCRIPTION } from './constants'
 import { log, type LoggerFn } from './logger'
-import { GET_MORE_TOOLS_NAME } from './tools'
 
 export const CONVERSATION_ID_PARAM_NAME = 'conversation_id'
 
@@ -70,60 +69,83 @@ export function addConversationIdToTool<TTool extends ConversationIdInjectableTo
   return modifiedTool
 }
 
+/**
+ * Injects `conversation_id` across a tool listing, including the virtual
+ * `get_more_tools` tool. Its calls publish `$mcp_missing_capability`, and a
+ * capability gap is only meaningful next to the work that hit it — so it belongs
+ * in the same session as the surrounding tool calls.
+ */
 export function addConversationIdToTools<TTool extends ConversationIdInjectableTool>(
   tools: TTool[],
-  missingCapabilityToolName: string = GET_MORE_TOOLS_NAME,
-  skipMissingCapabilityTool = true,
   logger: LoggerFn = log
 ): TTool[] {
-  return tools.map((tool) => {
-    if (skipMissingCapabilityTool && tool.name === missingCapabilityToolName) {
-      return tool
-    }
-    return addConversationIdToTool(tool, logger)
-  })
+  return tools.map((tool) => addConversationIdToTool(tool, logger))
 }
 
+/**
+ * The shape of every id we mint: a uuidv7. Used to tell an echo of our own handle
+ * from a value the agent made up.
+ */
+const MINTED_CONVERSATION_ID = /^[0-9a-f]{8}-[0-9a-f]{4}-7[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$/i
+
 export type ConversationIdResolution =
   | { minted: false; conversationId: string | undefined }
   | { minted: true; conversationId: string }
 
 /**
  * Decides which conversation_id to use for a tool call:
- *   - disabled or get_more_tools → none
- *   - agent supplied a value → use it
- *   - agent omitted → mint a UUID
+ *   - disabled → none
+ *   - agent echoed a handle we could have minted → use it
+ *   - anything else → mint a fresh one
+ *
+ * The shape check matters because this value becomes `$session_id`, and the
+ * derivation is deterministic so that two pods agree — which also means two
+ * *callers* sending the same string land in the same session. The strings agents
+ * invent are not random (`conv-1`, `1`, `session`), so trusting them verbatim
+ * would silently merge unrelated conversations, potentially across users.
+ *
+ * A compliant agent echoes the uuidv7 we minted and is unaffected. Anything else
+ * is treated exactly as if the handle were absent: mint, and prompt back. Shape
+ * rather than a registry of issued ids, because a per-request server has no
+ * memory of what it minted.
+ *
+ * This narrows the problem, it does not prove provenance: a caller supplying a
+ * well-formed uuidv7 is trusted, so two that pick the *same* one still merge. That
+ * residue is far smaller than the case it replaces — an agent ignoring the
+ * parameter description reaches for `conv-1` or `1`, not a conforming uuidv7 — and
+ * closing it needs a signed handle, i.e. a secret shared by every pod. Worth doing
+ * if this ever anchors something security-bearing; `$session_id` is an analytics
+ * grouping key, so it does not today.
  */
-export function resolveConversationId(
-  enabled: boolean,
-  args: unknown,
-  toolName: string | undefined,
-  missingCapabilityToolName: string = GET_MORE_TOOLS_NAME
-): ConversationIdResolution {
-  if (!enabled || toolName === missingCapabilityToolName) {
+export function resolveConversationId(enabled: boolean, args: unknown): ConversationIdResolution {
+  if (!enabled) {
     return { minted: false, conversationId: undefined }
   }
   const supplied = extractConversationId(args)
-  if (supplied) {
-    return { minted: false, conversationId: supplied }
+  if (supplied && MINTED_CONVERSATION_ID.test(supplied)) {
+    // Lowercased because the shape test is case-insensitive but the hash behind
+    // `$session_id` is not. Some hosts normalise uuids to uppercase, and an
+    // uppercased echo of our own handle would otherwise clear the gate and then
+    // land in a different session than the call that minted it.
+    return { minted: false, conversationId: supplied.toLowerCase() }
   }
   return { minted: true, conversationId: uuidv7() }
 }
 
 /**
- * Predicate matching the same eligibility checks injectConversationIdPromptBack uses,
- * exposed so callers can pre-decide whether the prompt-back will actually land
- * (and clear event.conversationId if not, to avoid orphan ids in analytics).
+ * Whether the prompt-back can ride this result's `content`. The only requirement
+ * is an array to append to.
+ *
+ * Errored results included on purpose. A tool that fails on the first call of a
+ * conversation is exactly when the agent needs the session handle: without it the
+ * retry starts a fresh conversation, so the failure and its fix land in different
+ * sessions.
  */
 export function canInjectConversationIdPromptBack(result: unknown): boolean {
   if (!(result && typeof result === 'object')) {
     return false
   }
-  const resultObj = result as { content?: unknown; isError?: unknown }
-  if (resultObj.isError === true) {
-    return false
-  }
-  return Array.isArray(resultObj.content)
+  return Array.isArray((result as { content?: unknown }).content)
 }
 
 export function extractConversationId(args: unknown): string | undefined {

package/src/extensions/instrument-highlevel.ts

@@ -262,7 +262,9 @@ async function handleToolCallRequest(
     request,
     extra,
     execute: (downstreamRequest) => originalCallToolHandler(downstreamRequest as MCPRequest, extra),
-    parameterOwnership: registeredTool ? getAnalyticsParameterOwnership(registeredTool.inputSchema) : undefined,
+    parameterOwnership: registeredTool
+      ? getAnalyticsParameterOwnership(registeredTool.inputSchema, registeredTool.outputSchema)
+      : undefined,
     takeCapturedError: () => {
       const captured = extra?.__mcp_analytics_error
       if (extra) {

package/src/extensions/instrumentation.ts

@@ -12,6 +12,7 @@ import type {
   MCPServerLike,
   McpEvent,
   ServerClientInfoLike,
+  SessionInfo,
 } from '../types'
 import { getAnalyticsParameterOwnership, stripOwnedAnalyticsArguments } from './analytics-parameters'
 import { addContextParameterToTools, getContextDescription, isContextEnabled } from './context-parameters'
@@ -23,15 +24,16 @@ import {
   resolveConversationId,
 } from './conversation-id'
 import { stampMetaClientInfo } from './client-identity'
+import { addInstructionsToOutputSchemas, mirrorInstructionsIntoStructuredContent } from './output-instructions'
 import { captureEvent } from './capture'
 import { MCPAnalyticsEventType } from './event-types'
 import { captureException } from './exceptions'
 import { resolveToolCallIntent, setEventIntent, setExplicitContextIntent } from './intent'
-import { getServerTrackingData, handleIdentify, setServerTrackingData } from './internal'
+import { getServerTrackingData, handleIdentify, setServerTrackingData, withIdentity } from './internal'
 import type { LoggerFn } from './logger'
 import { buildCapturedMcpParameters } from './mcp-payloads'
 import { getLiteralValue, getObjectShape } from './mcp-sdk-compat'
-import { getSessionId, newSessionId } from './session'
+import { getSessionId, getSessionInfo, newSessionId } from './session'
 import { encodeSessionId, readMcpSessionHeader, writeSessionIdToTransport } from './session-token'
 import { getReportMissingToolDescriptor, resolveMissingCapabilityToolName } from './tools'
 import { applyResolvedMetadata, isToolResultError } from './tracing-helpers'
@@ -105,20 +107,13 @@ export async function captureToolCall(params: TraceToolCallParams): Promise<unkn
     parameterOwnership,
     resolvedEventType === MCPAnalyticsEventType.mcpMissingCapability
   )
-  const conversation = resolveConversationId(
-    ownership.conversationId,
-    request.params?.arguments,
-    request.params?.name,
-    resolvedEventType === MCPAnalyticsEventType.mcpMissingCapability
-      ? resolveMissingCapabilityToolName(data.options)
-      : ''
-  )
+  const conversation = resolveConversationId(ownership.conversationId, request.params?.arguments)
   const downstreamRequest = cloneRequestWithoutOwnedAnalyticsArguments(request, ownership)
 
   // Prepare the event in isolation: if identity/metadata/intent resolution
   // throws, we drop instrumentation for this call but still run the tool.
   const startTime = new Date()
-  const event = await prepareToolCallEvent(
+  const preparedEvent = await prepareToolCallEvent(
     server,
     data,
     request,
@@ -129,34 +124,56 @@ export async function captureToolCall(params: TraceToolCallParams): Promise<unkn
     ownership,
     resolvedEventType
   )
-  if (event && explicitContextIntent) {
-    setExplicitContextIntent(event, explicitContextIntent)
+  if (preparedEvent && explicitContextIntent) {
+    setExplicitContextIntent(preparedEvent.event, explicitContextIntent)
   }
 
   let result: unknown
   try {
     result = await execute(downstreamRequest)
   } catch (error) {
-    publishFailedToolEvent(server, event, error, startTime, conversation, data.logger)
+    publishFailedToolEvent(server, preparedEvent, error, startTime, conversation, data.logger)
     throw error
   }
 
-  const finalResult = applyConversationPromptBack(event, result, conversation)
-  publishSuccessfulToolEvent(server, event, finalResult, startTime, data.logger, takeCapturedError)
+  const finalResult = applyConversationInstructions(preparedEvent?.event ?? null, result, conversation, ownership)
+  publishSuccessfulToolEvent(server, preparedEvent, finalResult, startTime, data.logger, takeCapturedError)
   return finalResult
 }
 
+interface PreparedToolEvent {
+  event: McpEvent
+  requestAttribution: SessionInfo
+}
+
 function getActiveAnalyticsParameterOwnership(
   data: MCPAnalyticsData,
   toolName: string | undefined,
   override: AnalyticsParameterOwnership | undefined,
   isMissingCapabilityTool: boolean
 ): AnalyticsParameterOwnership {
-  const ownership = override ?? (toolName ? data.toolAnalyticsParameterOwnership.get(toolName) : undefined)
+  const listed = toolName ? data.toolAnalyticsParameterOwnership.get(toolName) : undefined
+  const ownership = override ?? listed
   return {
     context: !isMissingCapabilityTool && isContextEnabled(data.options.context) && ownership?.context === true,
-    conversationId:
-      !isMissingCapabilityTool && data.options.enableConversationId === true && ownership?.conversationId === true,
+    conversationId: data.options.enableConversationId === true && ownership?.conversationId === true,
+    // Deliberately read off `listed`, never the override. This asks whether
+    // `tools/list` actually declared `_mcp_instructions`, and only the advertised
+    // JSON Schema can answer it — an override is built from the live registry,
+    // which on the high-level path holds Zod.
+    //
+    // The cache is per-instance, so this is not merely "before the first
+    // `tools/list`": an instance that never serves a listing never writes the
+    // mirror at all. That is the per-request server pattern — `tools/list` lands
+    // on one instance, `tools/call` on a cold one — where the handle falls back
+    // to the `content` block and a structuredContent-only client misses it.
+    //
+    // Failing closed is deliberate. Writing a key the advertised schema did not
+    // declare fails the *entire* tool result under `additionalProperties: false`,
+    // so guessing costs the caller their result, while not guessing costs us one
+    // delivery channel. The fix is a process-scoped ownership cache, so a listing
+    // served by any instance answers for the rest — not a per-call guess.
+    outputInstructions: data.options.enableConversationId === true && listed?.outputInstructions === true,
   }
 }
 
@@ -188,10 +205,22 @@ async function prepareToolCallEvent(
   conversation: ConversationIdResolution,
   ownership: AnalyticsParameterOwnership,
   eventType: MCPAnalyticsEventType
-): Promise<McpEvent | null> {
+): Promise<PreparedToolEvent | null> {
   try {
-    const sessionId = getSessionId(server, extra)
-    await handleIdentify(server, data, sessionId, request, extra)
+    const sessionId = getSessionId(server, extra, conversation.conversationId)
+    // Snapshot token/client/protocol metadata synchronously, before identify or
+    // metadata callbacks can yield and let another request replace shared state.
+    const sessionInfo = getSessionInfo(server, data, sessionId)
+    const identity = await handleIdentify(
+      server,
+      data,
+      sessionId,
+      request,
+      sessionInfo,
+      extra,
+      !!conversation.conversationId
+    )
+    const requestAttribution = withIdentity(sessionInfo, identity)
 
     const toolName = request.params?.name
     const event: McpEvent = {
@@ -211,7 +240,7 @@ async function prepareToolCallEvent(
 
     await applyResolvedMetadata(event, data, request, extra)
     setEventIntent(event, await resolveToolCallIntent(data, request, ownership.context, extra))
-    return event
+    return { event, requestAttribution }
   } catch (error) {
     data.logger(
       `Warning: PostHog MCP analytics instrumentation failed for tool ${request.params?.name}, the tool will still run - ${error}`
@@ -221,38 +250,65 @@ async function prepareToolCallEvent(
 }
 
 /**
- * When we minted a conversation id, append the prompt-back so the agent echoes
- * it on subsequent calls. If the result can't carry it, clear the id off the
- * event so analytics doesn't show an orphan the agent never received.
+ * Delivers the conversation session handle back to the agent, over both channels a tool
+ * result has:
+ *
+ * - `structuredContent`, on *every* response for a tool whose output schema we
+ *   declared the key on. Clients that read structured results never see the text
+ *   block, and re-sending it each time lets an agent that dropped the session handle read
+ *   it back.
+ * - `content`, as a text block, only on the response that minted the session handle.
+ *   Repeating it every call would put a `[SERVER]:` line in front of the user on
+ *   every single tool result.
+ *
+ * If neither channel could carry a session handle we minted, the agent never received it,
+ * so clear it off the event rather than showing analytics an id nobody has.
  */
-function applyConversationPromptBack(
+function applyConversationInstructions(
   event: McpEvent | null,
   result: unknown,
-  conversation: ConversationIdResolution
+  conversation: ConversationIdResolution,
+  ownership: AnalyticsParameterOwnership
 ): unknown {
-  if (!conversation.minted) {
+  const conversationId = conversation.conversationId
+  if (!conversationId) {
     return result
   }
-  if (canInjectConversationIdPromptBack(result)) {
-    return injectConversationIdPromptBack(result, conversation.conversationId)
+
+  let updated = result
+  let delivered = false
+
+  if (ownership.outputInstructions) {
+    const mirrored = mirrorInstructionsIntoStructuredContent(updated, conversationId)
+    delivered = mirrored !== updated
+    updated = mirrored
+  }
+
+  if (conversation.minted && canInjectConversationIdPromptBack(updated)) {
+    updated = injectConversationIdPromptBack(updated, conversationId)
+    delivered = true
   }
-  if (event) {
+
+  // Only a minted session handle can be lost this way — one the agent supplied, it has.
+  if (!delivered && conversation.minted && event) {
     event.conversationId = undefined
   }
-  return result
+
+  return updated
 }
 
 function publishSuccessfulToolEvent(
   server: MCPServerLike,
-  event: McpEvent | null,
+  preparedEvent: PreparedToolEvent | null,
   result: unknown,
   startTime: Date,
   logger: LoggerFn,
   takeCapturedError?: () => unknown
 ): void {
-  if (!event) {
+  if (!preparedEvent) {
     return
   }
+  const { event, requestAttribution } = preparedEvent
   try {
     if (isToolResultError(result)) {
       event.isError = true
@@ -263,7 +319,7 @@ function publishSuccessfulToolEvent(
     }
     event.response = result
     event.duration = Date.now() - startTime.getTime()
-    captureEvent(server, event, logger)
+    captureEvent(server, event, logger, requestAttribution)
   } catch (error) {
     logger(`Warning: PostHog MCP analytics failed to publish tool event - ${error}`)
   }
@@ -271,15 +327,16 @@ function publishSuccessfulToolEvent(
 
 function publishFailedToolEvent(
   server: MCPServerLike,
-  event: McpEvent | null,
+  preparedEvent: PreparedToolEvent | null,
   error: unknown,
   startTime: Date,
   conversation: ConversationIdResolution,
   logger: LoggerFn
 ): void {
-  if (!event) {
+  if (!preparedEvent) {
     return
   }
+  const { event, requestAttribution } = preparedEvent
   try {
     if (conversation.minted) {
       event.conversationId = undefined
@@ -287,7 +344,7 @@ function publishFailedToolEvent(
     event.isError = true
     event.error = captureException(error)
     event.duration = Date.now() - startTime.getTime()
-    captureEvent(server, event, logger)
+    captureEvent(server, event, logger, requestAttribution)
   } catch (publishError) {
     logger(`Warning: PostHog MCP analytics failed to publish failed tool event - ${publishError}`)
   }
@@ -407,8 +464,12 @@ export async function handleListToolsRequest(
 ): Promise<{ tools: ListToolsResult['tools'] }> {
   const data = getServerTrackingData(server)
   const startTime = new Date()
+  const sessionId = getSessionId(server, extra)
+  // Snapshot before metadata resolution or the list handler can yield to a
+  // concurrent request using the same instrumented server.
+  const requestAttribution = getSessionInfo(server, data, sessionId)
   const event: McpEvent = {
-    sessionId: getSessionId(server, extra),
+    sessionId,
     parameters: buildCapturedMcpParameters(request),
     eventType: MCPAnalyticsEventType.mcpToolsList,
     timestamp: startTime,
@@ -419,7 +480,15 @@ export async function handleListToolsRequest(
     await applyResolvedMetadata(event, data, request, extra)
   }
 
-  const tools = await getTracedToolsList(server, originalListToolsHandler, request, extra, event, logger)
+  const tools = await getTracedToolsList(
+    server,
+    originalListToolsHandler,
+    request,
+    extra,
+    event,
+    logger,
+    requestAttribution
+  )
 
   if (!data) {
     logger(
@@ -435,7 +504,7 @@ export async function handleListToolsRequest(
     event.error = captureException('No tools were sent to MCP client.')
     event.isError = true
     event.duration = Date.now() - startTime.getTime()
-    captureEvent(server, event, data.logger)
+    captureEvent(server, event, data.logger, requestAttribution)
     return { tools }
   }
 
@@ -443,7 +512,7 @@ export async function handleListToolsRequest(
   event.listedToolNames = collectListedToolNames(tools)
   event.isError = false
   event.duration = Date.now() - startTime.getTime()
-  captureEvent(server, event, data.logger)
+  captureEvent(server, event, data.logger, requestAttribution)
   return { tools }
 }
 
@@ -454,7 +523,7 @@ function cacheToolAnalyticsParameterOwnership(
   // Merge pages and concurrent enumerations; repeated tool names overwrite stale schemas.
   for (const tool of tools) {
     if (tool?.name) {
-      cache.set(tool.name, getAnalyticsParameterOwnership(tool.inputSchema))
+      cache.set(tool.name, getAnalyticsParameterOwnership(tool.inputSchema, tool.outputSchema))
     }
   }
 }
@@ -473,7 +542,8 @@ async function getTracedToolsList(
   request: MCPRequestLike,
   extra: CompatibleRequestHandlerExtra | undefined,
   event: McpEvent,
-  logger: LoggerFn
+  logger: LoggerFn,
+  requestAttribution: SessionInfo
 ): Promise<ListToolsResult['tools']> {
   try {
     const data = getServerTrackingData(server)
@@ -490,7 +560,6 @@ async function getTracedToolsList(
 
     if (data) {
       const missingToolName = resolveMissingCapabilityToolName(data.options)
-      let injectedMissingCapabilityTool = false
       if (data.options.reportMissing) {
         const alreadyPresent = tools.some((tool) => tool?.name === missingToolName)
         if (alreadyPresent) {
@@ -498,13 +567,17 @@ async function getTracedToolsList(
             `Warning: Cannot inject missing-capability tool "${missingToolName}" because a real tool already uses that name. The real tool will not be intercepted.`
           )
         } else {
-          tools.push(getReportMissingToolDescriptor(missingToolName))
-          injectedMissingCapabilityTool = true
+          const virtualTool = getReportMissingToolDescriptor(missingToolName)
+          tools.push(virtualTool)
+          // Cached separately because the virtual tool is added after the listing
+          // was cached, and its calls need ownership like any other tool's.
+          cacheToolAnalyticsParameterOwnership(data.toolAnalyticsParameterOwnership, [virtualTool])
         }
       }
 
       if (data.options.enableConversationId) {
-        tools = addConversationIdToTools(tools, missingToolName, injectedMissingCapabilityTool, data.logger)
+        tools = addConversationIdToTools(tools, data.logger)
+        tools = addInstructionsToOutputSchemas(tools, data.logger)
       }
     }
 
@@ -521,7 +594,7 @@ async function getTracedToolsList(
     event.error = captureException(error)
     event.isError = true
     event.duration = event.timestamp ? Date.now() - event.timestamp.getTime() : 0
-    captureEvent(server, event, logger)
+    captureEvent(server, event, logger, requestAttribution)
     throw error
   }
 }
@@ -630,9 +703,10 @@ function mintStatelessSessionOnInitialize(
  */
 function upgradeMintedTokenToNegotiated(
   server: MCPServerLike,
-  data: MCPAnalyticsData,
   mintedSessionId: string,
-  negotiatedProtocolVersion: string | undefined
+  sessionInfo: SessionInfo,
+  negotiatedProtocolVersion: string | undefined,
+  logger: LoggerFn
 ): void {
   try {
     const transport = server.transport
@@ -641,13 +715,13 @@ function upgradeMintedTokenToNegotiated(
     }
     const token = encodeSessionId({
       sessionId: mintedSessionId,
-      clientName: data.sessionInfo.clientName,
-      clientVersion: data.sessionInfo.clientVersion,
+      clientName: sessionInfo.clientName,
+      clientVersion: sessionInfo.clientVersion,
       protocolVersion: negotiatedProtocolVersion,
     })
     writeSessionIdToTransport(transport, token)
   } catch (error) {
-    data.logger(`Warning: PostHog MCP analytics failed to upgrade the stateless session token - ${error}`)
+    logger(`Warning: PostHog MCP analytics failed to upgrade the stateless session token - ${error}`)
   }
 }
 
@@ -689,7 +763,17 @@ export async function handleInitializeRequest(
   // Mint first so the `$mcp_initialize` event below already carries the minted id.
   const mintedSessionId = mintStatelessSessionOnInitialize(server, data, request, extra)
   const sessionId = getSessionId(server, extra)
-  await handleIdentify(server, data, sessionId, request, extra)
+  // Snapshot before identify, metadata, or the initialize handler can yield to
+  // another request using the same instrumented server.
+  const sessionInfo = getSessionInfo(server, data, sessionId)
+  const initializeClientInfo = readInitializeClientInfo(request)
+  const requestSessionInfo: SessionInfo = {
+    ...sessionInfo,
+    clientName: initializeClientInfo?.name ?? sessionInfo.clientName,
+    clientVersion: initializeClientInfo?.version ?? sessionInfo.clientVersion,
+  }
+  const identity = await handleIdentify(server, data, sessionId, request, requestSessionInfo, extra)
+  const requestAttribution = withIdentity(requestSessionInfo, identity)
 
   const event: McpEvent = {
     sessionId,
@@ -712,12 +796,19 @@ export async function handleInitializeRequest(
   // re-mint the token so pods replaying it report the negotiated version too.
   const negotiatedProtocolVersion = readProtocolVersion(result, request)
   event.protocolVersion = negotiatedProtocolVersion
-  data.sessionInfo.protocolVersion = negotiatedProtocolVersion
-  setServerTrackingData(server, data)
+  // Do not let a delayed initialize overwrite whichever session became current
+  // while its callbacks or the original handler were awaiting.
+  if (data.sessionId === sessionId) {
+    data.sessionInfo = { ...data.sessionInfo, protocolVersion: negotiatedProtocolVersion }
+    setServerTrackingData(server, data)
+  }
   if (mintedSessionId) {
-    upgradeMintedTokenToNegotiated(server, data, mintedSessionId, negotiatedProtocolVersion)
+    upgradeMintedTokenToNegotiated(server, mintedSessionId, requestSessionInfo, negotiatedProtocolVersion, data.logger)
   }
-  captureEvent(server, event, data.logger)
+  captureEvent(server, event, data.logger, {
+    ...requestAttribution,
+    protocolVersion: negotiatedProtocolVersion,
+  })
   return result
 }
 

package/src/extensions/internal.ts

@@ -9,6 +9,7 @@ import type {
   MCPRequestLike,
   MCPServerLike,
   McpEvent,
+  SessionInfo,
   UserIdentity,
 } from '../types'
 import { MCPAnalyticsEventType } from './event-types'
@@ -128,10 +129,20 @@ export async function handleIdentify(
   data: MCPAnalyticsData,
   sessionId: string,
   request: MCPRequestLike,
-  extra?: CompatibleRequestHandlerExtra
-): Promise<void> {
+  requestAttribution: SessionInfo,
+  extra?: CompatibleRequestHandlerExtra,
+  /**
+   * True when this session id came from an agent-carried `conversation_id`
+   * rather than the transport. Such a session is brand new and was never
+   * announced at `initialize`, whatever `data.sessionSource` still says about
+   * the connection.
+   */
+  sessionFromConversation = false
+): Promise<UserIdentity | undefined> {
+  const identityBeforeRequest = data.identifiedSessions.get(sessionId)
+  const sessionSourceBeforeIdentify = data.sessionSource
   if (!data.options.identify) {
-    return
+    return identityBeforeRequest
   }
 
   const identifyEvent: McpEvent = {
@@ -162,23 +173,35 @@ export async function handleIdentify(
       // announced); revisit with the stateless-by-default rework.
       const changed = previousIdentity !== undefined && !areIdentitiesEqual(previousIdentity, mergedIdentity)
       const firstSeen = previousIdentity === undefined
-      const announcedAtInitialize = data.sessionSource === 'token' && request.method !== 'initialize'
+      const announcedAtInitialize =
+        !sessionFromConversation && sessionSourceBeforeIdentify === 'token' && request.method !== 'initialize'
       const shouldPublish = changed || (firstSeen && !announcedAtInitialize)
 
       data.identifiedSessions.set(sessionId, mergedIdentity)
 
       if (shouldPublish) {
         data.logger(`Identified session ${sessionId}`)
-        captureEvent(server, identifyEvent, data.logger)
+        captureEvent(server, identifyEvent, data.logger, withIdentity(requestAttribution, mergedIdentity))
       }
-    } else {
-      data.logger(`Warning: Supplied identify function returned null for session ${sessionId}`)
+      return mergedIdentity
     }
+
+    data.logger(`Warning: Supplied identify function returned null for session ${sessionId}`)
   } catch (error) {
     data.logger(
       `Error: User supplied identify function threw an error while identifying session ${sessionId} - ${error}`
     )
   }
+  return identityBeforeRequest
+}
+
+export function withIdentity(sessionInfo: SessionInfo, identity: UserIdentity | undefined): SessionInfo {
+  return {
+    ...sessionInfo,
+    identifyActorGivenId: identity?.distinctId,
+    identifyActorData: identity?.properties || {},
+    identifyActorGroups: identity?.groups,
+  }
 }
 
 /**