@salesforce/sfdx-agent-harness-mastra 0.6.0 → 0.8.0

package/README.md

@@ -109,6 +109,30 @@ Mastra-internal pool / index / loaded-tools state is intentionally not on the ha
 are `getStateStats()`, `clearState(threadId?)`, `clearAllState()`, and `cleanupNow()` on `ToolSearchHandle`, plus `id`
 on both handles.
 
+### MCP tool-event enrichment
+
+The harness enriches every `tool-call`, `tool-result`, and `tool-approval-request` event with the originating MCP server
+name and any annotations the server declared, so consumers can branch on tool hints or group by server without parsing
+the namespaced tool name themselves. `serverName` is `undefined` for non-MCP tools (workspace tools, consumer-declared
+tools); `annotations` is `undefined` when the server omitted the field (per the MCP spec).
+
+```ts
+for await (const event of result.eventStream) {
+  if (event.type === 'tool-approval-request') {
+    if (event.annotations?.readOnlyHint) {
+      // Auto-approve safe reads.
+      await session.approveToolCall(event.toolCall.toolCallId);
+    } else {
+      // Group/label by server when present, e.g. "sfdx wants to: deploy_metadata".
+      const approved = await promptUser(event.serverName, event.toolCall);
+      approved
+        ? await session.approveToolCall(event.toolCall.toolCallId)
+        : await session.declineToolCall(event.toolCall.toolCallId);
+    }
+  }
+}
+```
+
 ### Using LLM Gateway Express
 
 ```ts

package/dist/mastra-harness.d.ts

@@ -28,9 +28,16 @@ interface ActiveRun {
 /**
  * Per-server MCP state. Each enabled server gets its own MCPClient and
  * independent discovery promise so readiness is tracked individually.
+ *
+ * `epoch` increments on every fresh discovery (initial create + each
+ * reconnect). Discovery callbacks captured against an earlier epoch drop their
+ * writes when a newer discovery has already started for the same server, so an
+ * orphaned in-flight `listToolsetsWithErrors()` cannot overwrite a fresher
+ * discovery's `tools` / `error` / `settled` state.
  */
 interface McpServerState {
     client: MCPClient;
+    epoch: number;
     error?: string;
     ready: Promise<void>;
     settled: boolean;
@@ -165,6 +172,18 @@ export declare class MastraHarness implements MastraAgentHarness {
      */
     getAgentIds(): Promise<string[]>;
     getMcpServerInfo(agentId: string): McpServerInfo[];
+    /**
+     * Eager implementation of {@link AgentHarness.reconnectMcpServer}: calls
+     * Mastra's `MCPClient.reconnectServer(name)` primitive, supersedes any
+     * in-flight discovery (replace-ready semantics — latest call wins), and
+     * awaits the new discovery before resolving.
+     *
+     * Programming errors (unknown agent / unknown server / disabled) reject.
+     * Transport / discovery failures are recorded on the per-server state
+     * (visible via {@link getMcpServerInfo}) and a `mcp-server-discovery-*`
+     * telemetry pair fires; the returned promise still resolves.
+     */
+    reconnectMcpServer(agentId: string, serverName: string): Promise<void>;
     /**
      * @requirements
      * - MUST throw an error if the specified `agentId` is not found in `activeAgents`.
@@ -215,6 +234,14 @@ export declare class MastraHarness implements MastraAgentHarness {
      *   User turns are treated as authoritative input.
      */
     compactThread(agentId: string, threadId: string): Promise<string>;
+    /**
+     * Best-effort wrapper around `destroyThread`: swallows the error so the
+     * caller's outer flow isn't aborted by a teardown blip, but logs the
+     * failure at warn level so a future regression in `destroyThread` is
+     * still observable. Used by `compactThread` on both the empty-source and
+     * non-empty paths.
+     */
+    private destroyThreadBestEffort;
     /**
      * @requirements
      * - MUST throw an error if the specified `agentId` is not found in `activeAgents`.
@@ -302,7 +329,6 @@ export declare class MastraHarness implements MastraAgentHarness {
      * directly.
      */
     private summarize;
-    private buildSummaryPrompt;
     /**
      * Two-pass summarization for transcripts above the char budget. The
      * transcript is split into roughly-equal chunks each ≤ MAX_TRANSCRIPT_CHARS,
@@ -317,9 +343,41 @@ export declare class MastraHarness implements MastraAgentHarness {
      */
     private splitIntoChunks;
     private getAgentStateOrThrow;
+    /**
+     * Runs MCP discovery for one server: emits the
+     * `mcp-server-discovery-started` / `…-completed` / `…-failed` telemetry
+     * triple and writes results into the corresponding entry in `serverStates`.
+     *
+     * The caller MUST `set()` the entry into `serverStates` either before
+     * calling this helper (the `createAgent` path constructs the entry
+     * synchronously after this returns the promise) or before any await
+     * boundary in the helper's chain — the chain reads back via
+     * `serverStates.get(serverName)!` to write `tools` / `error`. Both call
+     * sites today put the entry in place before the first microtask boundary.
+     *
+     * Errors observed during discovery are recorded on the entry's `error`
+     * field; the returned promise resolves either way (it only rejects when
+     * an aborted operation throws after we have already cleaned up). The
+     * `settled` flag is pinned by `.finally()` so `getMcpServerInfo()` can
+     * tell `Connecting` from `Connected`/`Error`.
+     */
+    private runMcpDiscovery;
     private toChatStreamResult;
     private getMemoryOrThrow;
     private agentResourceId;
+    /**
+     * Build a `RequestContext` carrying the thread + resource ids under the
+     * reserved keys Mastra reads when input processors (e.g. `ToolSearchProcessor`,
+     * `SkillSearchProcessor`) call `requestContext.get(MASTRA_THREAD_ID_KEY)`.
+     *
+     * `memory: { thread, resource }` populates Mastra's memory layer but is NOT
+     * propagated into the request context, so processors that key per-thread
+     * state by `MASTRA_THREAD_ID_KEY` would otherwise see `undefined` and fall
+     * back to the `'default'` bucket — folding every consumer thread into one.
+     * Surfacing the ids here lets `ToolSearchHandle.clearState(threadId)`
+     * actually match what the processor stored.
+     */
+    private buildRequestContext;
     private createEventStream;
     private readableStreamToAsyncGenerator;
     /**

package/dist/mastra-harness.js

@@ -8,6 +8,7 @@ import { Agent } from '@mastra/core/agent';
 import { Mastra } from '@mastra/core/mastra';
 import { Memory } from '@mastra/memory';
 import { SkillSearchProcessor, ToolSearchProcessor } from '@mastra/core/processors';
+import { MASTRA_RESOURCE_ID_KEY, MASTRA_THREAD_ID_KEY, RequestContext } from '@mastra/core/request-context';
 import { LibSQLStore } from '@mastra/libsql';
 import { MCPClient } from '@mastra/mcp';
 import { AgentSDKError, AgentSDKErrorType, DEFAULT_MAX_STEPS, HarnessBusOwner, McpServerStatus, } from '@salesforce/sfdx-agent-sdk';
@@ -18,7 +19,7 @@ import { mapToolDefinitionsToMastraTools } from './mastra-tool-mapper.js';
 import { hasEnabledServers, mapToMastraMCPServers } from './mcp-config-mapper.js';
 import { DefaultMastraModelConfigFactory } from './default-model-config-factory.js';
 import { partitionPool, } from './tool-search.js';
-import { Clock, getErrorMessage, RealClock, splitFrontmatterAndBody, UUIDGenerator, } from '@salesforce/agentic-common';
+import { buildSummaryPrompt, Clock, getErrorMessage, RealClock, splitFrontmatterAndBody, UUIDGenerator, } from '@salesforce/agentic-common';
 import { makeAbortable } from './utils.js';
 import { Workspace, LocalFilesystem, LocalSandbox, WORKSPACE_TOOLS, WORKSPACE_TOOLS_PREFIX, } from '@mastra/core/workspace';
 /**
@@ -168,6 +169,63 @@ async function listMarkdownFilesInDirectory(dir) {
     }
     return names.sort().map((name) => path.join(dir, name));
 }
+/**
+ * Tool ChatEvents that carry the optional `serverName` + `annotations`
+ * enrichment fields. Centralizing the predicate lets `createEventStream`
+ * narrow once and apply the lookup uniformly across the three variants.
+ */
+function isEnrichableToolEvent(event) {
+    return event.type === 'tool-call' || event.type === 'tool-result' || event.type === 'tool-approval-request';
+}
+/**
+ * Resolve the MCP server name and (optionally) annotations for a Mastra
+ * namespaced tool name (`${serverName}_${toolName}`). Returns `undefined` for
+ * tools that are not in any MCP server's catalog (consumer-declared client
+ * tools, workspace tools).
+ *
+ * Selects the **longest matching prefix whose catalog actually contains the
+ * tool**. Prefix matching alone is unsafe when two server names are
+ * prefix-related (e.g. OOTB `salesforce_dx` and a user-config
+ * `salesforce_dx_custom`): naive iteration would attribute every
+ * `salesforce_dx_custom_*` call to the OOTB server because `Map` iteration is
+ * insertion order and OOTB usually registers first. The existence check on the
+ * candidate's `tools` map closes that hole; the longest-prefix tiebreak handles
+ * the case where two prefix-related servers both legitimately catalog a tool
+ * with the same suffix.
+ *
+ * `serverState.tools` is keyed by the **namespaced** name (see the
+ * `namespacedTools[\`${serverName}_${toolName}\`] = toolConfig` line in
+ * `createAgent`), so the existence check uses `namespacedToolName` directly.
+ *
+ * Design debt: this walk runs per `tool-call` / `tool-result` /
+ * `tool-approval-request` event and reconstructs the tool→server linkage from
+ * the namespaced string. The cost is O(servers) on a typically-small map, but
+ * three call sites now share the same lookup and consumers (e.g. AFV's W4.8
+ * trust policy) consume the enriched `serverName` for trust-relevant
+ * decisions, so the carry-on-the-record fix — attaching an `McpToolInfo`
+ * reference (with `serverName`) on the in-memory tool-call record from the
+ * moment Mastra namespaces the name — becomes proportionally more valuable.
+ * That fix touches `MastraEventAdapter`, `ToolCallInfo`, and the per-stream
+ * pipeline. If this walk shows up in profiles before then, build a
+ * `Map<namespacedToolName, { serverName, McpToolInfo }>` index keyed off
+ * discovery completion.
+ */
+function lookupMcpToolMetadata(agentState, namespacedToolName) {
+    let best;
+    for (const [serverName, serverState] of agentState.mcpServers) {
+        if (!namespacedToolName.startsWith(`${serverName}_`))
+            continue;
+        if (!(namespacedToolName in serverState.tools))
+            continue;
+        if (!best || serverName.length > best.serverName.length) {
+            best = { serverName, serverState };
+        }
+    }
+    if (!best)
+        return undefined;
+    const annotations = best.serverState.tools[namespacedToolName]?.mcp?.annotations;
+    return { serverName: best.serverName, ...(annotations ? { annotations } : {}) };
+}
 /**
  * Projects the namespaced Mastra `ToolAction` map for one MCP server into the
  * SDK's harness-agnostic `McpToolInfo[]` shape.
@@ -361,78 +419,18 @@ export class MastraHarness {
                     // Fallback only — per-server MCPServerConfig.timeout (in serverDef) takes precedence.
                     timeout: DEFAULT_MCP_TIMEOUT_MS,
                 });
-                const startTime = this.clock.now().getTime();
-                this.buses.emitTelemetry({
-                    type: 'mcp-server-discovery-started',
-                    timestamp: this.clock.now(),
-                    agentId,
-                    serverName,
-                });
-                const ready = makeAbortable(client.listToolsetsWithErrors(), options?.abortSignal)
-                    .then(({ toolsets, errors }) => {
-                    const serverError = errors[serverName];
-                    if (serverError) {
-                        throw new Error(serverError);
-                    }
-                    const tools = toolsets[serverName] ?? {};
-                    const namespacedTools = {};
-                    for (const [toolName, toolConfig] of Object.entries(tools)) {
-                        namespacedTools[`${serverName}_${toolName}`] = toolConfig;
-                    }
-                    mcpServerStates.get(serverName).tools = namespacedTools;
-                    // listToolsetsWithErrors() runs in the background and can resolve
-                    // after `shutdown()` has disposed the buses. Swallow only `DISPOSED`
-                    // so a real listener exception still surfaces.
-                    try {
-                        this.buses.emitTelemetry({
-                            type: 'mcp-server-discovery-completed',
-                            timestamp: this.clock.now(),
-                            agentId,
-                            serverName,
-                            toolCount: Object.keys(namespacedTools).length,
-                            durationMs: this.clock.now().getTime() - startTime,
-                        });
-                    }
-                    catch (emitErr) {
-                        if (!isDisposedError(emitErr))
-                            throw emitErr;
-                    }
-                })
-                    .catch(async (err) => {
-                    if (options?.abortSignal?.aborted) {
-                        await client.disconnect().catch(() => undefined);
-                        mcpServerStates.delete(serverName);
-                    }
-                    else {
-                        const s = mcpServerStates.get(serverName);
-                        if (s)
-                            s.error = err instanceof Error ? err.message : String(err);
-                        this.buses.getLogBus()?.warn('MCP server discovery failed', {
-                            serverName,
-                            error: err instanceof Error ? err.message : String(err),
-                        });
-                        try {
-                            this.buses.emitTelemetry({
-                                type: 'mcp-server-discovery-failed',
-                                timestamp: this.clock.now(),
-                                agentId,
-                                serverName,
-                                durationMs: this.clock.now().getTime() - startTime,
-                                error: err instanceof Error ? err : new Error(String(err)),
-                            });
-                        }
-                        catch (emitErr) {
-                            if (!isDisposedError(emitErr))
-                                throw emitErr;
-                        }
-                    }
-                })
-                    .finally(() => {
-                    const s = mcpServerStates.get(serverName);
-                    if (s)
-                        s.settled = true;
-                });
-                mcpServerStates.set(serverName, { client, ready, settled: false, tools: {} });
+                // Place the entry first so runMcpDiscovery's epoch-gated closures
+                // see a matching epoch on first lookup. The promise chain only runs
+                // after the synchronous `set` and `ready` assignment complete.
+                const entry = {
+                    client,
+                    epoch: 0,
+                    ready: Promise.resolve(),
+                    settled: false,
+                    tools: {},
+                };
+                mcpServerStates.set(serverName, entry);
+                entry.ready = this.runMcpDiscovery(agentId, serverName, client, mcpServerStates, 0, options?.abortSignal);
             }
         }
         const customTools = config.tools?.length ? mapToolDefinitionsToMastraTools(config.tools) : {};
@@ -452,6 +450,13 @@ export class MastraHarness {
             const localFs = new LocalFilesystem({ basePath: projectRoot, allowedPaths });
             const sandbox = new LocalSandbox({
                 workingDirectory: projectRoot,
+                // Forward the parent process env so spawned commands see HOME,
+                // locale, etc. (LocalSandbox's default forwards only PATH, which
+                // breaks tools like `sf org list` that resolve config from HOME).
+                // TODO: filter sensitive vars (e.g. ANTHROPIC_*, AWS creds) here
+                // if the harness ever runs alongside secrets it shouldn't leak
+                // into spawned tool subprocesses.
+                env: process.env,
             });
             mastraWorkspace = new Workspace({
                 filesystem: localFs,
@@ -594,6 +599,87 @@ export class MastraHarness {
         }
         return results;
     }
+    /**
+     * Eager implementation of {@link AgentHarness.reconnectMcpServer}: calls
+     * Mastra's `MCPClient.reconnectServer(name)` primitive, supersedes any
+     * in-flight discovery (replace-ready semantics — latest call wins), and
+     * awaits the new discovery before resolving.
+     *
+     * Programming errors (unknown agent / unknown server / disabled) reject.
+     * Transport / discovery failures are recorded on the per-server state
+     * (visible via {@link getMcpServerInfo}) and a `mcp-server-discovery-*`
+     * telemetry pair fires; the returned promise still resolves.
+     */
+    async reconnectMcpServer(agentId, serverName) {
+        const state = this.getAgentStateOrThrow(agentId);
+        const serverState = state.mcpServers.get(serverName);
+        if (!serverState) {
+            const disabled = state.mcpConfig?.[serverName]?.enabled === false;
+            throw new AgentSDKError(disabled
+                ? `MCP server "${serverName}" is disabled on agent "${agentId}"; reconnect is not valid for disabled servers`
+                : `MCP server "${serverName}" is not configured on agent "${agentId}"`, disabled ? AgentSDKErrorType.MCP_SERVER_DISABLED : AgentSDKErrorType.MCP_SERVER_NOT_FOUND);
+        }
+        // Replace-ready: bump the epoch first so any in-flight discovery from a
+        // previous epoch will drop its writes when its `currentEntry()` lookup fails.
+        // Then reset visible state for the fresh attempt.
+        const epoch = serverState.epoch + 1;
+        serverState.epoch = epoch;
+        serverState.error = undefined;
+        serverState.tools = {};
+        serverState.settled = false;
+        // Emit `started` BEFORE awaiting the transport reconnect so subscribers
+        // see started → (completed | failed) timestamps in chronological order.
+        // The previous shape had `started` emitted from the catch block AFTER
+        // `reconnectServer()` had already thrown — making the started timestamp
+        // appear later than the failure event.
+        const startTime = this.clock.now().getTime();
+        this.buses.emitTelemetry({
+            type: 'mcp-server-discovery-started',
+            timestamp: this.clock.now(),
+            agentId,
+            serverName,
+        });
+        // Transport-level reconnect can throw (stdio spawn failure, remote handshake
+        // failure). Record on serverState and emit only `failed` (not started+failed
+        // — `started` already fired above). This keeps the public contract
+        // ("rejects only on programming errors") honest.
+        try {
+            await serverState.client.reconnectServer(serverName);
+        }
+        catch (err) {
+            // Drop telemetry + state writes if a newer epoch has superseded us.
+            if (state.mcpServers.get(serverName)?.epoch !== epoch)
+                return;
+            const message = err instanceof Error ? err.message : String(err);
+            serverState.error = message;
+            serverState.settled = true;
+            this.buses.getLogBus()?.warn('MCP server reconnect failed', {
+                serverName,
+                error: message,
+            });
+            try {
+                this.buses.emitTelemetry({
+                    type: 'mcp-server-discovery-failed',
+                    timestamp: this.clock.now(),
+                    agentId,
+                    serverName,
+                    durationMs: this.clock.now().getTime() - startTime,
+                    error: err instanceof Error ? err : new Error(message),
+                });
+            }
+            catch (emitErr) {
+                if (!isDisposedError(emitErr))
+                    throw emitErr;
+            }
+            return;
+        }
+        // Pass the existing startTime so runMcpDiscovery does NOT emit its own
+        // `started` (we already did) and the eventual completed/failed event's
+        // durationMs covers the full reconnect+discovery span — one `started`,
+        // one terminal event per public reconnect call.
+        serverState.ready = this.runMcpDiscovery(agentId, serverName, serverState.client, state.mcpServers, epoch, undefined, startTime);
+        await serverState.ready;
+    }
     // ── Thread (Chat Session) Lifecycle ──────────────────────────────
     /**
      * @requirements
@@ -673,7 +759,12 @@ export class MastraHarness {
         this.getAgentStateOrThrow(agentId);
         const messages = await this.getMessages(agentId, threadId);
         if (messages.length === 0) {
-            return this.createThread(agentId);
+            const newThreadId = await this.createThread(agentId);
+            // Empty source still gets compacted away — leaving it attached
+            // would be a leak. Mirrors the non-empty path's destroyThread
+            // call below; both paths exit with the source detached.
+            await this.destroyThreadBestEffort(agentId, threadId);
+            return newThreadId;
         }
         const transcript = messages
             .map((m) => {
@@ -684,7 +775,7 @@ export class MastraHarness {
         let newThreadId;
         try {
             const summaryText = transcript.length <= MAX_TRANSCRIPT_CHARS
-                ? await this.summarize(agentId, this.buildSummaryPrompt(transcript))
+                ? await this.summarize(agentId, buildSummaryPrompt(transcript))
                 : await this.summarizeInChunks(agentId, transcript);
             newThreadId = await this.createThread(agentId);
             const userTs = this.clock.now();
@@ -710,9 +801,26 @@ export class MastraHarness {
         // Source-thread cleanup is best-effort. The user's data is already
         // safely seeded into newThreadId; a failed delete leaves an orphan
         // but does not invalidate the compaction outcome.
-        await this.destroyThread(agentId, threadId).catch(() => undefined);
+        await this.destroyThreadBestEffort(agentId, threadId);
         return newThreadId;
     }
+    /**
+     * Best-effort wrapper around `destroyThread`: swallows the error so the
+     * caller's outer flow isn't aborted by a teardown blip, but logs the
+     * failure at warn level so a future regression in `destroyThread` is
+     * still observable. Used by `compactThread` on both the empty-source and
+     * non-empty paths.
+     */
+    async destroyThreadBestEffort(agentId, threadId) {
+        try {
+            await this.destroyThread(agentId, threadId);
+        }
+        catch (err) {
+            this.buses
+                .getLogBus()
+                ?.warn(`Best-effort destroyThread failed during compactThread for thread "${threadId}" on agent "${agentId}": ${getErrorMessage(err)}`, { agentId, threadId }, err instanceof Error ? err : new Error(String(err)));
+        }
+    }
     // ── Core Messaging ───────────────────────────────────────────────
     /**
      * @requirements
@@ -730,6 +838,7 @@ export class MastraHarness {
                 thread: threadId,
                 resource: this.agentResourceId(agentId),
             },
+            requestContext: this.buildRequestContext(agentId, threadId),
             maxSteps: options?.maxSteps ?? DEFAULT_MAX_STEPS,
             abortSignal: options?.abortSignal,
             requireToolApproval: options?.requireToolApproval,
@@ -767,6 +876,7 @@ export class MastraHarness {
                 thread: threadId,
                 resource: this.agentResourceId(agentId),
             },
+            requestContext: this.buildRequestContext(agentId, threadId),
         });
         return this.toChatStreamResult(agentState, threadId, output);
     }
@@ -804,6 +914,7 @@ export class MastraHarness {
                 thread: threadId,
                 resource: this.agentResourceId(agentId),
             },
+            requestContext: this.buildRequestContext(agentId, threadId),
         });
         return this.toChatStreamResult(agentState, threadId, output);
     }
@@ -835,6 +946,7 @@ export class MastraHarness {
                 thread: threadId,
                 resource: this.agentResourceId(agentId),
             },
+            requestContext: this.buildRequestContext(agentId, threadId),
         });
         return this.toChatStreamResult(agentState, threadId, output);
     }
@@ -898,15 +1010,14 @@ export class MastraHarness {
      */
     async summarize(agentId, prompt) {
         const agentState = this.getAgentStateOrThrow(agentId);
+        // No `requestContext` here on purpose: this is a thread-less one-shot
+        // generate, so per-thread input processors (e.g. ToolSearchProcessor)
+        // harmlessly fall back to the `'default'` bucket. The "every streaming
+        // entry-point passes a RequestContext" invariant in ARCHITECTURE.md
+        // applies to thread-bound entry points only.
         const result = await agentState.agent.generate(prompt, { maxSteps: 1 });
         return result.text;
     }
-    buildSummaryPrompt(transcript) {
-        return ('Summarize the following conversation into a concise context summary. ' +
-            'Preserve key decisions, facts, and action items. ' +
-            'Write in third person as a factual record.\n\n' +
-            transcript);
-    }
     /**
      * Two-pass summarization for transcripts above the char budget. The
      * transcript is split into roughly-equal chunks each ≤ MAX_TRANSCRIPT_CHARS,
@@ -917,7 +1028,7 @@ export class MastraHarness {
         const chunks = this.splitIntoChunks(transcript, MAX_TRANSCRIPT_CHARS);
         const chunkSummaries = [];
         for (const chunk of chunks) {
-            const summary = await this.summarize(agentId, this.buildSummaryPrompt(chunk));
+            const summary = await this.summarize(agentId, buildSummaryPrompt(chunk));
             chunkSummaries.push(summary);
         }
         const merged = chunkSummaries.map((s, i) => `Part ${i + 1}:\n${s}`).join('\n\n');
@@ -925,7 +1036,7 @@ export class MastraHarness {
         // long chat could in principle still exceed budget. If so, recurse;
         // otherwise final-summarize.
         return merged.length <= MAX_TRANSCRIPT_CHARS
-            ? this.summarize(agentId, this.buildSummaryPrompt(merged))
+            ? this.summarize(agentId, buildSummaryPrompt(merged))
             : this.summarizeInChunks(agentId, merged);
     }
     /**
@@ -969,6 +1080,117 @@ export class MastraHarness {
         }
         return state;
     }
+    /**
+     * Runs MCP discovery for one server: emits the
+     * `mcp-server-discovery-started` / `…-completed` / `…-failed` telemetry
+     * triple and writes results into the corresponding entry in `serverStates`.
+     *
+     * The caller MUST `set()` the entry into `serverStates` either before
+     * calling this helper (the `createAgent` path constructs the entry
+     * synchronously after this returns the promise) or before any await
+     * boundary in the helper's chain — the chain reads back via
+     * `serverStates.get(serverName)!` to write `tools` / `error`. Both call
+     * sites today put the entry in place before the first microtask boundary.
+     *
+     * Errors observed during discovery are recorded on the entry's `error`
+     * field; the returned promise resolves either way (it only rejects when
+     * an aborted operation throws after we have already cleaned up). The
+     * `settled` flag is pinned by `.finally()` so `getMcpServerInfo()` can
+     * tell `Connecting` from `Connected`/`Error`.
+     */
+    runMcpDiscovery(agentId, serverName, client, serverStates, epoch, abortSignal, 
+    // When set, reuse the caller's startTime and skip emitting `started` —
+    // the caller (e.g. reconnectMcpServer) has already emitted it. The
+    // resulting completed/failed event still uses the caller's startTime
+    // for `durationMs`, so subscribers see one started → terminal pair.
+    existingStartTime) {
+        const startTime = existingStartTime ?? this.clock.now().getTime();
+        if (existingStartTime === undefined) {
+            this.buses.emitTelemetry({
+                type: 'mcp-server-discovery-started',
+                timestamp: this.clock.now(),
+                agentId,
+                serverName,
+            });
+        }
+        // Returns the entry if and only if it exists AND still belongs to this
+        // epoch. If a fresher discovery has bumped the epoch, this discovery's
+        // writes / telemetry are dropped to prevent orphan-write races with
+        // reconnect's replace-ready semantics.
+        const currentEntry = () => {
+            const s = serverStates.get(serverName);
+            return s && s.epoch === epoch ? s : undefined;
+        };
+        return makeAbortable(client.listToolsetsWithErrors(), abortSignal)
+            .then(({ toolsets, errors }) => {
+            const serverError = errors[serverName];
+            if (serverError) {
+                throw new Error(serverError);
+            }
+            const tools = toolsets[serverName] ?? {};
+            const namespacedTools = {};
+            for (const [toolName, toolConfig] of Object.entries(tools)) {
+                namespacedTools[`${serverName}_${toolName}`] = toolConfig;
+            }
+            const s = currentEntry();
+            if (!s)
+                return; // superseded — drop write + telemetry
+            s.tools = namespacedTools;
+            // listToolsetsWithErrors() runs in the background and can resolve
+            // after `shutdown()` has disposed the buses. Swallow only `DISPOSED`
+            // so a real listener exception still surfaces.
+            try {
+                this.buses.emitTelemetry({
+                    type: 'mcp-server-discovery-completed',
+                    timestamp: this.clock.now(),
+                    agentId,
+                    serverName,
+                    toolCount: Object.keys(namespacedTools).length,
+                    durationMs: this.clock.now().getTime() - startTime,
+                });
+            }
+            catch (emitErr) {
+                if (!isDisposedError(emitErr))
+                    throw emitErr;
+            }
+        })
+            .catch(async (err) => {
+            if (abortSignal?.aborted) {
+                await client.disconnect().catch(() => undefined);
+                if (currentEntry())
+                    serverStates.delete(serverName);
+            }
+            else {
+                const s = currentEntry();
+                if (!s)
+                    return; // superseded — drop write + telemetry
+                s.error = err instanceof Error ? err.message : String(err);
+                this.buses.getLogBus()?.warn('MCP server discovery failed', {
+                    serverName,
+                    error: err instanceof Error ? err.message : String(err),
+                });
+                try {
+                    this.buses.emitTelemetry({
+                        type: 'mcp-server-discovery-failed',
+                        timestamp: this.clock.now(),
+                        agentId,
+                        serverName,
+                        durationMs: this.clock.now().getTime() - startTime,
+                        error: err instanceof Error ? err : new Error(String(err)),
+                    });
+                }
+                catch (emitErr) {
+                    if (!isDisposedError(emitErr))
+                        throw emitErr;
+                }
+            }
+        })
+            .finally(() => {
+            const s = currentEntry();
+            if (s)
+                s.settled = true;
+        });
+    }
     toChatStreamResult(agentState, threadId, output) {
         // Entries are keyed by threadId so each thread holds at most one entry.
         // Stale entries are overwritten here on the next stream; explicit cleanup
@@ -978,7 +1200,7 @@ export class MastraHarness {
         agentState.activeRuns.set(threadId, { runId: output.runId, fullStream: output.fullStream });
         const adapter = new MastraEventAdapter(this.buses.getLogBus());
         return {
-            eventStream: this.createEventStream(output.fullStream, adapter),
+            eventStream: this.createEventStream(output.fullStream, adapter, agentState),
             textStream: this.readableStreamToAsyncGenerator(output.textStream),
         };
     }
@@ -993,11 +1215,41 @@ export class MastraHarness {
     agentResourceId(agentId) {
         return `AGENT_INSTANCE:${agentId}`;
     }
-    async *createEventStream(fullStream, adapter) {
+    /**
+     * Build a `RequestContext` carrying the thread + resource ids under the
+     * reserved keys Mastra reads when input processors (e.g. `ToolSearchProcessor`,
+     * `SkillSearchProcessor`) call `requestContext.get(MASTRA_THREAD_ID_KEY)`.
+     *
+     * `memory: { thread, resource }` populates Mastra's memory layer but is NOT
+     * propagated into the request context, so processors that key per-thread
+     * state by `MASTRA_THREAD_ID_KEY` would otherwise see `undefined` and fall
+     * back to the `'default'` bucket — folding every consumer thread into one.
+     * Surfacing the ids here lets `ToolSearchHandle.clearState(threadId)`
+     * actually match what the processor stored.
+     */
+    buildRequestContext(agentId, threadId) {
+        const ctx = new RequestContext();
+        ctx.set(MASTRA_THREAD_ID_KEY, threadId);
+        ctx.set(MASTRA_RESOURCE_ID_KEY, this.agentResourceId(agentId));
+        return ctx;
+    }
+    async *createEventStream(fullStream, adapter, agentState) {
         for await (const chunk of this.readableStreamToAsyncGenerator(fullStream)) {
             const event = adapter.toChatEvent(chunk);
             if (!event)
                 continue;
+            if (isEnrichableToolEvent(event)) {
+                const toolName = event.type === 'tool-approval-request' ? event.toolCall.toolName : event.toolName;
+                const metadata = lookupMcpToolMetadata(agentState, toolName);
+                if (metadata) {
+                    yield {
+                        ...event,
+                        serverName: metadata.serverName,
+                        ...(metadata.annotations ? { annotations: metadata.annotations } : {}),
+                    };
+                    continue;
+                }
+            }
             yield event;
         }
     }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@salesforce/sfdx-agent-harness-mastra",
-  "version": "0.6.0",
+  "version": "0.8.0",
   "description": "Mastra-backed AgentHarness implementation for @salesforce/sfdx-agent-sdk",
   "type": "module",
   "main": "dist/index.js",
@@ -33,15 +33,15 @@
     "@mastra/mcp": "^1.8.0",
     "@mastra/memory": "^1.19.0",
     "@mastra/schema-compat": "^1.2.10",
-    "@salesforce/agentic-common": "0.4.0",
-    "@salesforce/llm-gateway-sdk": "0.4.0"
+    "@salesforce/agentic-common": "0.5.0",
+    "@salesforce/llm-gateway-sdk": "0.6.0"
   },
   "peerDependencies": {
-    "@salesforce/sfdx-agent-sdk": "0.6.0"
+    "@salesforce/sfdx-agent-sdk": "0.9.0"
   },
   "devDependencies": {
     "@eslint/js": "^10.0.1",
-    "@salesforce/sfdx-agent-sdk": "0.6.0",
+    "@salesforce/sfdx-agent-sdk": "0.9.0",
     "@types/node": "^22.19.17",
     "@vitest/coverage-istanbul": "^4.1.7",
     "@vitest/eslint-plugin": "^1.6.17",