From fc43e2d6edf28aa8861b8c52d935ee54d4c8a99b Mon Sep 17 00:00:00 2001 From: Kian Jones <11655409+kianjones9@users.noreply.github.com> Date: Thu, 18 Dec 2025 15:50:19 -0500 Subject: [PATCH] fix: update model in ui when switched in the ADE (#297) --- src/cli/App.tsx | 56 +++++++++++++++++++++++++++++++++++++++ src/cli/helpers/stream.ts | 25 ++++++++++++++++- 2 files changed, 80 insertions(+), 1 deletion(-) diff --git a/src/cli/App.tsx b/src/cli/App.tsx index 9fb7f4b..590b4dc 100644 --- a/src/cli/App.tsx +++ b/src/cli/App.tsx @@ -433,6 +433,10 @@ export default function App({ | null >(null); const [llmConfig, setLlmConfig] = useState(null); + const llmConfigRef = useRef(llmConfig); + useEffect(() => { + llmConfigRef.current = llmConfig; + }, [llmConfig]); const [currentModelId, setCurrentModelId] = useState(null); const [agentName, setAgentName] = useState(null); const [agentDescription, setAgentDescription] = useState(null); @@ -831,12 +835,64 @@ export default function App({ agentIdRef.current, currentInput, ); + + // Define callback to sync agent state on first message chunk + // This ensures the UI shows the correct model as early as possible + const syncAgentState = async () => { + try { + const client = await getClient(); + const agent = await client.agents.retrieve(agentIdRef.current); + + // Check if the model has changed by comparing llm_config + const currentModel = llmConfigRef.current?.model; + const currentEndpoint = llmConfigRef.current?.model_endpoint_type; + const agentModel = agent.llm_config.model; + const agentEndpoint = agent.llm_config.model_endpoint_type; + + if ( + currentModel !== agentModel || + currentEndpoint !== agentEndpoint + ) { + // Model has changed - update local state + setLlmConfig(agent.llm_config); + + // Derive model ID from llm_config for ModelSelector + // Try to find matching model by handle in models.json + const { getModelInfo } = await import("../agent/model"); + const agentModelHandle = + agent.llm_config.model_endpoint_type && agent.llm_config.model + ? `${agent.llm_config.model_endpoint_type}/${agent.llm_config.model}` + : agent.llm_config.model; + + const modelInfo = getModelInfo(agentModelHandle || ""); + if (modelInfo) { + setCurrentModelId(modelInfo.id); + } else { + // Model not in models.json (e.g., BYOK model) - use handle as ID + setCurrentModelId(agentModelHandle || null); + } + + // Also update agent state if other fields changed + setAgentName(agent.name); + setAgentDescription(agent.description ?? null); + const lastRunCompletion = ( + agent as { last_run_completion?: string } + ).last_run_completion; + setAgentLastRunAt(lastRunCompletion ?? null); + } + } catch (error) { + // Silently fail - don't interrupt the conversation flow + console.error("Failed to sync agent state:", error); + } + }; + const { stopReason, approval, approvals, apiDurationMs, lastRunId } = await drainStreamWithResume( stream, buffersRef.current, refreshDerivedThrottled, abortControllerRef.current?.signal, + syncAgentState, ); // Track API duration diff --git a/src/cli/helpers/stream.ts b/src/cli/helpers/stream.ts index f5462b4..0b946ae 100644 --- a/src/cli/helpers/stream.ts +++ b/src/cli/helpers/stream.ts @@ -31,6 +31,7 @@ export async function drainStream( buffers: ReturnType, refresh: () => void, abortSignal?: AbortSignal, + onFirstMessage?: () => void, ): Promise { const startTime = performance.now(); @@ -47,6 +48,7 @@ export async function drainStream( let stopReason: StopReasonType | null = null; let lastRunId: string | null = null; let lastSeqId: number | null = null; + let hasCalledFirstMessage = false; for await (const chunk of stream) { // console.log("chunk", chunk); @@ -71,6 +73,18 @@ export async function drainStream( if (chunk.message_type === "ping") continue; + // Call onFirstMessage callback on the first agent response chunk + if ( + !hasCalledFirstMessage && + onFirstMessage && + (chunk.message_type === "reasoning_message" || + chunk.message_type === "assistant_message") + ) { + hasCalledFirstMessage = true; + // Call async in background - don't block stream processing + queueMicrotask(() => onFirstMessage()); + } + // Remove tool from pending approvals when it completes (server-side execution finished) // This means the tool was executed server-side and doesn't need approval if (chunk.message_type === "tool_return_message") { @@ -218,6 +232,7 @@ export async function drainStream( * @param buffers - Buffer to accumulate chunks * @param refresh - Callback to refresh UI * @param abortSignal - Optional abort signal for cancellation + * @param onFirstMessage - Optional callback to invoke on first message chunk * @returns Result with stop_reason, approval info, and timing */ export async function drainStreamWithResume( @@ -225,11 +240,18 @@ export async function drainStreamWithResume( buffers: ReturnType, refresh: () => void, abortSignal?: AbortSignal, + onFirstMessage?: () => void, ): Promise { const overallStartTime = performance.now(); // Attempt initial drain - let result = await drainStream(stream, buffers, refresh, abortSignal); + let result = await drainStream( + stream, + buffers, + refresh, + abortSignal, + onFirstMessage, + ); // If stream ended without proper stop_reason and we have resume info, try once to reconnect if ( @@ -247,6 +269,7 @@ export async function drainStreamWithResume( }); // Continue draining from where we left off + // Note: Don't pass onFirstMessage again - already called in initial drain const resumeResult = await drainStream( resumeStream, buffers,