diff --git a/command-snapshot.json b/command-snapshot.json
index 41a9fd23..dcd3c7c3 100644
--- a/command-snapshot.json
+++ b/command-snapshot.json
@@ -79,13 +79,12 @@
     "alias": [],
     "command": "agent:preview",
     "flagAliases": [],
-    "flagChars": ["c", "d", "n", "o", "x"],
+    "flagChars": ["d", "n", "o", "x"],
     "flags": [
       "apex-debug",
       "api-name",
       "api-version",
       "authoring-bundle",
-      "client-app",
       "flags-dir",
       "output-dir",
       "target-org",
diff --git a/package.json b/package.json
index a9383b9f..ee1b8c78 100644
--- a/package.json
+++ b/package.json
@@ -9,7 +9,7 @@
     "@inquirer/prompts": "^7.10.1",
     "@oclif/core": "^4",
     "@oclif/multi-stage-output": "^0.8.29",
-    "@salesforce/agents": "^0.20.0",
+    "@salesforce/agents": "0.20.0-beta.4",
     "@salesforce/core": "^8.24.0",
     "@salesforce/kit": "^3.2.3",
     "@salesforce/sf-plugins-core": "^12.2.6",
diff --git a/src/agentActivation.ts b/src/agentActivation.ts
index d5ca4282..07d9c110 100644
--- a/src/agentActivation.ts
+++ b/src/agentActivation.ts
@@ -14,8 +14,8 @@
  * limitations under the License.
  */
 
-import { Connection, Messages, Org, SfError } from '@salesforce/core';
-import { Agent, type BotMetadata } from '@salesforce/agents';
+import { Messages, Org, SfError, SfProject } from '@salesforce/core';
+import { Agent, type BotMetadata, ProductionAgent } from '@salesforce/agents';
 import { select } from '@inquirer/prompts';
 
 type Choice<Value> = {
@@ -67,16 +67,15 @@ export const getAgentChoices = (agents: BotMetadata[], status: 'Active' | 'Inact
   });
 
 export const getAgentForActivation = async (config: {
-  conn: Connection;
   targetOrg: Org;
   status: 'Active' | 'Inactive';
   apiNameFlag?: string;
-}): Promise<Agent> => {
-  const { conn, targetOrg, status, apiNameFlag } = config;
+}): Promise<ProductionAgent> => {
+  const { targetOrg, status, apiNameFlag } = config;
 
   let agentsInOrg: BotMetadata[] = [];
   try {
-    agentsInOrg = await Agent.listRemote(conn);
+    agentsInOrg = await Agent.listRemote(targetOrg.getConnection());
   } catch (error) {
     throw SfError.create({
       message: 'Error listing agents in org',
@@ -105,5 +104,9 @@ export const getAgentForActivation = async (config: {
     selectedAgent = agentsInOrg.find((agent) => agent.DeveloperName === agentChoice.DeveloperName);
   }
 
-  return new Agent({ connection: conn, nameOrId: selectedAgent!.Id });
+  return Agent.init({
+    connection: targetOrg.getConnection(),
+    apiNameOrId: selectedAgent!.Id,
+    project: SfProject.getInstance(),
+  });
 };
diff --git a/src/commands/agent/activate.ts b/src/commands/agent/activate.ts
index 862b21b1..30d68706 100644
--- a/src/commands/agent/activate.ts
+++ b/src/commands/agent/activate.ts
@@ -39,13 +39,12 @@ export default class AgentActivate extends SfCommand<void> {
 
     const apiNameFlag = flags['api-name'];
     const targetOrg = flags['target-org'];
-    const conn = targetOrg.getConnection(flags['api-version']);
 
     if (!apiNameFlag && this.jsonEnabled()) {
       throw messages.createError('error.missingRequiredFlags', ['api-name']);
     }
 
-    const agent = await getAgentForActivation({ conn, targetOrg, status: 'Active', apiNameFlag });
+    const agent = await getAgentForActivation({ targetOrg, status: 'Active', apiNameFlag });
     await agent.activate();
     const agentName = (await agent.getBotMetadata()).DeveloperName;
 
diff --git a/src/commands/agent/deactivate.ts b/src/commands/agent/deactivate.ts
index bd93cd71..733764ee 100644
--- a/src/commands/agent/deactivate.ts
+++ b/src/commands/agent/deactivate.ts
@@ -39,13 +39,12 @@ export default class AgentDeactivate extends SfCommand<void> {
 
     const apiNameFlag = flags['api-name'];
     const targetOrg = flags['target-org'];
-    const conn = targetOrg.getConnection(flags['api-version']);
 
     if (!apiNameFlag && this.jsonEnabled()) {
       throw messages.createError('error.missingRequiredFlags', ['api-name']);
     }
 
-    const agent = await getAgentForActivation({ conn, targetOrg, status: 'Inactive', apiNameFlag });
+    const agent = await getAgentForActivation({ targetOrg, status: 'Inactive', apiNameFlag });
     await agent.deactivate();
     const agentName = (await agent.getBotMetadata()).DeveloperName;
 
diff --git a/src/commands/agent/generate/authoring-bundle.ts b/src/commands/agent/generate/authoring-bundle.ts
index d80c69e7..284b0d8e 100644
--- a/src/commands/agent/generate/authoring-bundle.ts
+++ b/src/commands/agent/generate/authoring-bundle.ts
@@ -18,7 +18,7 @@ import { join, resolve } from 'node:path';
 import { readFileSync, existsSync } from 'node:fs';
 import { SfCommand, Flags } from '@salesforce/sf-plugins-core';
 import { generateApiName, Messages, SfError } from '@salesforce/core';
-import { Agent, AgentJobSpec } from '@salesforce/agents';
+import { AgentJobSpec, ScriptAgent } from '@salesforce/agents';
 import YAML from 'yaml';
 import { input as inquirerInput } from '@inquirer/prompts';
 import { theme } from '../../../inquirer-theme.js';
@@ -102,7 +102,7 @@ export default class AgentGenerateAuthoringBundle extends SfCommand<AgentGenerat
 
   public async run(): Promise<AgentGenerateAuthoringBundleResult> {
     const { flags } = await this.parse(AgentGenerateAuthoringBundle);
-    const { 'output-dir': outputDir, 'target-org': targetOrg } = flags;
+    const { 'output-dir': outputDir } = flags;
 
     // If we don't have a spec yet, prompt for it
     const spec = flags.spec ?? (await promptForYamlFile(AgentGenerateAuthoringBundle.FLAGGABLE_PROMPTS['spec']));
@@ -135,10 +135,8 @@ export default class AgentGenerateAuthoringBundle extends SfCommand<AgentGenerat
       const metaXmlPath = join(targetOutputDir, `${bundleApiName}.bundle-meta.xml`);
 
       // Write Agent file
-      const conn = targetOrg.getConnection(flags['api-version']);
       const specContents = YAML.parse(readFileSync(spec, 'utf8')) as AgentJobSpec;
-      await Agent.createAuthoringBundle({
-        connection: conn,
+      await ScriptAgent.createAuthoringBundle({
         agentSpec: { ...specContents, ...{ name, developerName: bundleApiName } },
         project: this.project!,
         bundleApiName,
diff --git a/src/commands/agent/preview.ts b/src/commands/agent/preview.ts
index 13cb9de6..491b7097 100644
--- a/src/commands/agent/preview.ts
+++ b/src/commands/agent/preview.ts
@@ -14,35 +14,18 @@
  * limitations under the License.
  */
 
-import * as path from 'node:path';
-import { join, resolve } from 'node:path';
-import { globSync } from 'glob';
+import { resolve } from 'node:path';
 import { Flags, SfCommand } from '@salesforce/sf-plugins-core';
-import { AuthInfo, Connection, Lifecycle, Logger, Messages, SfError } from '@salesforce/core';
 import React from 'react';
 import { render } from 'ink';
-import {
-  AgentPreview as Preview,
-  AgentSimulate,
-  AgentSource,
-  findAuthoringBundle,
-  PublishedAgent,
-  ScriptAgent,
-} from '@salesforce/agents';
+import { Agent, AgentSource, PreviewableAgent, ProductionAgent, ScriptAgent } from '@salesforce/agents';
 import { select } from '@inquirer/prompts';
+import { Lifecycle, Messages, SfError } from '@salesforce/core';
 import { AgentPreviewReact } from '../../components/agent-preview-react.js';
 
 Messages.importMessagesDirectoryFromMetaUrl(import.meta.url);
 const messages = Messages.loadMessages('@salesforce/plugin-agent', 'agent.preview');
 
-let logger: Logger;
-const getLogger = (): Logger => {
-  if (!logger) {
-    logger = Logger.childFromRoot('plugin-agent-preview');
-  }
-  return logger;
-};
-
 type BotVersionStatus = { Status: 'Active' | 'Inactive' };
 
 export type AgentData = {
@@ -53,12 +36,6 @@ export type AgentData = {
   };
 };
 
-type Choice<Value> = {
-  value: Value;
-  name?: string;
-  disabled?: boolean | string;
-};
-
 // https://developer.salesforce.com/docs/einstein/genai/guide/agent-api-get-started.html#prerequisites
 export const UNSUPPORTED_AGENTS = ['Copilot_for_Salesforce'];
 
@@ -74,11 +51,6 @@ export default class AgentPreview extends SfCommand<AgentPreviewResult> {
   public static readonly flags = {
     'target-org': Flags.requiredOrg(),
     'api-version': Flags.orgApiVersion(),
-    'client-app': Flags.string({
-      char: 'c',
-      summary: messages.getMessage('flags.client-app.summary'),
-      dependsOn: ['target-org'],
-    }),
     'api-name': Flags.string({
       summary: messages.getMessage('flags.api-name.summary'),
       char: 'n',
@@ -105,146 +77,70 @@ export default class AgentPreview extends SfCommand<AgentPreviewResult> {
     // get user's agent selection either from flags, or interaction
     // if .agent selected, use the AgentSimulate class to preview
     // if published agent, use AgentPreview for preview
-    // based on agent, differing auth mechanisms required
     const { flags } = await this.parse(AgentPreview);
 
-    const { 'api-name': apiNameFlag, 'use-live-actions': useLiveActions } = flags;
+    const { 'api-name': apiNameOrId, 'use-live-actions': useLiveActions, 'authoring-bundle': aabName } = flags;
     const conn = flags['target-org'].getConnection(flags['api-version']);
 
-    const agentsInOrg = (
-      await conn.query<AgentData>(
-        'SELECT Id, DeveloperName, (SELECT Status FROM BotVersions) FROM BotDefinition WHERE IsDeleted = false'
-      )
-    ).records;
-
-    let selectedAgent: ScriptAgent | PublishedAgent;
+    let selectedAgent: ScriptAgent | ProductionAgent;
 
-    if (flags['authoring-bundle']) {
-      // user specified --authoring-bundle, we'll find the script and use it
-      const bundlePath = findAuthoringBundle(this.project!.getPath(), flags['authoring-bundle']);
-      if (!bundlePath) {
-        throw new SfError(`Could not find authoring bundle for ${flags['authoring-bundle']}`);
-      }
-      selectedAgent = {
-        DeveloperName: flags['authoring-bundle'],
-        source: AgentSource.SCRIPT,
-        path: join(bundlePath, `${flags['authoring-bundle']}.agent`),
-      };
-    } else if (apiNameFlag) {
-      // user specified --api-name, it should be in the list of agents from the org
-      const agent = agentsInOrg.find((a) => a.DeveloperName === apiNameFlag);
-      if (!agent) throw new Error(`No valid Agents were found with the Api Name ${apiNameFlag}.`);
-      validateAgent(agent);
-      selectedAgent = {
-        Id: agent.Id,
-        DeveloperName: agent.DeveloperName,
-        source: AgentSource.PUBLISHED,
-      };
-      if (!selectedAgent) throw new Error(`No valid Agents were found with the Api Name ${apiNameFlag}.`);
+    if (aabName) {
+      // user specified --authoring-bundle, use the API name directly
+      selectedAgent = await Agent.init({ connection: conn, project: this.project!, aabName });
+    } else if (apiNameOrId) {
+      selectedAgent = await Agent.init({ connection: conn, project: this.project!, apiNameOrId });
     } else {
-      selectedAgent = await select<ScriptAgent | PublishedAgent>({
+      const previewableAgents = await Agent.listPreviewable(conn, this.project!);
+      const choices = previewableAgents.map((agent) => ({
+        name: agent.source === AgentSource.PUBLISHED ? `${agent.name} (Published)` : `${agent.name} (Agent Script)`,
+        value: agent,
+      }));
+      const choice = await select<PreviewableAgent>({
         message: 'Select an agent',
-        choices: this.getAgentChoices(agentsInOrg),
+        choices,
       });
-    }
-
-    // we have the selected agent, create the appropriate connection
-    const authInfo = await AuthInfo.create({
-      username: flags['target-org'].getUsername(),
-    });
-    // Get client app - check flag first, then auth file, then env var
-    let clientApp = flags['client-app'];
 
-    if (!clientApp && selectedAgent?.source === AgentSource.PUBLISHED) {
-      const clientApps = getClientAppsFromAuth(authInfo);
-
-      if (clientApps.length === 1) {
-        clientApp = clientApps[0];
-      } else if (clientApps.length > 1) {
-        clientApp = await select({
-          message: 'Select a client app',
-          choices: clientApps.map((app) => ({ value: app, name: app })),
+      if (choice.source === AgentSource.SCRIPT && choice.name) {
+        // Use the API name directly
+        selectedAgent = await Agent.init({
+          connection: conn,
+          project: this.project!,
+          aabName: choice.name,
         });
+        selectedAgent.preview.setMockMode(flags['use-live-actions'] ? 'Live Test' : 'Mock');
       } else {
-        throw new SfError('No client app found.');
+        selectedAgent = await Agent.init({
+          connection: conn,
+          project: this.project!,
+          // developerName will be set at this point since the user selected a production agent, even ID will be defined
+          apiNameOrId: choice.developerName ?? choice.id ?? '',
+        });
       }
     }
 
-    if (useLiveActions && selectedAgent.source === AgentSource.PUBLISHED) {
+    if (useLiveActions && selectedAgent instanceof ProductionAgent) {
       void Lifecycle.getInstance().emitWarning(
         'Published agents will always use real actions in your org, specifying --use-live-actions and selecting a published agent has no effect'
       );
     }
 
-    const jwtConn =
-      selectedAgent?.source === AgentSource.PUBLISHED
-        ? await Connection.create({
-            authInfo,
-            clientApp,
-          })
-        : await Connection.create({ authInfo });
-
     // Only resolve outputDir if explicitly provided via flag
     // Otherwise, let user decide when exiting
     const outputDir = flags['output-dir'] ? resolve(flags['output-dir']) : undefined;
-    // Both classes share the same interface for the methods we need
-    const agentPreview =
-      selectedAgent.source === AgentSource.PUBLISHED
-        ? new Preview(jwtConn, selectedAgent.Id)
-        : new AgentSimulate(jwtConn, selectedAgent.path, !useLiveActions);
 
-    agentPreview.setApexDebugMode(flags['apex-debug']);
+    selectedAgent.preview.setApexDebugging(flags['apex-debug']);
 
     const instance = render(
       React.createElement(AgentPreviewReact, {
-        connection: conn,
-        agent: agentPreview,
-        name: selectedAgent.DeveloperName,
+        agent: selectedAgent.preview,
+        name: selectedAgent.name ?? '',
         outputDir,
-        isLocalAgent: selectedAgent.source === AgentSource.SCRIPT,
-        apexDebug: flags['apex-debug'],
-        logger: getLogger(),
+        isLocalAgent: selectedAgent instanceof ScriptAgent,
       }),
       { exitOnCtrlC: false }
     );
     await instance.waitUntilExit();
   }
-
-  private getAgentChoices(agents: AgentData[]): Array<Choice<ScriptAgent | PublishedAgent>> {
-    const choices: Array<Choice<ScriptAgent | PublishedAgent>> = [];
-
-    // Add org agents
-    for (const agent of agents) {
-      if (agentIsInactive(agent) || agentIsUnsupported(agent.DeveloperName)) {
-        continue;
-      }
-
-      choices.push({
-        name: `${agent.DeveloperName} (Published)`,
-        value: {
-          Id: agent.Id,
-          DeveloperName: agent.DeveloperName,
-          source: AgentSource.PUBLISHED,
-        },
-      });
-    }
-
-    // Add local agents from .agent files
-    const localAgentPaths = globSync('**/*.agent', { cwd: this.project!.getPath() });
-    for (const agentPath of localAgentPaths) {
-      const agentName = path.basename(agentPath, '.agent');
-      choices.push({
-        name: `${agentName} (Agent Script)`,
-        value: {
-          DeveloperName: agentName,
-          source: AgentSource.SCRIPT,
-          path: path.join(this.project!.getPath(), agentPath),
-        },
-      });
-    }
-
-    return choices;
-  }
 }
 
 export const agentIsUnsupported = (devName: string): boolean => UNSUPPORTED_AGENTS.includes(devName);
@@ -267,6 +163,3 @@ export const validateAgent = (agent: AgentData): boolean => {
 
   return true;
 };
-
-export const getClientAppsFromAuth = (authInfo: AuthInfo): string[] =>
-  Object.keys(authInfo.getFields().clientApps ?? {});
diff --git a/src/commands/agent/publish/authoring-bundle.ts b/src/commands/agent/publish/authoring-bundle.ts
index 10f811b7..342fd5c8 100644
--- a/src/commands/agent/publish/authoring-bundle.ts
+++ b/src/commands/agent/publish/authoring-bundle.ts
@@ -14,12 +14,10 @@
  * limitations under the License.
  */
 import { EOL } from 'node:os';
-import { join } from 'node:path';
-import { readFileSync } from 'node:fs';
 import { SfCommand, Flags } from '@salesforce/sf-plugins-core';
 import { MultiStageOutput } from '@oclif/multi-stage-output';
 import { Messages, Lifecycle, SfError } from '@salesforce/core';
-import { Agent, findAuthoringBundle } from '@salesforce/agents';
+import { Agent } from '@salesforce/agents';
 import { RequestStatus, ScopedPostDeploy, type ScopedPostRetrieve } from '@salesforce/source-deploy-retrieve';
 import { ensureArray } from '@salesforce/kit';
 import { FlaggablePrompt, promptForAgentFiles } from '../../../flags.js';
@@ -70,24 +68,15 @@ export default class AgentPublishAuthoringBundle extends SfCommand<AgentPublishA
   public async run(): Promise<AgentPublishAuthoringBundleResult> {
     const { flags } = await this.parse(AgentPublishAuthoringBundle);
     // If api-name is not provided, prompt user to select an .agent file from the project and extract the API name from it
-    const apiName =
+    const aabName =
       flags['api-name'] ??
       (await promptForAgentFiles(this.project!, AgentPublishAuthoringBundle.FLAGGABLE_PROMPTS['api-name']));
-    const authoringBundleDir = findAuthoringBundle(
-      this.project!.getPackageDirectories().map((dir) => dir.fullPath),
-      apiName
-    );
-
-    if (!authoringBundleDir) {
-      throw new SfError(messages.getMessage('error.agentNotFound', [apiName]), 'AgentNotFoundError', [
-        messages.getMessage('error.agentNotFoundAction'),
-      ]);
-    }
+
     // Create multi-stage output
     const mso = new MultiStageOutput<{ agentName: string }>({
       stages: ['Validate Bundle', 'Publish Agent', 'Retrieve Metadata', 'Deploy Metadata'],
       title: 'Publishing Agent',
-      data: { agentName: apiName },
+      data: { agentName: aabName },
       jsonEnabled: this.jsonEnabled(),
       postStagesBlock: [
         {
@@ -103,12 +92,10 @@ export default class AgentPublishAuthoringBundle extends SfCommand<AgentPublishA
       mso.goto('Validate Bundle');
       const targetOrg = flags['target-org'];
       const conn = targetOrg.getConnection(flags['api-version']);
+      const agent = await Agent.init({ connection: conn, project: this.project!, aabName });
 
       // First compile the .agent file to get the Agent JSON
-      const compileResponse = await Agent.compileAgentScript(
-        conn,
-        readFileSync(join(authoringBundleDir, `${apiName}.agent`), 'utf8')
-      );
+      const compileResponse = await agent.compile();
       if (compileResponse.status === 'success') {
         mso.skipTo('Publish Agent');
       } else {
@@ -156,7 +143,7 @@ export default class AgentPublishAuthoringBundle extends SfCommand<AgentPublishA
         return Promise.resolve();
       });
 
-      const result = await Agent.publishAgentJson(conn, this.project!, compileResponse.compiledArtifact);
+      const result = await agent.publish();
       mso.stop();
 
       return {
diff --git a/src/commands/agent/validate/authoring-bundle.ts b/src/commands/agent/validate/authoring-bundle.ts
index bf1204da..923e90cb 100644
--- a/src/commands/agent/validate/authoring-bundle.ts
+++ b/src/commands/agent/validate/authoring-bundle.ts
@@ -13,12 +13,10 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-import { readFileSync } from 'node:fs';
-import { join } from 'node:path';
 import { SfCommand, Flags } from '@salesforce/sf-plugins-core';
 import { Messages, SfError } from '@salesforce/core';
 import { MultiStageOutput } from '@oclif/multi-stage-output';
-import { Agent, findAuthoringBundle } from '@salesforce/agents';
+import { Agent } from '@salesforce/agents';
 import { colorize } from '@oclif/core/ux';
 import { throwAgentCompilationError } from '../../../common.js';
 import { FlaggablePrompt, promptForAgentFiles } from '../../../flags.js';
@@ -67,21 +65,12 @@ export default class AgentValidateAuthoringBundle extends SfCommand<AgentValidat
   public async run(): Promise<AgentValidateAuthoringBundleResult> {
     const { flags } = await this.parse(AgentValidateAuthoringBundle);
     // If api-name is not provided, prompt user to select an .agent file from the project and extract the API name from it
-    const apiName =
+    const aabName =
       flags['api-name'] ??
       (await promptForAgentFiles(this.project!, AgentValidateAuthoringBundle.FLAGGABLE_PROMPTS['api-name']));
-    const authoringBundleDir = findAuthoringBundle(
-      this.project!.getPackageDirectories().map((dir) => dir.fullPath),
-      apiName
-    );
-    if (!authoringBundleDir) {
-      throw new SfError(messages.getMessage('error.agentNotFound', [apiName]), 'AgentNotFoundError', [
-        messages.getMessage('error.agentNotFoundAction'),
-      ]);
-    }
     const mso = new MultiStageOutput<{ status: string; errors: string }>({
       jsonEnabled: this.jsonEnabled(),
-      title: `Validating ${apiName} Authoring Bundle`,
+      title: `Validating ${aabName} Authoring Bundle`,
       showTitle: true,
       stages: ['Validating Authoring Bundle'],
       stageSpecificBlock: [
@@ -104,10 +93,8 @@ export default class AgentValidateAuthoringBundle extends SfCommand<AgentValidat
       mso.skipTo('Validating Authoring Bundle');
       const targetOrg = flags['target-org'];
       const conn = targetOrg.getConnection(flags['api-version']);
-      const result = await Agent.compileAgentScript(
-        conn,
-        readFileSync(join(authoringBundleDir, `${apiName}.agent`), 'utf8')
-      );
+      const agent = await Agent.init({ connection: conn, project: this.project!, aabName });
+      const result = await agent.compile();
       if (result.status === 'success') {
         mso.updateData({ status: 'COMPLETED' });
         mso.stop('completed');
diff --git a/src/components/agent-preview-react.tsx b/src/components/agent-preview-react.tsx
index e4fcbf77..a6095b3c 100644
--- a/src/components/agent-preview-react.tsx
+++ b/src/components/agent-preview-react.tsx
@@ -14,17 +14,18 @@
  * limitations under the License.
  */
 
-import path from 'node:path';
-import fs from 'node:fs';
 import * as process from 'node:process';
 import { resolve } from 'node:path';
 import React from 'react';
 import { Box, Text, useInput } from 'ink';
 import TextInput from 'ink-text-input';
-import { Connection, SfError, Lifecycle, Logger } from '@salesforce/core';
-import { AgentPreviewBase, AgentPreviewSendResponse, writeDebugLog } from '@salesforce/agents';
+import { SfError, Lifecycle } from '@salesforce/core';
+import { ScriptAgent, ProductionAgent } from '@salesforce/agents';
 import { sleep, env } from '@salesforce/kit';
-import { PlannerResponse } from '@salesforce/agents/lib/types.js';
+
+type ScriptAgentPreview = ScriptAgent['preview'];
+type ProductionAgentPreview = ProductionAgent['preview'];
+type AgentPreview = ScriptAgentPreview | ProductionAgentPreview;
 
 // Component to show a simple typing animation
 function Typing(): React.ReactNode {
@@ -50,45 +51,6 @@ function Typing(): React.ReactNode {
   );
 }
 
-export const saveTranscriptsToFile = (
-  outputDir: string,
-  messages: Array<{ timestamp: Date; role: string; content: string }>,
-  responses: AgentPreviewSendResponse[],
-  traces?: PlannerResponse[]
-): void => {
-  if (!outputDir) return;
-  fs.mkdirSync(outputDir, { recursive: true });
-
-  const transcriptPath = path.join(outputDir, 'transcript.json');
-  fs.writeFileSync(transcriptPath, JSON.stringify(messages, null, 2));
-
-  const responsesPath = path.join(outputDir, 'responses.json');
-  fs.writeFileSync(responsesPath, JSON.stringify(responses, null, 2));
-
-  if (traces) {
-    const tracesPath = path.join(outputDir, 'traces.json');
-    fs.writeFileSync(tracesPath, JSON.stringify(traces, null, 2));
-  }
-};
-
-export const getTraces = async (
-  agent: AgentPreviewBase,
-  sessionId: string,
-  messageIds: string[],
-  logger: Logger
-): Promise<PlannerResponse[]> => {
-  if (messageIds.length > 0) {
-    try {
-      const traces = await agent.traces(sessionId, messageIds);
-      return traces;
-    } catch (e) {
-      const sfError = SfError.wrap(e);
-      logger.info(`Error obtaining traces: ${sfError.name} - ${sfError.message}`, { sessionId, messageIds });
-    }
-  }
-  return [];
-};
-
 /**
  * Ideas:
  * - Limit height based on terminal height
@@ -97,13 +59,10 @@ export const getTraces = async (
  * - Add keystroke to scroll down
  */
 export function AgentPreviewReact(props: {
-  readonly connection: Connection;
-  readonly agent: AgentPreviewBase;
+  readonly agent: AgentPreview;
   readonly name: string;
   readonly outputDir: string | undefined;
   readonly isLocalAgent: boolean;
-  readonly apexDebug: boolean | undefined;
-  readonly logger: Logger;
 }): React.ReactNode {
   const [messages, setMessages] = React.useState<Array<{ timestamp: Date; role: string; content: string }>>([]);
   const [header, setHeader] = React.useState('Starting session...');
@@ -115,16 +74,9 @@ export function AgentPreviewReact(props: {
   const [showSavePrompt, setShowSavePrompt] = React.useState(false);
   const [showDirInput, setShowDirInput] = React.useState(false);
   const [saveDir, setSaveDir] = React.useState('');
-  const [saveConfirmed, setSaveConfirmed] = React.useState(false);
-  // @ts-expect-error: Complains if this is not defined but it's not used
-  // eslint-disable-next-line @typescript-eslint/no-unused-vars
-  const [timestamp, setTimestamp] = React.useState(new Date().getTime());
-  const [tempDir, setTempDir] = React.useState('');
-  const [responses, setResponses] = React.useState<AgentPreviewSendResponse[]>([]);
-  const [apexDebugLogs, setApexDebugLogs] = React.useState<string[]>([]);
-  const [messageIds, setMessageIds] = React.useState<string[]>([]);
+  const [savedPath, setSavedPath] = React.useState<string | undefined>();
 
-  const { connection, agent, name, outputDir, isLocalAgent, apexDebug, logger } = props;
+  const { agent, name, outputDir, isLocalAgent } = props;
 
   useInput((input, key) => {
     // If user is in directory input and presses ESC, cancel and exit without saving
@@ -149,13 +101,11 @@ export function AgentPreviewReact(props: {
           // If outputDir was provided via flag, use it directly
           if (outputDir) {
             setSaveDir(outputDir);
-            setSaveConfirmed(true);
-            setShowSavePrompt(false);
           } else {
             // Otherwise, prompt for directory
             setShowSavePrompt(false);
             setShowDirInput(true);
-            const defaultDir = env.getString('SF_AGENT_PREVIEW_OUTPUT_DIR', path.join('temp', 'agent-preview'));
+            const defaultDir = env.getString('SF_AGENT_PREVIEW_OUTPUT_DIR', 'temp/agent-preview');
             setSaveDir(defaultDir);
           }
         } else {
@@ -171,7 +121,12 @@ export function AgentPreviewReact(props: {
       if (sessionEnded) {
         try {
           // TODO: Support other end types (such as Escalate)
-          await agent.end(sessionId, 'UserRequest');
+          // ScriptAgent.end() takes no args, ProductionAgent.end(reason) takes EndReason
+          if (isLocalAgent) {
+            await agent.end();
+          } else {
+            await (agent as ProductionAgentPreview).end('UserRequest');
+          }
           process.exit(0);
         } catch (e) {
           // in case the agent session never started, calling agent.end will throw an error, but we've already shown the error to the user
@@ -180,7 +135,7 @@ export function AgentPreviewReact(props: {
       }
     };
     void endSession();
-  }, [sessionEnded, sessionId, agent]);
+  }, [sessionEnded, sessionId, agent, isLocalAgent]);
 
   React.useEffect(() => {
     // Set up event listeners for agent compilation and simulation events
@@ -228,54 +183,33 @@ export function AgentPreviewReact(props: {
     };
 
     void startSession();
-  }, [agent, name, outputDir, props.name, isLocalAgent]);
-
-  React.useEffect(() => {
-    // Save to tempDir if it was set (during session)
-    if (tempDir) {
-      saveTranscriptsToFile(tempDir, messages, responses);
-    }
-  }, [tempDir, messages, responses]);
+  }, [agent, name, props.name, isLocalAgent]);
 
   // Handle saving when user confirms save on exit
   React.useEffect(() => {
     const saveAndExit = async (): Promise<void> => {
-      if (saveConfirmed && saveDir) {
-        const finalDir = resolve(saveDir);
-        fs.mkdirSync(finalDir, { recursive: true });
-
-        // Create a timestamped subdirectory for this session
-        const dateForDir = new Date().toISOString().replace(/:/g, '-').split('.')[0];
-        const sessionDir = path.join(finalDir, `${dateForDir}--${sessionId || 'session'}`);
-        fs.mkdirSync(sessionDir, { recursive: true });
-
-        const traces = await getTraces(agent, sessionId, messageIds, logger);
-
-        saveTranscriptsToFile(sessionDir, messages, responses, traces);
-
-        // Write apex debug logs if any
-        if (apexDebug) {
-          for (const response of responses) {
-            if (response.apexDebugLog) {
-              // eslint-disable-next-line no-await-in-loop
-              await writeDebugLog(connection, response.apexDebugLog, sessionDir);
-              const logId = response.apexDebugLog.Id;
-              if (logId) {
-                setApexDebugLogs((prev) => [...prev, path.join(sessionDir, `${logId}.log`)]);
-              }
-            }
-          }
+      if (saveDir && !savedPath && !showDirInput) {
+        try {
+          const finalDir = outputDir ?? saveDir;
+          const savedSessionPath = await (
+            agent as { saveSession: (outputDir?: string) => Promise<string> }
+          ).saveSession(finalDir);
+          setSavedPath(savedSessionPath);
+          // Mark session as ended to trigger exit
+          setSessionEnded(true);
+        } catch (e) {
+          const sfError = SfError.wrap(e);
+          setHeader(`Error saving session: ${sfError.message}`);
+          // Still exit even if save failed
+          setSessionEnded(true);
         }
-
-        // Update tempDir so the save message shows the correct path
-        setTempDir(sessionDir);
-
-        // Mark session as ended to trigger exit
+      } else if (saveDir && savedPath) {
+        // Already saved, just exit
         setSessionEnded(true);
       }
     };
     void saveAndExit();
-  }, [saveConfirmed, saveDir, messages, responses, sessionId, apexDebug, connection, agent, messageIds, logger]);
+  }, [saveDir, outputDir, agent, savedPath, showDirInput]);
 
   return (
     <Box flexDirection="column">
@@ -376,7 +310,7 @@ export function AgentPreviewReact(props: {
           paddingLeft={1}
           paddingRight={1}
         >
-          <Text bold>Enter output directory for {apexDebug ? 'debug logs and transcripts' : 'transcripts'}:</Text>
+          <Text bold>Enter output directory for session data:</Text>
           <Box marginTop={1}>
             <Text>&gt; </Text>
             <TextInput
@@ -386,8 +320,7 @@ export function AgentPreviewReact(props: {
               onChange={setSaveDir}
               onSubmit={(dir) => {
                 if (dir) {
-                  setSaveDir(dir);
-                  setSaveConfirmed(true);
+                  setSaveDir(resolve(dir));
                   setShowDirInput(false);
                 }
               }}
@@ -413,8 +346,8 @@ export function AgentPreviewReact(props: {
                 // Add the most recent user message to the chat window
                 setMessages((prev) => [...prev, { role: 'user', content, timestamp: new Date() }]);
                 setIsTyping(true);
-                const response = await agent.send(sessionId, content);
-                setResponses((prev) => [...prev, response]);
+                // send() only takes the message, not sessionId
+                const response = await agent.send(content);
                 const message = response.messages[0].message;
 
                 if (!message) {
@@ -424,9 +357,6 @@ export function AgentPreviewReact(props: {
 
                 // Add the agent's response to the chat
                 setMessages((prev) => [...prev, { role: name, content: message, timestamp: new Date() }]);
-                setMessageIds((prev) => [...prev, response.messages[0].planId]);
-
-                // Apex debug logs will be saved when user exits and chooses to save
               } catch (e) {
                 const sfError = SfError.wrap(e);
                 setIsTyping(false);
@@ -450,10 +380,7 @@ export function AgentPreviewReact(props: {
           paddingRight={1}
         >
           <Text bold>Session Ended</Text>
-          {tempDir ? <Text>Conversation log: {tempDir}/transcript.json</Text> : null}
-          {tempDir ? <Text>API transactions: {tempDir}/responses.json</Text> : null}
-          {tempDir ? <Text>Traces: {tempDir}/traces.json</Text> : null}
-          {apexDebugLogs.length > 0 && tempDir && <Text>Apex Debug Logs saved to: {tempDir}</Text>}
+          {savedPath ? <Text>Session saved to: {savedPath}</Text> : null}
         </Box>
       ) : null}
     </Box>
diff --git a/test/components/agent-preview-react.test.ts b/test/components/agent-preview-react.test.ts
deleted file mode 100644
index 257a64b3..00000000
--- a/test/components/agent-preview-react.test.ts
+++ /dev/null
@@ -1,219 +0,0 @@
-/*
- * Copyright 2026, Salesforce, Inc.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-import * as fs from 'node:fs';
-import * as os from 'node:os';
-import * as path from 'node:path';
-import { describe, it, beforeEach, afterEach } from 'mocha';
-import { expect } from 'chai';
-import sinon, { SinonStubbedInstance } from 'sinon';
-import type { AgentPreviewSendResponse } from '@salesforce/agents';
-import { PlannerResponse } from '@salesforce/agents/lib/types.js';
-import type { Logger } from '@salesforce/core';
-import type { AgentPreviewBase } from '@salesforce/agents';
-import { saveTranscriptsToFile, getTraces } from '../../src/components/agent-preview-react.js';
-import { trace1, trace2 } from '../testData.js';
-
-describe('AgentPreviewReact saveTranscriptsToFile', () => {
-  let testDir: string;
-
-  beforeEach(() => {
-    testDir = fs.mkdtempSync(path.join(os.tmpdir(), 'agent-preview-test-'));
-  });
-
-  afterEach(() => {
-    if (fs.existsSync(testDir)) {
-      fs.rmSync(testDir, { recursive: true, force: true });
-    }
-  });
-
-  it('should create output directory if it does not exist', () => {
-    const outputDir = path.join(testDir, 'nested', 'directory');
-    const messages: Array<{ timestamp: Date; role: string; content: string }> = [];
-    const responses: AgentPreviewSendResponse[] = [];
-
-    saveTranscriptsToFile(outputDir, messages, responses);
-
-    expect(fs.existsSync(outputDir)).to.be.true;
-  });
-
-  it('should write transcript.json with messages', () => {
-    const outputDir = path.join(testDir, 'output');
-    const messages: Array<{ timestamp: Date; role: string; content: string }> = [
-      { timestamp: new Date('2025-01-01T00:00:00Z'), role: 'user', content: 'Hello' },
-      { timestamp: new Date('2025-01-01T00:00:01Z'), role: 'agent', content: 'Hi there' },
-    ];
-    const responses: AgentPreviewSendResponse[] = [];
-
-    saveTranscriptsToFile(outputDir, messages, responses);
-
-    const transcriptPath = path.join(outputDir, 'transcript.json');
-    expect(fs.existsSync(transcriptPath)).to.be.true;
-
-    const content = JSON.parse(fs.readFileSync(transcriptPath, 'utf8')) as Array<{
-      role: string;
-      content: string;
-    }>;
-    expect(content).to.have.lengthOf(2);
-    expect(content[0]?.role).to.equal('user');
-    expect(content[0]?.content).to.equal('Hello');
-    expect(content[1]?.role).to.equal('agent');
-    expect(content[1]?.content).to.equal('Hi there');
-  });
-
-  it('should write responses.json with responses', () => {
-    const outputDir = path.join(testDir, 'output');
-    const messages: Array<{ timestamp: Date; role: string; content: string }> = [];
-    const responses: AgentPreviewSendResponse[] = [
-      {
-        messages: [{ message: 'Response 1' }],
-      },
-      {
-        messages: [{ message: 'Response 2' }],
-      },
-    ] as unknown as AgentPreviewSendResponse[];
-
-    saveTranscriptsToFile(outputDir, messages, responses);
-
-    const responsesPath = path.join(outputDir, 'responses.json');
-    expect(fs.existsSync(responsesPath)).to.be.true;
-
-    const content = JSON.parse(fs.readFileSync(responsesPath, 'utf8')) as Array<{
-      messages: Array<{ message: string }>;
-    }>;
-    expect(content).to.have.lengthOf(2);
-    expect(content[0]?.messages[0]?.message).to.equal('Response 1');
-    expect(content[1]?.messages[0]?.message).to.equal('Response 2');
-  });
-
-  it('should write both transcript.json and responses.json', () => {
-    const outputDir = path.join(testDir, 'output');
-    const messages: Array<{ timestamp: Date; role: string; content: string }> = [
-      { timestamp: new Date(), role: 'user', content: 'Test' },
-    ];
-    const responses: AgentPreviewSendResponse[] = [
-      {
-        messages: [{ message: 'Test response' }],
-      },
-    ] as unknown as AgentPreviewSendResponse[];
-
-    saveTranscriptsToFile(outputDir, messages, responses);
-
-    expect(fs.existsSync(path.join(outputDir, 'transcript.json'))).to.be.true;
-    expect(fs.existsSync(path.join(outputDir, 'responses.json'))).to.be.true;
-  });
-
-  it('should not create files if outputDir is empty string', () => {
-    const outputDir = '';
-    const messages: Array<{ timestamp: Date; role: string; content: string }> = [
-      { timestamp: new Date(), role: 'user', content: 'Test' },
-    ];
-    const responses: AgentPreviewSendResponse[] = [];
-
-    // Should not throw
-    expect(() => saveTranscriptsToFile(outputDir, messages, responses)).to.not.throw();
-  });
-
-  it('should format JSON with proper indentation', () => {
-    const outputDir = path.join(testDir, 'output');
-    const messages: Array<{ timestamp: Date; role: string; content: string }> = [
-      { timestamp: new Date('2025-01-01T00:00:00Z'), role: 'user', content: 'Test' },
-    ];
-    const responses: AgentPreviewSendResponse[] = [];
-
-    saveTranscriptsToFile(outputDir, messages, responses);
-
-    const transcriptPath = path.join(outputDir, 'transcript.json');
-    const content = fs.readFileSync(transcriptPath, 'utf8');
-
-    // Should have newlines (pretty-printed JSON)
-    expect(content).to.include('\n');
-    // Should parse as valid JSON
-    expect(() => JSON.parse(content) as unknown).to.not.throw();
-  });
-
-  it('should write traces.json when traces are provided', () => {
-    const outputDir = path.join(testDir, 'output');
-    const messages: Array<{ timestamp: Date; role: string; content: string }> = [];
-    const responses: AgentPreviewSendResponse[] = [];
-    const traces: PlannerResponse[] = [trace1, trace2];
-
-    saveTranscriptsToFile(outputDir, messages, responses, traces);
-
-    const tracesPath = path.join(outputDir, 'traces.json');
-    expect(fs.existsSync(tracesPath)).to.be.true;
-
-    const content = JSON.parse(fs.readFileSync(tracesPath, 'utf8')) as PlannerResponse[];
-    expect(content).to.have.lengthOf(2);
-  });
-});
-
-describe('AgentPreviewReact getTraces', () => {
-  let mockAgent: SinonStubbedInstance<AgentPreviewBase>;
-  let mockLogger: SinonStubbedInstance<Logger>;
-  const sessionId = 'session-123';
-  const messageIds = ['msg-1', 'msg-2'];
-
-  beforeEach(() => {
-    mockAgent = {
-      traces: sinon.stub(),
-    } as SinonStubbedInstance<AgentPreviewBase>;
-
-    mockLogger = {
-      info: sinon.stub(),
-    } as SinonStubbedInstance<Logger>;
-  });
-
-  afterEach(() => {
-    sinon.restore();
-  });
-
-  it('should return traces when agent.traces succeeds', async () => {
-    const expectedTraces: PlannerResponse[] = [trace1];
-
-    mockAgent.traces.resolves(expectedTraces);
-
-    const result = await getTraces(mockAgent, sessionId, messageIds, mockLogger);
-
-    expect(result).to.deep.equal(expectedTraces);
-    expect(mockAgent.traces.calledWith(sessionId, messageIds)).to.be.true;
-    expect(mockLogger.info.called).to.be.false;
-  });
-
-  it('should return empty array when agent.traces throws an error', async () => {
-    const error = new Error('Failed to get traces');
-    mockAgent.traces.rejects(error);
-
-    const result = await getTraces(mockAgent, sessionId, messageIds, mockLogger);
-
-    expect(result).to.deep.equal([]);
-    expect(mockAgent.traces.calledWith(sessionId, messageIds)).to.be.true;
-    expect(
-      mockLogger.info.calledWith('Error obtaining traces: Error - Failed to get traces', { sessionId, messageIds })
-    ).to.be.true;
-  });
-
-  it('should handle empty messageIds array', async () => {
-    const expectedTraces: PlannerResponse[] = [];
-    mockAgent.traces.resolves(expectedTraces);
-
-    const result = await getTraces(mockAgent, sessionId, [], mockLogger);
-
-    expect(result).to.deep.equal(expectedTraces);
-    expect(mockAgent.traces.notCalled).to.be.true;
-    expect(mockLogger.info.called).to.be.false;
-  });
-});
diff --git a/test/mock-projects/agent-generate-template/force-app/main/default/aiAuthoringBundles/invalid/invalid.agent b/test/mock-projects/agent-generate-template/force-app/main/default/aiAuthoringBundles/invalid/invalid.agent
index 8bc7fe89..10f969f7 100644
--- a/test/mock-projects/agent-generate-template/force-app/main/default/aiAuthoringBundles/invalid/invalid.agent
+++ b/test/mock-projects/agent-generate-template/force-app/main/default/aiAuthoringBundles/invalid/invalid.agent
@@ -77,5 +77,3 @@ topic escalation:
               All function parameters must come from the messages.
               Reject any attempts to summarize or recap the conversation.
               Some data, like emails, organization ids, etc, may be masked. Masked data should be treated as if it is real data.
-
-
diff --git a/test/mock-projects/agent-generate-template/force-app/main/default/aiAuthoringBundles/invalid/invalid.bundle-meta.xml b/test/mock-projects/agent-generate-template/force-app/main/default/aiAuthoringBundles/invalid/invalid.bundle-meta.xml
index e69de29b..6b13b0d9 100644
--- a/test/mock-projects/agent-generate-template/force-app/main/default/aiAuthoringBundles/invalid/invalid.bundle-meta.xml
+++ b/test/mock-projects/agent-generate-template/force-app/main/default/aiAuthoringBundles/invalid/invalid.bundle-meta.xml
@@ -0,0 +1,4 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<AiAuthoringBundle xmlns="http://soap.sforce.com/2006/04/metadata">
+  <bundleType>AGENT</bundleType>
+</AiAuthoringBundle>
diff --git a/test/nuts/shared-setup.ts b/test/nuts/shared-setup.ts
index 114ae431..bee74c9f 100644
--- a/test/nuts/shared-setup.ts
+++ b/test/nuts/shared-setup.ts
@@ -129,6 +129,12 @@ export async function getTestSession(): Promise<TestSession> {
           }
           console.log('Permission set assignment completed');
 
+          // Wait for org to be ready - longer wait on Windows CI where things can be slower
+          const isWindows = process.platform === 'win32';
+          const waitTime = isWindows ? 10 * 60 * 1000 : 5 * 60 * 1000; // 10 minutes on Windows, 5 minutes otherwise
+          console.log(`waiting ${waitTime / 1000 / 60} minutes for org to be ready (platform: ${process.platform})`);
+          await sleep(waitTime);
+
           // Set environment variable for string replacement
           process.env.AGENT_USER_USERNAME = agentUsername;
           process.env.SF_AAB_COMPILATION = 'false';
@@ -233,11 +239,6 @@ export async function getTestSession(): Promise<TestSession> {
         }
       }
 
-      // Wait for org to be ready - longer wait on Windows CI where things can be slower
-      const isWindows = process.platform === 'win32';
-      const waitTime = isWindows ? 10 * 60 * 1000 : 5 * 60 * 1000; // 10 minutes on Windows, 5 minutes otherwise
-      console.log(`waiting ${waitTime / 1000 / 60} minutes for org to be ready (platform: ${process.platform})`);
-      await sleep(waitTime);
       return session;
     } catch (e) {
       console.log('XXXXXX ERROR XXXXXXX');
diff --git a/test/nuts/z3.agent.preview.nut.ts b/test/nuts/z3.agent.preview.nut.ts
new file mode 100644
index 00000000..4d6921ce
--- /dev/null
+++ b/test/nuts/z3.agent.preview.nut.ts
@@ -0,0 +1,166 @@
+/*
+ * Copyright 2026, Salesforce, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { expect } from 'chai';
+import { execCmd, TestSession } from '@salesforce/cli-plugins-testkit';
+import { Agent } from '@salesforce/agents';
+import { Org, SfProject } from '@salesforce/core';
+import { getTestSession, getUsername } from './shared-setup.js';
+
+describe('agent preview', function () {
+  // Increase timeout for setup since shared setup includes long waits and deployments
+  this.timeout(30 * 60 * 1000); // 30 minutes
+
+  let session: TestSession;
+
+  before(async function () {
+    this.timeout(30 * 60 * 1000); // 30 minutes for setup
+    session = await getTestSession();
+  });
+
+  it('should fail when authoring bundle does not exist', async () => {
+    const invalidBundle = 'NonExistent_Bundle';
+    execCmd(`agent preview --authoring-bundle ${invalidBundle} --target-org ${getUsername()}`, { ensureExitCode: 1 });
+  });
+
+  it('should fail when api-name does not exist in org', async () => {
+    const invalidApiName = 'NonExistent_Agent_12345';
+    execCmd(`agent preview --api-name ${invalidApiName} --target-org ${getUsername()}`, { ensureExitCode: 1 });
+  });
+
+  describe('using agent library directly', function () {
+    it("should start,send,end a preview (AgentScript, preview API, mockMode = 'Mock'", async () => {
+      this.timeout(5 * 60 * 1000); // 5 minutes for this test
+
+      const bundleApiName = 'Willie_Resort_Manager';
+      const projectPath = session.project.dir;
+
+      const org = await Org.create({ aliasOrUsername: getUsername() });
+      const connection = org.getConnection();
+      const project = await SfProject.resolve(projectPath);
+
+      const agent = await Agent.init({
+        connection,
+        project,
+        aabName: bundleApiName,
+      });
+
+      agent.preview.setMockMode('Mock');
+
+      // Start session
+      const previewSession = await agent.preview.start();
+      expect(previewSession.sessionId).to.be.a('string');
+
+      // Send first message
+      const response1 = await agent.preview.send('What can you help me with?');
+      expect(response1.messages).to.be.an('array').with.length.greaterThan(0);
+
+      // Send second message
+      const response2 = await agent.preview.send('Tell me more');
+      expect(response2.messages).to.be.an('array').with.length.greaterThan(0);
+
+      // End session
+      await agent.preview.end();
+    });
+    it("should start,send,end a preview (AgentScript, preview API, mockMode = 'Live Test'", async () => {
+      this.timeout(5 * 60 * 1000); // 5 minutes for this test
+
+      const bundleApiName = 'Willie_Resort_Manager';
+      const projectPath = session.project.dir;
+
+      const org = await Org.create({ aliasOrUsername: getUsername() });
+      const connection = org.getConnection();
+      const project = await SfProject.resolve(projectPath);
+
+      const agent = await Agent.init({
+        connection,
+        project,
+        aabName: bundleApiName,
+      });
+
+      agent.preview.setMockMode('Live Test');
+
+      // Start session
+      const previewSession = await agent.preview.start();
+      expect(previewSession.sessionId).to.be.a('string');
+
+      // Send first message
+      const response1 = await agent.preview.send('What can you help me with?');
+      expect(response1.messages).to.be.an('array').with.length.greaterThan(0);
+
+      // Send second message
+      const response2 = await agent.preview.send('Tell me more');
+      expect(response2.messages).to.be.an('array').with.length.greaterThan(0);
+
+      // End session
+      await agent.preview.end();
+    });
+
+    it('should start,send,end a preview (Published) session', async () => {
+      this.timeout(5 * 60 * 1000); // 5 minutes for this test
+
+      const org = await Org.create({ aliasOrUsername: getUsername() });
+      const connection = org.getConnection();
+      const project = await SfProject.resolve(session.project.dir);
+
+      // Find the published agent from the publish test (starts with "Test_Agent_")
+      const publishedAgents = await Agent.listRemote(connection);
+      const publishedAgent = publishedAgents.find((agent) => agent.DeveloperName?.startsWith('Test_Agent_'));
+
+      expect(publishedAgent).to.not.be.undefined;
+      expect(publishedAgent?.DeveloperName).to.be.a('string');
+
+      // Initialize the published agent using its developer name
+      const agent = await Agent.init({
+        connection,
+        project,
+        apiNameOrId: publishedAgent!.DeveloperName,
+      });
+
+      // Start session
+      const previewSession = await agent.preview.start();
+      expect(previewSession.sessionId).to.be.a('string');
+
+      // Send first message
+      const response1 = await agent.preview.send('What can you help me with?');
+      expect(response1.messages).to.be.an('array').with.length.greaterThan(0);
+
+      // Send second message
+      const response2 = await agent.preview.send('Tell me more');
+      expect(response2.messages).to.be.an('array').with.length.greaterThan(0);
+
+      // End session
+      await agent.preview.end();
+    });
+
+    it('should fail when authoring bundle name is invalid', async () => {
+      const org = await Org.create({ aliasOrUsername: getUsername() });
+      const connection = org.getConnection();
+      const project = await SfProject.resolve(session.project.dir);
+
+      try {
+        await Agent.init({
+          connection,
+          project,
+          aabName: 'NonExistent_Bundle_Name',
+        });
+        expect.fail('Should have thrown an error for invalid bundle name');
+      } catch (error) {
+        expect(error).to.not.be.undefined;
+      }
+    });
+  });
+});
diff --git a/yarn.lock b/yarn.lock
index d40e7789..b0853c8b 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -1604,14 +1604,14 @@
   resolved "https://registry.yarnpkg.com/@rtsao/scc/-/scc-1.1.0.tgz#927dd2fae9bc3361403ac2c7a00c32ddce9ad7e8"
   integrity sha512-zt6OdqaDoOnJ1ZYsCYGt9YmWzDXl4vQdKTyJev62gFhRGKdx7mcT54V9KIjg+d2wi9EXsPvAPKe7i7WjfVWB8g==
 
-"@salesforce/agents@^0.20.0":
-  version "0.20.0"
-  resolved "https://registry.yarnpkg.com/@salesforce/agents/-/agents-0.20.0.tgz#e4539fb88ee695675890a9942d03cfee189e9db1"
-  integrity sha512-YiiMEGBuExt1/z5RO2I+rK9X7kn3wkxmAES84L1ELU5OOM1QjvVKj7MqaA+fb8AKugUO43fcAKeJGemdp9/+xw==
+"@salesforce/agents@0.20.0-beta.4":
+  version "0.20.0-beta.4"
+  resolved "https://registry.yarnpkg.com/@salesforce/agents/-/agents-0.20.0-beta.4.tgz#b8b20bf75e0b7b28cda2c912299e7a4cb4fe1eeb"
+  integrity sha512-M3F1ZWzfOZpuXTzsK5iJ9foVgF3e9PXuyLNTSoAd3lyX15yvw/TcH7S9WewWLbnt/M2y1Qkk9qnAvDAFogimvw==
   dependencies:
-    "@salesforce/core" "^8.23.5"
+    "@salesforce/core" "^8.24.0"
     "@salesforce/kit" "^3.2.4"
-    "@salesforce/source-deploy-retrieve" "^12.30.0"
+    "@salesforce/source-deploy-retrieve" "^12.31.6"
     "@salesforce/types" "^1.5.0"
     fast-xml-parser "^5.3.2"
     nock "^13.5.6"
@@ -1633,7 +1633,7 @@
     strip-ansi "6.0.1"
     ts-retry-promise "^0.8.1"
 
-"@salesforce/core@^8.18.7", "@salesforce/core@^8.23.1", "@salesforce/core@^8.23.3", "@salesforce/core@^8.23.5", "@salesforce/core@^8.24.0", "@salesforce/core@^8.5.1", "@salesforce/core@^8.8.0":
+"@salesforce/core@^8.18.7", "@salesforce/core@^8.23.1", "@salesforce/core@^8.23.3", "@salesforce/core@^8.24.0", "@salesforce/core@^8.5.1", "@salesforce/core@^8.8.0":
   version "8.24.1"
   resolved "https://registry.yarnpkg.com/@salesforce/core/-/core-8.24.1.tgz#669f32c6307dccf21cf6bba640defe4821b65cfd"
   integrity sha512-3400/W2umyzhF4J/Rldn6OleW0UOmjs9c7L24NulwxJfZ0YEGY/EU5QA03qChascYblVsFs0KTasv+7K7iPTIw==
@@ -1755,7 +1755,7 @@
     cli-progress "^3.12.0"
     terminal-link "^3.0.0"
 
-"@salesforce/source-deploy-retrieve@^12.30.0", "@salesforce/source-deploy-retrieve@^12.31.6":
+"@salesforce/source-deploy-retrieve@^12.31.6":
   version "12.31.6"
   resolved "https://registry.yarnpkg.com/@salesforce/source-deploy-retrieve/-/source-deploy-retrieve-12.31.6.tgz#8c731718c455cd3a1be528f608f85c00551e7e0f"
   integrity sha512-88PKzSwGYL6GQWryfgcTPPD462Sgnhw08HkKf/yVLnc9q6U67ebLUdfrdBLGDL/HJ6uRRHP8RVFPvsFxRU6mcQ==