Agents Manager: Restore agent ID and version query param override

packages/agents-manager/src/components/agent-dock/index.tsx

@@ -147,7 +147,7 @@ export default function AgentDock( {
 
 			// Sync local session ID with the server's
 			if ( sessionId !== serverSessionId ) {
-				setSessionId( serverSessionId );
+				setSessionId( serverSessionId, agentId );
 				navigate( '/chat', { state: { sessionId: serverSessionId }, replace: true } );
 			}
 		},
@@ -200,7 +200,7 @@ export default function AgentDock( {
 		},
 		// This ensures the same session ID is used between Big Sky and Calypso agents,
 		// so that messages will be stored in the same conversation.
-		getSessionId: () => sessionId || getStoredSessionId(),
+		getSessionId: () => sessionId || getStoredSessionId( agentId ),
 		setIsBuildingSite,
 		setThinkingMessage,
 	} );
@@ -225,7 +225,7 @@ export default function AgentDock( {
 			const sessionId = conversation.session_id || '';
 
 			abortCurrentRequest();
-			setSessionId( sessionId );
+			setSessionId( sessionId, agentId );
 			navigate( '/chat', { state: { sessionId } } );
 		}
 	};

packages/agents-manager/src/components/unified-ai-agent/index.tsx

@@ -2,7 +2,7 @@ import { getAgentManager } from '@automattic/agenttic-client';
 import { QueryClient, QueryClientProvider } from '@tanstack/react-query';
 import { useEffect, useState, useRef } from '@wordpress/element';
 import { useLocation, useNavigate } from 'react-router-dom';
-import { ORCHESTRATOR_AGENT_ID } from '../../constants';
+import { getAgentConfig } from '../../constants';
 import { useAgentsManagerContext } from '../../contexts';
 import '../../types'; // Import for Window type augmentation
 import { useEmptyViewSuggestions } from '../../hooks/use-empty-view-suggestions';
@@ -43,24 +43,25 @@ function AgentSetup( { currentRoute }: UnifiedAIAgentProps ): JSX.Element | null
 	const isChatRoute = pathname.startsWith( '/chat' );
 	const isNewChat = isChatRoute && !! state?.isNewChat;
 	const routeSessionId = isChatRoute && state?.sessionId;
-	// Use empty `sessionId` for new chat, otherwise use route or stored session ID
-	const sessionId = isNewChat ? '' : routeSessionId || getSessionId();
+	// Read agent/version overrides from browser URL (?agent=, ?version=).
+	// PersistentRouter (memory router) does not track window.location.search.
+	const { agentId, version } = getAgentConfig();
+	const sessionId = isNewChat ? '' : routeSessionId || getSessionId( agentId );
 
 	useEffect( () => {
 		async function initializeAgent(): Promise< void > {
 			// Handle new chat: clear existing session and navigate to clean state
 			if ( isNewChat ) {
 				const agentManager = getAgentManager();
 
-				if ( agentManager.hasAgent( ORCHESTRATOR_AGENT_ID ) ) {
-					// Abort any ongoing requests
-					await agentManager.abortCurrentRequest( ORCHESTRATOR_AGENT_ID );
-					// Remove existing agent to start fresh
-					agentManager.removeAgent( ORCHESTRATOR_AGENT_ID );
+				if ( agentManager.hasAgent( agentId ) ) {
+					// eslint-disable-next-line @typescript-eslint/await-thenable -- ensure abort completes before teardown
+					await agentManager.abortCurrentRequest( agentId );
+					agentManager.removeAgent( agentId );
 				}
 
 				// Clear stored session ID
-				clearSessionId();
+				clearSessionId( agentId );
 				// Clear route state to prevent repeated new chat initialization
 				navigate( '/chat', { replace: true } );
 				return;
@@ -82,13 +83,15 @@ function AgentSetup( { currentRoute }: UnifiedAIAgentProps ): JSX.Element | null
 				toolProvider: providers.toolProvider,
 				contextProvider: providers.contextProvider,
 				environment: 'calypso',
+				agentId,
+				version,
 			} );
 
 			setAgentConfig( config );
 		}
 
 		initializeAgent();
-	}, [ currentRoute, isNewChat, navigate, sessionId, site?.ID ] );
+	}, [ agentId, version, currentRoute, isNewChat, navigate, sessionId, site?.ID ] );
 
 	// Expose agentManager on window for cross-bundle access (e.g., Image Studio)
 	useEffect( () => {

packages/agents-manager/src/utils/agent-session.ts

@@ -9,13 +9,21 @@
  * 3. Client stores UUID in localStorage via setSessionId()
  * 4. Subsequent loads: retrieve UUID from localStorage
  */
+import { ORCHESTRATOR_AGENT_ID } from '../constants';
 
 export const SESSION_STORAGE_KEY = 'agents-manager-session-id';
 const SESSION_EXPIRY_MS = 24 * 60 * 60 * 1000; // 24 hours
 
 /**
  * Session storage format.
  */
+export function getSessionStorageKey( agentId?: string ): string {
+	if ( agentId && agentId !== ORCHESTRATOR_AGENT_ID ) {
+		return `${ SESSION_STORAGE_KEY }-${ agentId }`;
+	}
+	return SESSION_STORAGE_KEY;
+}
+
 interface StoredSession {
 	sessionId: string;
 	timestamp: number;
@@ -26,9 +34,10 @@ interface StoredSession {
  * Returns empty string if no session exists or session expired.
  * @returns The current session ID, or an empty string if no valid session exists.
  */
-export function getSessionId(): string {
+export function getSessionId( agentId?: string ): string {
 	try {
-		const stored = localStorage.getItem( SESSION_STORAGE_KEY );
+		const key = getSessionStorageKey( agentId );
+		const stored = localStorage.getItem( key );
 		if ( stored ) {
 			const session: StoredSession = JSON.parse( stored );
 
@@ -54,13 +63,13 @@ export function getSessionId(): string {
  * Save session ID to localStorage.
  * @param sessionId - The session ID to save.
  */
-export function setSessionId( sessionId: string ): void {
+export function setSessionId( sessionId: string, agentId?: string ): void {
 	try {
 		const session: StoredSession = {
 			sessionId,
 			timestamp: Date.now(),
 		};
-		localStorage.setItem( SESSION_STORAGE_KEY, JSON.stringify( session ) );
+		localStorage.setItem( getSessionStorageKey( agentId ), JSON.stringify( session ) );
 	} catch ( error ) {
 		// eslint-disable-next-line no-console
 		console.error( '[agent-session] Error storing session ID:', error );
@@ -71,9 +80,9 @@ export function setSessionId( sessionId: string ): void {
  * Reset to a new chat (clear session).
  * Returns empty string - server will generate UUID on first message.
  */
-export function clearSessionId(): void {
+export function clearSessionId( agentId?: string ): void {
 	try {
-		localStorage.removeItem( SESSION_STORAGE_KEY );
+		localStorage.removeItem( getSessionStorageKey( agentId ) );
 	} catch ( error ) {
 		// eslint-disable-next-line no-console
 		console.error( '[agent-session] Error clearing session ID:', error );

packages/agents-manager/src/utils/create-agent-config.ts

@@ -7,7 +7,7 @@
 
 import { createCalypsoAuthProvider } from '../auth/calypso-auth-provider';
 import { ORCHESTRATOR_AGENT_ID, ORCHESTRATOR_AGENT_URL } from '../constants';
-import { SESSION_STORAGE_KEY } from './agent-session';
+import { getSessionStorageKey } from './agent-session';
 import type { ContextEntry, ToolProvider, ContextProvider } from '../extension-types';
 import type { UseAgentChatConfig, Ability as AgenticAbility } from '@automattic/agenttic-client';
 
@@ -18,16 +18,15 @@ export interface CreateAgentConfigOptions {
 	toolProvider?: ToolProvider;
 	contextProvider?: ContextProvider;
 	environment?: 'calypso' | 'wp-admin';
+	/** Override the agent ID (e.g., from query string). Defaults to ORCHESTRATOR_AGENT_ID. */
+	agentId?: string;
+	/** Override the agent version (e.g., from query string). Passed via constructorArguments. */
+	version?: string;
 }
 
 /**
- * Resolve context entries by calling `getData()` closures.
- *
- * Takes context entries with optional `getData()` closures and resolves them
- * by calling `getData()` to populate the `data` field. The `getData` function
- * is removed from the resolved entries.
- *
- * This allows fetching live data as needed.
+ * Resolve context entries by calling their `getData()` closures
+ * to populate the `data` field.
  */
 export function resolveContextEntries( entries: ContextEntry[] ): ContextEntry[] {
 	return entries.map( ( entry ) => {
@@ -80,20 +79,27 @@ function wrapToolProvider( toolProvider: ToolProvider ): UseAgentChatConfig[ 'to
  * Create a context provider that resolves context entries.
  */
 function createWrappedContextProvider(
-	contextProvider: ContextProvider
+	contextProvider: ContextProvider,
+	version?: string
 ): UseAgentChatConfig[ 'contextProvider' ] {
 	return {
 		getClientContext: () => {
 			const pluginContext = contextProvider.getClientContext();
 
-			if ( pluginContext.contextEntries?.length ) {
-				return {
-					...pluginContext,
-					contextEntries: resolveContextEntries( pluginContext.contextEntries ),
-				};
-			}
+			const resolvedContext = pluginContext.contextEntries?.length
+				? {
+						...pluginContext,
+						contextEntries: resolveContextEntries( pluginContext.contextEntries ),
+				  }
+				: pluginContext;
 
-			return pluginContext;
+			return {
+				...resolvedContext,
+				constructorArguments: {
+					...( resolvedContext.constructorArguments || {} ),
+					...( version && { version } ),
+				},
+			};
 		},
 	};
 }
@@ -103,14 +109,17 @@ function createWrappedContextProvider(
  */
 function createDefaultContextProvider(
 	currentRoute: string | undefined,
-	environment: string
+	environment: string,
+	version?: string
 ): UseAgentChatConfig[ 'contextProvider' ] {
 	return {
 		getClientContext: () => ( {
 			url: window.location.href,
 			pathname: currentRoute || window.location.pathname,
 			search: window.location.search,
 			environment,
+			// TODO: Remove once agenttic-client supports top-level constructorArguments
+			...( version && { constructorArguments: { version } } ),
 		} ),
 	};
 }
@@ -129,13 +138,15 @@ export function createAgentConfig( options: CreateAgentConfigOptions ): UseAgent
 		toolProvider,
 		contextProvider,
 		environment = 'calypso',
+		agentId = ORCHESTRATOR_AGENT_ID,
+		version,
 	} = options;
 
 	const config: UseAgentChatConfig = {
-		agentId: ORCHESTRATOR_AGENT_ID,
+		agentId,
 		agentUrl: ORCHESTRATOR_AGENT_URL,
 		sessionId,
-		sessionIdStorageKey: SESSION_STORAGE_KEY,
+		sessionIdStorageKey: getSessionStorageKey( agentId ),
 		authProvider: createCalypsoAuthProvider( siteId ),
 		enableStreaming: true,
 	};
@@ -145,9 +156,9 @@ export function createAgentConfig( options: CreateAgentConfigOptions ): UseAgent
 	}
 
 	if ( contextProvider ) {
-		config.contextProvider = createWrappedContextProvider( contextProvider );
+		config.contextProvider = createWrappedContextProvider( contextProvider, version );
 	} else {
-		config.contextProvider = createDefaultContextProvider( currentRoute, environment );
+		config.contextProvider = createDefaultContextProvider( currentRoute, environment, version );
 	}
 
 	return config;