[2/3] Add guarded structured execution and resumable interactions

ts/packages/agentRpc/src/client.ts

@@ -631,6 +631,28 @@ export async function createAgentRpcClient(
 
     // The shim needs to implement all the APIs regardless whether the actual agent
     // has that API.  We remove remove it the one that is not necessary below.
+    async function invokeWithActionCancellation<T>(
+        context: ActionContext<ShimContext>,
+        contextParams: ActionContextParams,
+        invoke: () => Promise<T>,
+    ): Promise<T> {
+        const signal = context.abortSignal;
+        signal?.throwIfAborted();
+        const onAbort = () =>
+            rpc.send("cancelAction", {
+                actionContextId: contextParams.actionContextId,
+            });
+        signal?.addEventListener("abort", onAbort, { once: true });
+        try {
+            const pending = invoke();
+            return await (context.waitForCompletionOnAbort
+                ? pending
+                : raceWithSignal(pending, signal));
+        } finally {
+            signal?.removeEventListener("abort", onAbort);
+        }
+    }
+
     const agent: Required<AppAgent> = {
         initializeAgentContext(settings?: AppAgentInitSettings) {
             return rpc.invoke("initializeAgentContext", {
@@ -653,32 +675,14 @@ export async function createAgentRpcClient(
             action: TypeAgentAction,
             context: ActionContext<ShimContext>,
         ) {
-            return withActionContextAsync(context, (contextParams) => {
-                const signal = context.abortSignal;
-                if (signal) {
-                    const onAbort = () =>
-                        rpc.send("cancelAction", {
-                            actionContextId: contextParams.actionContextId,
-                        });
-                    signal.addEventListener("abort", onAbort, { once: true });
-                    return raceWithSignal(
-                        rpc.invoke("executeAction", {
-                            ...contextParams,
-                            action,
-                        }),
-                        signal,
-                    ).finally(() => {
-                        signal.removeEventListener("abort", onAbort);
-                    });
-                }
-                return raceWithSignal(
+            return withActionContextAsync(context, (contextParams) =>
+                invokeWithActionCancellation(context, contextParams, () =>
                     rpc.invoke("executeAction", {
                         ...contextParams,
                         action,
                     }),
-                    signal,
-                );
-            });
+                ),
+            );
         },
         validateWildcardMatch(
             action: AppAction,
@@ -809,6 +813,12 @@ export async function createAgentRpcClient(
                 entityTypeName,
             });
         },
+        cancelChoice(choiceId: string, context: SessionContext<ShimContext>) {
+            return rpc.invoke("cancelChoice", {
+                ...getContextParam(context),
+                choiceId,
+            });
+        },
         handleChoice(
             choiceId: string,
             response:
@@ -819,11 +829,13 @@ export async function createAgentRpcClient(
             context: ActionContext<ShimContext>,
         ) {
             return withActionContextAsync(context, (contextParams) =>
-                rpc.invoke("handleChoice", {
-                    ...contextParams,
-                    choiceId,
-                    response,
-                }),
+                invokeWithActionCancellation(context, contextParams, () =>
+                    rpc.invoke("handleChoice", {
+                        ...contextParams,
+                        choiceId,
+                        response,
+                    }),
+                ),
             );
         },
         getDynamicSchema(

ts/packages/agentRpc/src/common.ts

@@ -138,20 +138,27 @@ export function createChannelProvider(
             return;
         }
         if (message.name === undefined) {
-            debugError(
-                `Missing channel name in message: ${JSON.stringify(message)}`,
-            );
+            debugError("Missing channel name in message");
             return;
         }
         const channelAdapter = channelAdapters.get(message.name);
         if (channelAdapter === undefined) {
             debugError(
-                `Invalid channel name ${message.name} in message (available: ${Array.from(channelAdapters.keys()).join(", ")})`,
+                `Invalid channel name in message (available channels: ${channelAdapters.size})`,
             );
             return;
         }
-        const msgType = message.message?.type || "unknown";
-        const callId = message.message?.callId ?? "n/a";
+        // Remote envelopes may contain capabilities or action parameters.
+        // Log only recognized routing metadata, never arbitrary wire values.
+        const type = message.message?.type;
+        const msgType =
+            typeof type === "string" &&
+            ["call", "invoke", "invokeResult", "invokeError"].includes(type)
+                ? type
+                : "unknown";
+        const callId = Number.isSafeInteger(message.message?.callId)
+            ? message.message.callId
+            : "n/a";
         debug(
             `routing message to channel: ${message.name} (type=${msgType}, callId=${callId})`,
         );

ts/packages/agentRpc/src/rpc.ts

@@ -92,6 +92,14 @@ export type RpcTracingOptions = {
     ) => RpcCorrelationFields | undefined;
 };
 
+/** The transport cannot establish whether an in-flight invocation completed. */
+export class RpcDisconnectedError extends Error {
+    constructor(message = "Agent channel disconnected") {
+        super(message);
+        this.name = "RpcDisconnectedError";
+    }
+}
+
 export type RpcOptions = {
     // When true, a disconnect rejects in-flight calls but leaves invoke/send
     // intact so the rpc can be reattached to a fresh channel via rebind().
@@ -185,11 +193,11 @@ export function createRpc<
     let connected = true;
     let bindGeneration = 0;
     const errorFunc = () => {
-        throw new Error("Agent channel disconnected");
+        throw new RpcDisconnectedError();
     };
     const rejectAllPending = (reason: string) => {
         for (const pendingInvoke of pending.values()) {
-            pendingInvoke.reject(new Error(reason));
+            pendingInvoke.reject(new RpcDisconnectedError(reason));
         }
         pending.clear();
     };
@@ -565,7 +573,7 @@ export function createRpc<
             methodName as string,
             nextCallId++,
         );
-        const error = new Error("Agent channel disconnected");
+        const error = new RpcDisconnectedError();
         emitStructuredStarted(options?.logger, lifecycle);
         emitStructuredCompleted(options?.logger, {
             ...lifecycle,
@@ -605,7 +613,7 @@ export function createRpc<
                 };
                 try {
                     if (!connected) {
-                        throw new Error("Agent channel disconnected");
+                        throw new RpcDisconnectedError();
                     }
                     const correlation = getOutboundCorrelation(
                         options?.tracing,
@@ -672,7 +680,7 @@ export function createRpc<
             invoke(methodName, args),
         send: (methodName: keyof CallTargetFunctions, ...args: any[]) => {
             if (!connected) {
-                throw new Error("Agent channel disconnected");
+                throw new RpcDisconnectedError();
             }
             out(
                 {

ts/packages/agentRpc/src/server.ts

@@ -311,15 +311,26 @@ export function createAgentRpcServer(
                 param.entityTypeName,
             );
         },
+        async cancelChoice(param) {
+            await agent.cancelChoice?.(
+                param.choiceId,
+                getSessionContextShim(param),
+            );
+        },
         async handleChoice(param) {
             if (agent.handleChoice === undefined) {
                 throw new Error("Invalid invocation of handleChoice");
             }
-            return agent.handleChoice(
-                param.choiceId,
-                param.response,
-                getActionContextShim(param),
-            );
+            try {
+                return await agent.handleChoice(
+                    param.choiceId,
+                    param.response,
+                    getActionContextShim(param),
+                );
+            } finally {
+                if (param.actionContextId !== undefined)
+                    actionAbortControllers.delete(param.actionContextId);
+            }
         },
         async getDynamicSchema(param) {
             if (agent.getDynamicSchema === undefined) {

ts/packages/agentRpc/src/types.ts

@@ -276,6 +276,9 @@ export type AgentInvokeFunctions = {
                 | QuestionFormResponse;
         },
     ): Promise<ActionResult | undefined>;
+    cancelChoice(
+        param: Partial<ContextParams> & { choiceId: string },
+    ): Promise<void>;
     getDynamicSchema(
         param: Partial<ContextParams> & { schemaName: string },
     ): Promise<SchemaContent | undefined>;

ts/packages/agentRpc/test/actionContext.spec.ts

@@ -12,8 +12,87 @@ import {
     type ChannelProviderAdapter,
 } from "../src/common.js";
 import { createAgentRpcServer } from "../src/server.js";
+import {
+    ChoiceManager,
+    createYesNoChoiceResult,
+} from "@typeagent/agent-sdk/helpers/action";
 
 describe("agent action context RPC", () => {
+    test("cancels a real SDK choice over agent RPC without invoking its callback", async () => {
+        let clientProvider: ChannelProviderAdapter;
+        let serverProvider: ChannelProviderAdapter;
+        clientProvider = createChannelProviderAdapter(
+            "choice-client",
+            (message, callback) => {
+                queueMicrotask(() =>
+                    serverProvider.notifyMessage(structuredClone(message)),
+                );
+                callback?.(null);
+            },
+        );
+        serverProvider = createChannelProviderAdapter(
+            "choice-server",
+            (message, callback) => {
+                queueMicrotask(() =>
+                    clientProvider.notifyMessage(structuredClone(message)),
+                );
+                callback?.(null);
+            },
+        );
+        const choices = new ChoiceManager();
+        let invoked = 0;
+        const agent: AppAgent = {
+            initializeAgentContext: async () => ({}),
+            executeAction: async () =>
+                createYesNoChoiceResult(choices, "Confirm", async () => {
+                    invoked++;
+                    return undefined;
+                }),
+            handleChoice: (id, response, context) =>
+                choices.handleChoice(id, response, context),
+            cancelChoice: async (id) => {
+                choices.cancelChoice(id);
+            },
+        };
+        const server = createAgentRpcServer("choice", agent, serverProvider);
+        const client = await createAgentRpcClient(
+            "choice",
+            clientProvider,
+            server.agentInterface,
+        );
+        try {
+            const agentContext = await client.initializeAgentContext?.();
+            const sessionContext = {
+                agentContext,
+                sessionContextId: "choice-session",
+            } as SessionContext<unknown>;
+            const actionContext = {
+                sessionContext,
+                isFromReasoningLoop: false,
+            } as ActionContext<unknown>;
+            const result = await client.executeAction!(
+                { schemaName: "choice", actionName: "test" },
+                actionContext,
+            );
+            if (
+                result === undefined ||
+                result.error !== undefined ||
+                result.pendingChoice === undefined
+            )
+                throw new Error("Expected a pending choice");
+            const id = result.pendingChoice.choiceId;
+            await client.cancelChoice!(id, sessionContext);
+            await expect(
+                client.handleChoice!(id, true, actionContext),
+            ).rejects.toThrow("Choice not found or expired");
+            expect(invoked).toBe(0);
+        } finally {
+            server.closeFn();
+            clientProvider.notifyDisconnected();
+            serverProvider.notifyDisconnected();
+        }
+    });
+
     test("propagates workingDirectory to the out-of-process agent", async () => {
         let clientProvider: ChannelProviderAdapter;
         let serverProvider: ChannelProviderAdapter;

ts/packages/agentRpc/test/channelLogging.spec.ts

@@ -0,0 +1,62 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import debug from "debug";
+import { format } from "node:util";
+import { createChannelProviderAdapter } from "../src/common.js";
+
+describe("channel diagnostic privacy", () => {
+    it("never logs payloads or untrusted routing fields, including malformed join envelopes", () => {
+        const previousNamespaces = debug.disable();
+        const previousLog = debug.log;
+        const logs: string[] = [];
+        const marker = "test-only-private-payload-marker";
+        debug.log = (...args: unknown[]) => {
+            logs.push(format(...args));
+        };
+        debug.enable("typeagent:channel-redaction:*");
+        try {
+            const sent: unknown[] = [];
+            const received: unknown[] = [];
+            const provider = createChannelProviderAdapter(
+                "channel-redaction",
+                (message) => {
+                    sent.push(message);
+                },
+            );
+            const channel = provider.createChannel("dispatcher");
+            channel.on("message", (message) => {
+                received.push(message);
+            });
+            const payload = { structuredActions: { resumeToken: marker } };
+            const invoke = {
+                type: "invoke",
+                name: "joinConversation",
+                callId: 1,
+                args: [payload],
+            };
+            const result = { type: "invokeResult", callId: 1, result: payload };
+            channel.send(invoke);
+            provider.notifyMessage({ name: "dispatcher", message: invoke });
+            provider.notifyMessage({ name: "dispatcher", message: result });
+            provider.notifyMessage({ ...payload, message: invoke });
+            provider.notifyMessage({ name: marker, message: invoke });
+            provider.notifyMessage({
+                name: "dispatcher",
+                message: { type: marker, callId: marker, result: payload },
+            });
+            provider.notifyDisconnected();
+
+            expect(sent).toEqual([{ name: "dispatcher", message: invoke }]);
+            expect(received.slice(0, 2)).toEqual([invoke, result]);
+            expect(logs.length).toBeGreaterThan(0);
+            expect(logs.join("\n")).toContain("Missing channel name");
+            expect(logs.join("\n")).toContain("type=invoke");
+            expect(logs.join("\n")).not.toContain(marker);
+            expect(logs.join("\n")).not.toContain("resumeToken");
+        } finally {
+            debug.log = previousLog;
+            debug.enable(previousNamespaces);
+        }
+    });
+});

ts/packages/agentSdk/src/agentInterface.ts

@@ -197,6 +197,7 @@ export interface AppAgent extends Partial<AppAgentCommandInterface> {
     ): Promise<ActionResult | undefined>;
 
     // Choice (yes/no confirmation, multi-select, or multi-question form)
+    cancelChoice?(choiceId: string, context: SessionContext): Promise<void>;
     handleChoice?(
         choiceId: string,
         response:
@@ -478,6 +479,9 @@ export interface ActionContext<T = void> {
     readonly actionIO: ActionIO;
     readonly sessionContext: SessionContext<T>;
     readonly abortSignal?: AbortSignal | undefined;
+    // Hosts retaining shared execution state require transports to await the
+    // actual handler after forwarding abort, rather than racing its response.
+    readonly waitForCompletionOnAbort?: boolean;
 
     // true when this action was dispatched from within the reasoning loop (via MCP execute_action),
     // false when dispatched directly from the translator. Agents can use this to decide whether