Trace client agent cleanup failures and pin the group's method set

android/samples/mobile-2/app/src/main/java/com/example/typeagentchat/AndroidDeviceAgent.kt

@@ -7,6 +7,24 @@ internal object AndroidDeviceAgent {
     const val NAME = "androidDevice"
     const val CHANNEL_NAME = "agent:$NAME"
     const val SCHEMA_ASSET = "typeagent/androidDeviceSchema.ts"
+
+    /**
+     * Methods this agent answers on its RPC channel, sent as `agentInterface`
+     * at registration.
+     *
+     * The server builds its proxy from this list and, when several devices host
+     * `androidDevice`, rejects one whose list differs from the others. So it has
+     * to describe what `handleAndroidDeviceInvoke` really dispatches: declaring
+     * a method the device cannot answer fails only later, at the call. Keeping
+     * one list for both the declaration and the dispatch guard is what stops the
+     * two from drifting - nothing else checks them against each other, and no CI
+     * job builds this module.
+     */
+    val SUPPORTED_METHODS = listOf("executeAction")
+
+    /** Whether [SUPPORTED_METHODS] covers an incoming RPC method. */
+    fun supports(methodName: String): Boolean = SUPPORTED_METHODS.contains(methodName)
+
     private const val AGENT_DESCRIPTION =
         "Acts on this Android device: sets alarms and countdown timers, shows the " +
             "alarm and timer lists, searches for nearby places, shows a place on the " +
@@ -34,11 +52,14 @@ internal object AndroidDeviceAgent {
             .put("actionDefaultEnabled", true)
             .put("schema", schema)
 
+        val agentInterface = JSONArray()
+        SUPPORTED_METHODS.forEach { agentInterface.put(it) }
+
         return JSONObject()
             .put("name", NAME)
             .put("conversationId", conversationId)
             .put("manifest", manifest)
-            .put("agentInterface", JSONArray().put("executeAction"))
+            .put("agentInterface", agentInterface)
             // Identifies this device so several devices can share one
             // `androidDevice` agent, and so a reconnect replaces this device
             // instead of adding another. `multiInstance` is the opt-in: without

android/samples/mobile-2/app/src/main/java/com/example/typeagentchat/WebSocketManager.kt

@@ -748,7 +748,10 @@ class WebSocketManager internal constructor(
             Log.e(TAG, "Android agent invocation is missing callId.")
             return
         }
-        if (methodName != "executeAction") {
+        // The same list registration declares as agentInterface, so the guard
+        // and the declaration cannot drift apart. It has one entry today; a
+        // second would need its own dispatch below, not just a line in the list.
+        if (!AndroidDeviceAgent.supports(methodName)) {
             sendRpcError(
                 channelName,
                 callId,

android/samples/mobile-2/app/src/test/java/com/example/typeagentchat/AndroidDeviceAgentTest.kt

@@ -23,10 +23,16 @@ class AndroidDeviceAgentTest {
         assertEquals("instance-1", registration.getString("instanceId"))
         assertEquals("Pixel 8", registration.getString("displayName"))
         assertEquals(true, registration.getBoolean("multiInstance"))
-        assertEquals(
-            "executeAction",
-            registration.getJSONArray("agentInterface").getString(0)
-        )
+        // The declared set must be exactly what the RPC dispatcher answers: the
+        // server builds its proxy from this and, with several devices hosting
+        // the agent, rejects one that declares a different set.
+        val declared = registration.getJSONArray("agentInterface")
+        assertEquals(AndroidDeviceAgent.SUPPORTED_METHODS.size, declared.length())
+        for (index in AndroidDeviceAgent.SUPPORTED_METHODS.indices) {
+            assertEquals(AndroidDeviceAgent.SUPPORTED_METHODS[index], declared.getString(index))
+        }
+        assertTrue(AndroidDeviceAgent.supports("executeAction"))
+        assertFalse(AndroidDeviceAgent.supports("getDynamicDisplay"))
         assertEquals(
             "export type AndroidDeviceAction = never;",
             registration

ts/packages/agentRpc/src/client.ts

@@ -46,8 +46,22 @@ import { getActiveTypeAgentSpanAttributes } from "@typeagent/telemetry/traceCont
 export type AgentRpcOptions = {
     trustedContextPropagation?: boolean;
     logger?: RpcStructuredLogger;
+    channelName?: string;
 };
 
+function getAgentChannelName(name: string, options?: AgentRpcOptions): string {
+    return options?.channelName ?? `agent:${name}`;
+}
+
+function getOptionsChannelName(
+    name: string,
+    options?: AgentRpcOptions,
+): string {
+    return options?.channelName === undefined
+        ? `options:${name}`
+        : `${options.channelName}:options`;
+}
+
 /**
  * Race a promise against an AbortSignal. If the signal fires before the
  * promise settles, throw an AbortError immediately (the underlying work
@@ -164,7 +178,9 @@ function createOptionsRpc(
     name: string,
     options?: AgentRpcOptions,
 ) {
-    const channel = channelProvider.createChannel(`options:${name}`);
+    const channel = channelProvider.createChannel(
+        getOptionsChannelName(name, options),
+    );
     const optionsMap = createObjectMap();
     return {
         optionsMap,
@@ -209,7 +225,9 @@ export async function createAgentRpcClient(
     agentInterface: AgentInterfaceFunctionName[],
     options?: AgentRpcOptions,
 ) {
-    const channel = channelProvider.createChannel(`agent:${name}`);
+    const channel = channelProvider.createChannel(
+        getAgentChannelName(name, options),
+    );
     const contextMap = createObjectMap<SessionContext<ShimContext>>();
     // Tracks port registration handles returned by sessionContext.registerPort
     // so the out-of-process agent can release them via the regId we sent back.
@@ -869,7 +887,7 @@ export async function createAgentRpcClient(
         // Options are agent-scoped (created once per initializeAgentContext call)
         // so they can be released when the context is torn down.
         if (optionsRpc !== undefined) {
-            channelProvider.deleteChannel(`options:${name}`);
+            channelProvider.deleteChannel(getOptionsChannelName(name, options));
             optionsRpc = undefined;
         }
         return result;

ts/packages/agentRpc/src/server.ts

@@ -43,8 +43,25 @@ import {
 export type AgentRpcServerOptions = {
     trustedContextPropagation?: boolean;
     logger?: RpcStructuredLogger;
+    channelName?: string;
 };
 
+function getAgentChannelName(
+    name: string,
+    options?: AgentRpcServerOptions,
+): string {
+    return options?.channelName ?? `agent:${name}`;
+}
+
+function getOptionsChannelName(
+    name: string,
+    options?: AgentRpcServerOptions,
+): string {
+    return options?.channelName === undefined
+        ? `options:${name}`
+        : `${options.channelName}:options`;
+}
+
 function getTrustedRpcOptions(
     options: AgentRpcServerOptions | undefined,
 ): RpcOptions | undefined {
@@ -73,7 +90,7 @@ function createOptionsRpc(
     options?: AgentRpcServerOptions,
 ) {
     const optionsChannel: RpcChannel = channelProvider.createChannel(
-        `options:${name}`,
+        getOptionsChannelName(name, options),
     );
     return createRpc<OptionsFunctionCallBack>(
         name,
@@ -107,7 +124,7 @@ export function createAgentRpcServer(
     channelProvider: ChannelProvider,
     options?: AgentRpcServerOptions,
 ) {
-    const channelName = `agent:${name}`;
+    const channelName = getAgentChannelName(name, options);
     const channel = channelProvider.createChannel(channelName);
     let optionsRpc: ReturnType<typeof createOptionsRpc> | undefined;
 

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

@@ -44,11 +44,13 @@ describe("agent action context RPC", () => {
             "test",
             serverAgent,
             serverProvider,
+            { channelName: "agent:test:registration-1" },
         );
         const clientAgent = await createAgentRpcClient(
             "test",
             clientProvider,
             server.agentInterface,
+            { channelName: "agent:test:registration-1" },
         );
 
         try {

ts/packages/agentServer/client/src/agentServerClient.ts

@@ -335,10 +335,40 @@ export function createAgentServerConnection(
         { dispatcher: Dispatcher; connectionId: string }
     >();
 
-    // Client-hosted agents registered on the server, name → agent-rpc server
-    // closeFn. Used to tear down the local rpc server when unregistering,
-    // re-registering, or closing the connection.
-    const clientAgentServers = new Map<string, () => void>();
+    // Client-hosted agents registered on the server. Registration details let
+    // re-registration remove the previous agent while its RPC endpoint is
+    // still alive, so dispatcher lifecycle teardown can reach it.
+    const clientAgentServers = new Map<
+        string,
+        {
+            closeFn: () => void;
+            conversationId: string;
+            instanceId?: string | undefined;
+        }
+    >();
+    let nextClientAgentRegistrationId = 0;
+
+    function resolveClientAgentConversationId(conversationId?: string): string {
+        if (conversationId !== undefined) {
+            if (!joinedConversations.has(conversationId)) {
+                throw new Error(
+                    `Not joined to conversation: ${conversationId}`,
+                );
+            }
+            return conversationId;
+        }
+        if (joinedConversations.size === 1) {
+            return joinedConversations.keys().next().value as string;
+        }
+        if (joinedConversations.size === 0) {
+            throw new Error(
+                "Cannot register client agent: no conversation joined",
+            );
+        }
+        throw new Error(
+            "Cannot register client agent: multiple conversations joined; specify conversationId",
+        );
+    }
 
     let closed = false;
 
@@ -624,22 +654,25 @@ export function createAgentServerConnection(
             conversationId?: string,
             identity?: ClientAgentIdentity,
         ): Promise<void> {
-            // Drop any previous rpc server for this name (e.g. re-registering
-            // after a reconnect, where the old server sat on a stale channel).
-            clientAgentServers.get(name)?.();
-            clientAgentServers.delete(name);
+            const resolvedConversationId =
+                resolveClientAgentConversationId(conversationId);
+            const previous = clientAgentServers.get(name);
+            const registrationId = String(++nextClientAgentRegistrationId);
+            const channelName = `agent:${name}:${registrationId}`;
 
             const { closeFn, agentInterface } = createAgentRpcServer(
                 name,
                 agent,
                 currentChannel,
+                { channelName },
             );
             try {
                 await rpc.invoke("registerClientAgent", {
                     name,
                     manifest,
                     agentInterface,
-                    ...(conversationId !== undefined ? { conversationId } : {}),
+                    conversationId: resolvedConversationId,
+                    registrationId,
                     ...(identity?.instanceId !== undefined
                         ? { instanceId: identity.instanceId }
                         : {}),
@@ -654,7 +687,44 @@ export function createAgentServerConnection(
                 closeFn();
                 throw e;
             }
-            clientAgentServers.set(name, closeFn);
+            if (
+                previous !== undefined &&
+                previous.conversationId !== resolvedConversationId
+            ) {
+                try {
+                    await rpc.invoke("unregisterClientAgent", {
+                        name,
+                        conversationId: previous.conversationId,
+                        ...(previous.instanceId !== undefined
+                            ? { instanceId: previous.instanceId }
+                            : {}),
+                    });
+                } catch (e) {
+                    try {
+                        await rpc.invoke("unregisterClientAgent", {
+                            name,
+                            conversationId: resolvedConversationId,
+                            ...(identity?.instanceId !== undefined
+                                ? { instanceId: identity.instanceId }
+                                : {}),
+                        });
+                    } catch (rollbackError) {
+                        closeFn();
+                        throw new AggregateError(
+                            [e, rollbackError],
+                            `Failed to move client agent '${name}' and roll back the new registration`,
+                        );
+                    }
+                    closeFn();
+                    throw e;
+                }
+            }
+            previous?.closeFn();
+            clientAgentServers.set(name, {
+                closeFn,
+                conversationId: resolvedConversationId,
+                instanceId: identity?.instanceId,
+            });
         },
 
         async unregisterClientAgent(
@@ -669,7 +739,7 @@ export function createAgentServerConnection(
                     ...(instanceId !== undefined ? { instanceId } : {}),
                 });
             } finally {
-                clientAgentServers.get(name)?.();
+                clientAgentServers.get(name)?.closeFn();
                 clientAgentServers.delete(name);
             }
         },
@@ -689,7 +759,7 @@ export function createAgentServerConnection(
             joinedConversations.clear();
             // Client-agent rpc servers were bound to the old channel; drop them
             // so the caller re-registers them on the new channel after re-join.
-            for (const closeFn of clientAgentServers.values()) {
+            for (const { closeFn } of clientAgentServers.values()) {
                 closeFn();
             }
             clientAgentServers.clear();
@@ -702,7 +772,7 @@ export function createAgentServerConnection(
             }
             closed = true;
             debug("Closing agent server connection");
-            for (const closeFn of clientAgentServers.values()) {
+            for (const { closeFn } of clientAgentServers.values()) {
                 closeFn();
             }
             clientAgentServers.clear();

ts/packages/agentServer/protocol/src/protocol.ts

@@ -280,12 +280,13 @@ export type AgentServerInvokeFunctions = {
      * The agent is removed automatically when the connection drops or leaves
      * the conversation.
      *
-     * The client must create its agent-rpc server on the `agent:<name>`
-     * channel (via createAgentRpcServer over the connection channel provider)
-     * before calling this. A client that opts in with `multiInstance` may share
-     * the agent name with other clients carrying the same schema: the server
-     * keeps one registration with an instance per client, and routes each
-     * action to one of them. Without it, a second client is rejected as before.
+     * The client must create its agent-rpc server before calling this. Current
+     * clients pass a `registrationId` and host it on the derived unique channel;
+     * older clients use `agent:<name>`. A client that opts in with
+     * `multiInstance` may share the agent name with other clients carrying the
+     * same schema: the server keeps one registration with an instance per
+     * client, and routes each action to one of them. Without it, a second client
+     * is rejected as before.
      */
     registerClientAgent: (param: RegisterClientAgentParams) => Promise<void>;
     /** Unregister a previously registered client-hosted agent. */
@@ -298,6 +299,12 @@ export type RegisterClientAgentParams = {
     name: string;
     manifest: AppAgentManifest;
     agentInterface: AgentInterfaceFunctionName[];
+    /**
+     * Identifies this RPC endpoint on the connection. When supplied, the
+     * endpoint is hosted on `agent:<name>:<registrationId>`, allowing a
+     * replacement to be validated while the previous endpoint remains live.
+     */
+    registrationId?: string;
     /**
      * Target conversation. If omitted, the server uses the connection's single
      * joined conversation (and errors if the connection has joined none or