diff --git a/src/plugin/chat-message.ts b/src/plugin/chat-message.ts
index ca8577d85..494530cf5 100644
--- a/src/plugin/chat-message.ts
+++ b/src/plugin/chat-message.ts
@@ -1,7 +1,6 @@
 import type { OhMyOpenCodeConfig } from "../config"
 import type { CreatedHooks } from "../create-hooks"
 
-import { getMainSessionID, setSessionAgent, subagentSessions } from "../features/claude-code-session-state"
 import { parseRalphLoopArguments } from "../hooks/ralph-loop/command-arguments"
 import {
   isModelCacheAvailable,
@@ -11,6 +10,7 @@ import {
 } from "../shared"
 import { getAgentConfigKey } from "../shared/agent-display-names"
 import { getSessionModel, setSessionModel } from "../shared/session-model-state"
+import { getMainSessionID, setSessionAgent, updateSessionAgent, subagentSessions } from "../features/claude-code-session-state"
 import { NATIVE_LOOP_TRIGGERED_FLAG } from "./command-execute-before"
 import type { PluginContext } from "./types"
 import { applyUltraworkModelOverrideOnMessage } from "./ultrawork-model-override"
@@ -203,7 +203,7 @@ export function createChatMessageHandler(args: {
     }
 
     if (input.agent) {
-      setSessionAgent(input.sessionID, input.agent)
+      updateSessionAgent(input.sessionID, input.agent)
     }
 
     const isFirstMessage = firstMessageVariantGate.shouldOverride(input.sessionID)