fix(coding-agent): sync agent state after newSession setup callback

Move setup callback handling from interactive/rpc modes into AgentSession.newSession().
After setup() runs, sync agent state via replaceMessages() so the LLM has context
and the UI renders the messages properly.

fixes #968

packages/coding-agent/CHANGELOG.md

@@ -5,6 +5,7 @@
 ### Fixed
 
 - Git extension updates now handle force-pushed remotes gracefully instead of failing ([#961](https://github.com/badlogic/pi-mono/pull/961) by [@aliou](https://github.com/aliou))
+- Extension `ctx.newSession({ setup })` now properly syncs agent state and renders messages after setup callback runs ([#968](https://github.com/badlogic/pi-mono/issues/968))
 
 ## [0.50.0] - 2026-01-26
 

packages/coding-agent/src/core/agent-session.ts

@@ -65,7 +65,7 @@ import type { BashExecutionMessage, CustomMessage } from "./messages.js";
 import type { ModelRegistry } from "./model-registry.js";
 import { expandPromptTemplate, type PromptTemplate } from "./prompt-templates.js";
 import type { ResourceLoader } from "./resource-loader.js";
-import type { BranchSummaryEntry, CompactionEntry, NewSessionOptions, SessionManager } from "./session-manager.js";
+import type { BranchSummaryEntry, CompactionEntry, SessionManager } from "./session-manager.js";
 import type { SettingsManager } from "./settings-manager.js";
 import { buildSystemPrompt } from "./system-prompt.js";
 import type { BashOperations } from "./tools/bash.js";
@@ -1049,10 +1049,14 @@ export class AgentSession {
 	 * Start a new session, optionally with initial messages and parent tracking.
 	 * Clears all messages and starts a new session.
 	 * Listeners are preserved and will continue receiving events.
-	 * @param options - Optional initial messages and parent session path
+	 * @param options.parentSession - Optional parent session path for tracking
+	 * @param options.setup - Optional callback to initialize session (e.g., append messages)
 	 * @returns true if completed, false if cancelled by extension
 	 */
-	async newSession(options?: NewSessionOptions): Promise<boolean> {
+	async newSession(options?: {
+		parentSession?: string;
+		setup?: (sessionManager: SessionManager) => Promise<void>;
+	}): Promise<boolean> {
 		const previousSessionFile = this.sessionFile;
 
 		// Emit session_before_switch event with reason "new" (can be cancelled)
@@ -1070,11 +1074,20 @@ export class AgentSession {
 		this._disconnectFromAgent();
 		await this.abort();
 		this.agent.reset();
-		this.sessionManager.newSession(options);
+		this.sessionManager.newSession({ parentSession: options?.parentSession });
 		this.agent.sessionId = this.sessionManager.getSessionId();
 		this._steeringMessages = [];
 		this._followUpMessages = [];
 		this._pendingNextTurnMessages = [];
+
+		// Run setup callback if provided (e.g., to append initial messages)
+		if (options?.setup) {
+			await options.setup(this.sessionManager);
+			// Sync agent state with session manager after setup
+			const sessionContext = this.sessionManager.buildSessionContext();
+			this.agent.replaceMessages(sessionContext.messages);
+		}
+
 		this._reconnectToAgent();
 
 		// Emit session_switch event with reason "new" to extensions

packages/coding-agent/src/modes/interactive/interactive-mode.ts

@@ -995,15 +995,13 @@ export class InteractiveMode {
 					}
 					this.statusContainer.clear();
 
-					const success = await this.session.newSession({ parentSession: options?.parentSession });
+					// Delegate to AgentSession (handles setup + agent state sync)
+					const success = await this.session.newSession(options);
 					if (!success) {
 						return { cancelled: true };
 					}
 
-					if (options?.setup) {
+					// Clear UI state
-						await options.setup(this.sessionManager);
-					}
-
 					this.chatContainer.clear();
 					this.pendingMessagesContainer.clear();
 					this.compactionQueuedMessages = [];
@@ -1011,8 +1009,8 @@ export class InteractiveMode {
 					this.streamingMessage = undefined;
 					this.pendingTools.clear();
 
-					this.chatContainer.addChild(new Spacer(1));
+					// Render any messages added via setup, or show empty session
-					this.chatContainer.addChild(new Text(`${theme.fg("accent", "✓ New session started")}`, 1, 1));
+					this.renderInitialMessages();
 					this.ui.requestRender();
 
 					return { cancelled: false };

packages/coding-agent/src/modes/rpc/rpc-mode.ts

@@ -261,10 +261,8 @@ export async function runRpcMode(session: AgentSession): Promise<never> {
 			commandContextActions: {
 				waitForIdle: () => session.agent.waitForIdle(),
 				newSession: async (options) => {
-					const success = await session.newSession({ parentSession: options?.parentSession });
+					// Delegate to AgentSession (handles setup + agent state sync)
-					if (success && options?.setup) {
+					const success = await session.newSession(options);
-						await options.setup(session.sessionManager);
-					}
 					return { cancelled: !success };
 				},
 				fork: async (entryId) => {