@superblocksteam/vite-plugin-file-sync 2.0.43-next.12 → 2.0.43-next.14

package/dist/ai-service/state-machine/handlers/llm-generating.js

@@ -1,5 +1,5 @@
 import { AiMode, } from "@superblocksteam/library-shared/types";
-import { hasToolCall, smoothStream, } from "ai";
+import { smoothStream, } from "ai";
 import { getErrorMeta, getLogger } from "../../../util/logger.js";
 import { buildBaseSystemPrompt } from "../../agent/prompts/build-base-system-prompt.js";
 import { getToolCallArguments } from "../../agent/tool-message-utils.js";
@@ -95,244 +95,285 @@ export const doLLMGenerating = (clark, services) => {
                 const { abortController } = clark.context;
                 const contextId = getContextId(clark, services);
                 const contextOptions = clark.context.llmConfig?.contextOptions;
-                const context = await services.contextManager.getContext(contextId, contextOptions);
-                const mode = clark.context.currentMode;
-                const userApprovedPlan = clark.context.userApprovedPlan;
-                await chatSessionStore.getMessages();
-                const latestSummary = await chatSessionStore.getLatestSummary();
-                if (latestSummary) {
-                    context.seedIfEmpty({
-                        role: "user",
-                        content: `Reference the context from the previous exchange with the user, if relevant:\n\n${formatSummaryForAgents(latestSummary)}`,
-                    });
-                }
-                const systemPrompt = buildBaseSystemPrompt(mode);
-                const system = {
-                    role: "system",
-                    content: systemPrompt,
-                };
-                context.setSystemPrompt(system);
-                const prompt = buildUserMessage(userPrompt, promptContext, mode, userApprovedPlan);
-                context.startTurn(prompt);
-                const messages = context.getMessages();
-                const model = services.llmProvider.modelForClassification("broad_edit");
-                // Process LLM configuration up front so we can log it once at the top
-                const llmConfig = clark.context.llmConfig;
-                const disabledTools = clark.context.llmConfig?.disabledTools;
-                const { headers, thinkingEnabled, thinkingBudgetTokens, interleavedThinking, providerOptions, } = processLLMConfig(llmConfig, 5000, // default budget tokens
-                `LLM Generating (model=${model.modelId}, disabledTools=${disabledTools?.length ? disabledTools.join(",") : "none"})`);
-                const conversationId = Date.now();
-                const startTimestamp = new Date().toISOString();
-                const runTimestamp = clark.context.runTimestamp || startTimestamp;
-                // Initialize log content
-                const logRef = {
-                    content: `=== LLM CONVERSATION START [${conversationId}] ===\n`,
-                };
-                logRef.content += `Timestamp: ${startTimestamp}\n`;
-                logRef.content += `Model: ${model.modelId}\n`;
-                // Log AI provider config
-                logRef.content += `Provider Config: provider=${llmConfig?.provider || "anthropic (default)"}, thinking=${thinkingEnabled}, budget=${thinkingBudgetTokens}, interleaved=${interleavedThinking}`;
-                if (disabledTools?.length) {
-                    logRef.content += `, disabledTools=${disabledTools.join(",")}`;
-                }
-                if (headers && Object.keys(headers).length > 0) {
-                    logRef.content += `, headers=${JSON.stringify(headers)}`;
-                }
-                logRef.content += `--- INPUT MESSAGES ---\n`;
-                logRef.content += `[SYSTEM] ${systemPrompt}\n\n`;
-                logRef.content += `[${system.role.toUpperCase()}] ${system.content}\n\n`;
-                // Note: Input token count will be logged once LLM call begins
-                // buildTools now returns filtered tools based on mode via registry
-                const actualTools = await buildTools(clark, services, promptContext, logRef, disabledTools, mode);
-                logRef.content += `Mode: ${mode}, Available tools: ${Object.keys(actualTools).length}\n`;
-                logRef.content += `Tools: ${Object.keys(actualTools).join(", ")}\n\n`;
-                let stepCount = 0;
-                // Track cumulative token usage across all steps
-                let totalInputTokens = 0;
-                let totalOutputTokens = 0;
-                let totalCachedTokens = 0;
-                services.clarkProfiler.startLLMWaiting({
-                    messages: messages,
-                    model: model.modelId,
-                });
-                let firstTokenReceived = false;
-                let thinkingSpanActive = false;
-                let textSpanActive = false;
-                const allReasoningDeltas = [];
-                let currentStepReasoningDeltas = [];
-                let currentStepTextDeltas = [];
-                getLogger().info(`CURRENT_MODE: ${mode}`);
-                const build = tracedStreamText({
-                    abortSignal: abortController?.signal,
-                    model,
-                    providerOptions,
-                    headers,
-                    experimental_transform: [smoothStream({ chunking: "line" })],
-                    messages,
-                    tools: actualTools,
-                    prepareStep: async (step) => {
-                        context.startStep();
-                        const messages = context.getMessages();
-                        return { ...step, messages };
+                // Generate a unique owner ID for this LLM generation session
+                const ownerId = `llm-generating-${Date.now()}-${Math.random().toString(36).slice(2, 9)}`;
+                // Acquire exclusive context access
+                const contextHandle = await services.contextManager.acquireContext(contextId, ownerId, {
+                    contextOptions,
+                    acquireOptions: {
+                        waitTimeoutMs: 30000, // Wait up to 30 seconds for context
+                        lockTimeoutMs: 600000, // Hold lock for up to 10 minutes
+                        retryIntervalMs: 100,
                     },
-                    stopWhen: [
-                        hasToolCall("build_finalize"),
-                        hasToolSuccess("exitPlanMode"),
-                        hasToolSuccess("askMultiChoice"),
-                    ],
-                    onChunk: (chunkData) => {
-                        if (!firstTokenReceived) {
-                            firstTokenReceived = true;
-                            services.clarkProfiler.startLLMStreaming({
-                                firstChunk: chunkData.chunk.type === "text-delta"
-                                    ? chunkData.chunk.text
-                                    : `[${chunkData.chunk.type}]`,
-                            });
-                        }
-                        if (chunkData.chunk.type === "reasoning-delta") {
-                            const thinkingTrack = "thinking";
-                            services.clarkProfiler
-                                .getProfiler()
-                                .createTrack(thinkingTrack, "AI Thinking/Reasoning", "llm");
-                            allReasoningDeltas.push(chunkData.chunk.text);
-                            currentStepReasoningDeltas.push(chunkData.chunk.text);
-                            if (!thinkingSpanActive) {
-                                thinkingSpanActive = true;
-                                services.clarkProfiler
-                                    .getProfiler()
-                                    .startFrame(`Thinking Step ${stepCount + 1}`, thinkingTrack, {
-                                    stepNumber: stepCount + 1,
-                                    chunkType: chunkData.chunk.type,
+                });
+                try {
+                    const context = contextHandle.context;
+                    const mode = clark.context.currentMode;
+                    const userApprovedPlan = clark.context.userApprovedPlan;
+                    await chatSessionStore.getMessages();
+                    const latestSummary = await chatSessionStore.getLatestSummary();
+                    if (latestSummary) {
+                        context.seedIfEmpty({
+                            role: "user",
+                            content: `Reference the context from the previous exchange with the user, if relevant:\n\n${formatSummaryForAgents(latestSummary)}`,
+                        });
+                    }
+                    const systemPrompt = buildBaseSystemPrompt(mode);
+                    const system = {
+                        role: "system",
+                        content: systemPrompt,
+                    };
+                    context.setSystemPrompt(system);
+                    const prompt = buildUserMessage(userPrompt, promptContext, mode, userApprovedPlan);
+                    context.startTurn(prompt);
+                    const messages = context.getMessages();
+                    const model = services.llmProvider.modelForClassification("broad_edit");
+                    // Process LLM configuration up front so we can log it once at the top
+                    const llmConfig = clark.context.llmConfig;
+                    const disabledTools = clark.context.llmConfig?.disabledTools;
+                    const { headers, thinkingEnabled, thinkingBudgetTokens, interleavedThinking, providerOptions, } = processLLMConfig(llmConfig, 5000, // default budget tokens
+                    `LLM Generating (model=${model.modelId}, disabledTools=${disabledTools?.length ? disabledTools.join(",") : "none"})`);
+                    const conversationId = Date.now();
+                    const startTimestamp = new Date().toISOString();
+                    const runTimestamp = clark.context.runTimestamp || startTimestamp;
+                    // Initialize log content
+                    const logRef = {
+                        content: `=== LLM CONVERSATION START [${conversationId}] ===\n`,
+                    };
+                    logRef.content += `Timestamp: ${startTimestamp}\n`;
+                    logRef.content += `Model: ${model.modelId}\n`;
+                    // Log AI provider config
+                    logRef.content += `Provider Config: provider=${llmConfig?.provider || "anthropic (default)"}, thinking=${thinkingEnabled}, budget=${thinkingBudgetTokens}, interleaved=${interleavedThinking}`;
+                    if (disabledTools?.length) {
+                        logRef.content += `, disabledTools=${disabledTools.join(",")}`;
+                    }
+                    if (headers && Object.keys(headers).length > 0) {
+                        logRef.content += `, headers=${JSON.stringify(headers)}`;
+                    }
+                    logRef.content += `--- INPUT MESSAGES ---\n`;
+                    logRef.content += `[SYSTEM] ${systemPrompt}\n\n`;
+                    logRef.content += `[${system.role.toUpperCase()}] ${system.content}\n\n`;
+                    // Note: Input token count will be logged once LLM call begins
+                    // buildTools now returns filtered tools based on mode via registry
+                    const actualTools = await buildTools(clark, services, promptContext, logRef, disabledTools, mode);
+                    logRef.content += `Mode: ${mode}, Available tools: ${Object.keys(actualTools).length}\n`;
+                    logRef.content += `Tools: ${Object.keys(actualTools).join(", ")}\n\n`;
+                    let stepCount = 0;
+                    // Track cumulative token usage across all steps
+                    let totalInputTokens = 0;
+                    let totalOutputTokens = 0;
+                    let totalCachedTokens = 0;
+                    services.clarkProfiler.startLLMWaiting({
+                        messages: messages,
+                        model: model.modelId,
+                    });
+                    let firstTokenReceived = false;
+                    let thinkingSpanActive = false;
+                    let textSpanActive = false;
+                    const allReasoningDeltas = [];
+                    let currentStepReasoningDeltas = [];
+                    let currentStepTextDeltas = [];
+                    getLogger().info(`CURRENT_MODE: ${mode}`);
+                    const build = tracedStreamText({
+                        abortSignal: abortController?.signal,
+                        model,
+                        providerOptions,
+                        headers,
+                        experimental_transform: [smoothStream({ chunking: "line" })],
+                        messages,
+                        tools: actualTools,
+                        prepareStep: async (step) => {
+                            context.startStep();
+                            const messages = context.getMessages();
+                            return { ...step, messages };
+                        },
+                        stopWhen: [
+                            hasToolSuccess("build_finalize"),
+                            hasToolSuccess("exitPlanMode"),
+                            hasToolSuccess("askMultiChoice"),
+                        ],
+                        onChunk: (chunkData) => {
+                            if (!firstTokenReceived) {
+                                firstTokenReceived = true;
+                                services.clarkProfiler.startLLMStreaming({
+                                    firstChunk: chunkData.chunk.type === "text-delta"
+                                        ? chunkData.chunk.text
+                                        : `[${chunkData.chunk.type}]`,
                                 });
                             }
-                        }
-                        if (chunkData.chunk.type === "text-delta") {
-                            const textTrack = "text_generation";
-                            services.clarkProfiler
-                                .getProfiler()
-                                .createTrack(textTrack, "Text Generation", "llm");
-                            currentStepTextDeltas.push(chunkData.chunk.text);
-                            if (!textSpanActive) {
-                                textSpanActive = true;
+                            if (chunkData.chunk.type === "reasoning-delta") {
+                                const thinkingTrack = "thinking";
                                 services.clarkProfiler
                                     .getProfiler()
-                                    .startFrame(`Text Generation Step ${stepCount + 1}`, textTrack, {
-                                    stepNumber: stepCount + 1,
-                                    firstTextDelta: chunkData.chunk.text.slice(0, 50) +
-                                        (chunkData.chunk.text.length > 50 ? "..." : ""),
-                                });
-                            }
-                        }
-                    },
-                    onStepFinish: async (step) => {
-                        stepCount++;
-                        const stepTimestamp = new Date().toISOString();
-                        context.endStep(step.response.messages, step.usage);
-                        logRef.content += `--- OUTPUT STEP ${stepCount} [${stepTimestamp}] ---\n`;
-                        if (step.reasoning && thinkingSpanActive) {
-                            const thinkingTrack = "thinking";
-                            services.clarkProfiler
-                                .getProfiler()
-                                .updateActiveFrameArgs(thinkingTrack, {
-                                completeReasoningText: currentStepReasoningDeltas.join(" "),
-                                reasoningLength: step.reasoning.length,
-                                stepComplete: true,
-                            });
-                            services.clarkProfiler.getProfiler().endFrame(thinkingTrack);
-                            thinkingSpanActive = false;
-                            currentStepReasoningDeltas = [];
-                        }
-                        // Log token usage for this step and accumulate totals
-                        if (step.usage) {
-                            const stepInputTokens = step.usage.inputTokens ?? 0;
-                            const stepOutputTokens = step.usage.outputTokens ?? 0;
-                            const stepCachedTokens = step.usage.cachedInputTokens ?? 0;
-                            // Accumulate totals
-                            totalInputTokens += stepInputTokens;
-                            totalOutputTokens += stepOutputTokens;
-                            totalCachedTokens += stepCachedTokens;
-                            logRef.content += `[TOKEN USAGE] Input: ${stepInputTokens}, Output: ${stepOutputTokens}, Total: ${step.usage.totalTokens ?? 0}`;
-                            if (stepCachedTokens) {
-                                logRef.content += `, Cached: ${stepCachedTokens}`;
+                                    .createTrack(thinkingTrack, "AI Thinking/Reasoning", "llm");
+                                allReasoningDeltas.push(chunkData.chunk.text);
+                                currentStepReasoningDeltas.push(chunkData.chunk.text);
+                                if (!thinkingSpanActive) {
+                                    thinkingSpanActive = true;
+                                    services.clarkProfiler
+                                        .getProfiler()
+                                        .startFrame(`Thinking Step ${stepCount + 1}`, thinkingTrack, {
+                                        stepNumber: stepCount + 1,
+                                        chunkType: chunkData.chunk.type,
+                                    });
+                                }
                             }
-                            logRef.content += `\n`;
-                        }
-                        const apiToolCall = step.toolCalls.find((toolCall) => toolCall.toolName === "build_generateApiSource");
-                        const apiName = apiToolCall?.input?.apiName;
-                        const apiGroup = apiName ? `api-${apiName}` : undefined;
-                        // Record tool calls FIRST before other messages to ensure proper ordering
-                        const toolsCalled = step.content
-                            .filter((c) => c.type === "tool-result")
-                            .map((c) => ({
-                            toolName: c.toolName,
-                            input: JSON.stringify(c.input),
-                            output: JSON.stringify(c.output, null, 2),
-                        }));
-                        if (toolsCalled.length > 0) {
-                            logRef.content += `[TOOLS CALLED]\n`;
-                            toolsCalled.forEach((tool, idx) => {
-                                logRef.content += `  Tool ${idx + 1}: ${tool.toolName}\n`;
-                                logRef.content += `    Input: ${tool.input}\n`;
-                                logRef.content += `    Output: ${tool.output}\n`;
-                            });
-                            try {
-                                await Promise.all(step.toolCalls.map(async (toolCall) => {
-                                    const args = await getToolCallArguments(toolCall.toolName, toolCall.input, clark);
-                                    await services.chatSessionStore.recordAssistant({
-                                        type: "tool",
-                                        tool: toolCall.toolName,
-                                        args: args,
-                                        ...(apiGroup && { group: apiGroup }),
+                            if (chunkData.chunk.type === "text-delta") {
+                                const textTrack = "text_generation";
+                                services.clarkProfiler
+                                    .getProfiler()
+                                    .createTrack(textTrack, "Text Generation", "llm");
+                                currentStepTextDeltas.push(chunkData.chunk.text);
+                                if (!textSpanActive) {
+                                    textSpanActive = true;
+                                    services.clarkProfiler
+                                        .getProfiler()
+                                        .startFrame(`Text Generation Step ${stepCount + 1}`, textTrack, {
+                                        stepNumber: stepCount + 1,
+                                        firstTextDelta: chunkData.chunk.text.slice(0, 50) +
+                                            (chunkData.chunk.text.length > 50 ? "..." : ""),
                                     });
-                                }));
+                                }
                             }
-                            catch (error) {
-                                getLogger().error("Failed to record tool calls", getErrorMeta(error));
+                        },
+                        onStepFinish: async (step) => {
+                            stepCount++;
+                            const stepTimestamp = new Date().toISOString();
+                            context.endStep(step.response.messages, step.usage);
+                            logRef.content += `--- OUTPUT STEP ${stepCount} [${stepTimestamp}] ---\n`;
+                            if (step.reasoning && thinkingSpanActive) {
+                                const thinkingTrack = "thinking";
+                                services.clarkProfiler
+                                    .getProfiler()
+                                    .updateActiveFrameArgs(thinkingTrack, {
+                                    completeReasoningText: currentStepReasoningDeltas.join(" "),
+                                    reasoningLength: step.reasoning.length,
+                                    stepComplete: true,
+                                });
+                                services.clarkProfiler.getProfiler().endFrame(thinkingTrack);
+                                thinkingSpanActive = false;
+                                currentStepReasoningDeltas = [];
                             }
-                            // Record profiling events
-                            toolsCalled.forEach((tool, idx) => {
-                                let parsedInput, parsedOutput;
-                                try {
-                                    parsedInput = JSON.parse(tool.input);
-                                }
-                                catch {
-                                    parsedInput = tool.input;
+                            // Log token usage for this step and accumulate totals
+                            if (step.usage) {
+                                const stepInputTokens = step.usage.inputTokens ?? 0;
+                                const stepOutputTokens = step.usage.outputTokens ?? 0;
+                                const stepCachedTokens = step.usage.cachedInputTokens ?? 0;
+                                // Accumulate totals
+                                totalInputTokens += stepInputTokens;
+                                totalOutputTokens += stepOutputTokens;
+                                totalCachedTokens += stepCachedTokens;
+                                logRef.content += `[TOKEN USAGE] Input: ${stepInputTokens}, Output: ${stepOutputTokens}, Total: ${step.usage.totalTokens ?? 0}`;
+                                if (stepCachedTokens) {
+                                    logRef.content += `, Cached: ${stepCachedTokens}`;
                                 }
+                                logRef.content += `\n`;
+                            }
+                            const apiToolCall = step.toolCalls.find((toolCall) => toolCall.toolName === "build_generateApiSource");
+                            const apiName = apiToolCall?.input?.apiName;
+                            const apiGroup = apiName ? `api-${apiName}` : undefined;
+                            // Record tool calls FIRST before other messages to ensure proper ordering
+                            const toolsCalled = step.content
+                                .filter((c) => c.type === "tool-result")
+                                .map((c) => ({
+                                toolName: c.toolName,
+                                input: JSON.stringify(c.input),
+                                output: JSON.stringify(c.output, null, 2),
+                            }));
+                            if (toolsCalled.length > 0) {
+                                logRef.content += `[TOOLS CALLED]\n`;
+                                toolsCalled.forEach((tool, idx) => {
+                                    logRef.content += `  Tool ${idx + 1}: ${tool.toolName}\n`;
+                                    logRef.content += `    Input: ${tool.input}\n`;
+                                    logRef.content += `    Output: ${tool.output}\n`;
+                                });
                                 try {
-                                    parsedOutput = JSON.parse(tool.output);
+                                    await Promise.all(step.toolCalls.map(async (toolCall) => {
+                                        const args = await getToolCallArguments(toolCall.toolName, toolCall.input, clark);
+                                        await services.chatSessionStore.recordAssistant({
+                                            type: "tool",
+                                            tool: toolCall.toolName,
+                                            args: args,
+                                            ...(apiGroup && { group: apiGroup }),
+                                        });
+                                    }));
                                 }
-                                catch {
-                                    parsedOutput = tool.output;
+                                catch (error) {
+                                    getLogger().error("Failed to record tool calls", getErrorMeta(error));
                                 }
+                                // Record profiling events
+                                toolsCalled.forEach((tool, idx) => {
+                                    let parsedInput, parsedOutput;
+                                    try {
+                                        parsedInput = JSON.parse(tool.input);
+                                    }
+                                    catch {
+                                        parsedInput = tool.input;
+                                    }
+                                    try {
+                                        parsedOutput = JSON.parse(tool.output);
+                                    }
+                                    catch {
+                                        parsedOutput = tool.output;
+                                    }
+                                    services.clarkProfiler
+                                        .getProfiler()
+                                        .addInstantEvent(`Tool Call: ${tool.toolName}`, "llm", {
+                                        step: stepCount,
+                                        toolIndex: idx + 1,
+                                        toolName: tool.toolName,
+                                        input: parsedInput,
+                                        output: parsedOutput,
+                                        inputSize: tool.input.length,
+                                        outputSize: tool.output.length,
+                                    });
+                                });
+                            }
+                            // Record reasoning messages AFTER tool calls
+                            if (step.reasoning) {
+                                const reasoningLines = [
+                                    "[REASONING]",
+                                    ...step.reasoning.map(({ text }) => text),
+                                    "",
+                                ];
+                                logRef.content += reasoningLines.join("\n");
+                                const reasoningText = step.reasoning
+                                    .map(({ text }) => text)
+                                    .join("");
+                                if (reasoningText) {
+                                    try {
+                                        await services.chatSessionStore.recordAssistant({
+                                            type: "reasoning",
+                                            text: reasoningText,
+                                            ...(apiGroup && { group: apiGroup }),
+                                        });
+                                    }
+                                    catch (error) {
+                                        getLogger().error("Failed to record message from LLM", getErrorMeta(error));
+                                    }
+                                }
+                            }
+                            if (step.text && textSpanActive) {
+                                const textTrack = "text_generation";
                                 services.clarkProfiler
                                     .getProfiler()
-                                    .addInstantEvent(`Tool Call: ${tool.toolName}`, "llm", {
-                                    step: stepCount,
-                                    toolIndex: idx + 1,
-                                    toolName: tool.toolName,
-                                    input: parsedInput,
-                                    output: parsedOutput,
-                                    inputSize: tool.input.length,
-                                    outputSize: tool.output.length,
+                                    .updateActiveFrameArgs(textTrack, {
+                                    completeTextContent: currentStepTextDeltas.join(""),
+                                    finalText: step.text,
+                                    textLength: step.text.length,
+                                    stepComplete: true,
                                 });
-                            });
-                        }
-                        // Record reasoning messages AFTER tool calls
-                        if (step.reasoning) {
-                            const reasoningLines = [
-                                "[REASONING]",
-                                ...step.reasoning.map(({ text }) => text),
-                                "",
-                            ];
-                            logRef.content += reasoningLines.join("\n");
-                            const reasoningText = step.reasoning
-                                .map(({ text }) => text)
-                                .join("");
-                            if (reasoningText) {
+                                services.clarkProfiler.getProfiler().endFrame(textTrack);
+                                textSpanActive = false;
+                                currentStepTextDeltas = [];
+                            }
+                            // Record text messages AFTER tool calls and reasoning
+                            if (step.text) {
+                                logRef.content += `[ASSISTANT TEXT] ${step.text}\n`;
                                 try {
                                     await services.chatSessionStore.recordAssistant({
-                                        type: "reasoning",
-                                        text: reasoningText,
+                                        type: "text",
+                                        text: step.text,
                                         ...(apiGroup && { group: apiGroup }),
                                     });
                                 }
@@ -340,126 +381,102 @@ export const doLLMGenerating = (clark, services) => {
                                     getLogger().error("Failed to record message from LLM", getErrorMeta(error));
                                 }
                             }
-                        }
-                        if (step.text && textSpanActive) {
-                            const textTrack = "text_generation";
-                            services.clarkProfiler
-                                .getProfiler()
-                                .updateActiveFrameArgs(textTrack, {
-                                completeTextContent: currentStepTextDeltas.join(""),
-                                finalText: step.text,
-                                textLength: step.text.length,
-                                stepComplete: true,
+                            logRef.content += `\n`;
+                        },
+                        onFinish: (result) => {
+                            context.endTurn(result.totalUsage);
+                        },
+                    }, clark.tracer, clark.logger);
+                    for await (const chunk of build.fullStream) {
+                        await processStreamChunk(chunk, clark, logRef);
+                    }
+                    if (firstTokenReceived) {
+                        services.clarkProfiler.endFrame();
+                    }
+                    services.clarkProfiler.endFrame();
+                    if (thinkingSpanActive) {
+                        services.clarkProfiler.getProfiler().endFrame("thinking");
+                    }
+                    if (textSpanActive) {
+                        services.clarkProfiler.getProfiler().endFrame("text_generation");
+                    }
+                    const endTimestamp = new Date().toISOString();
+                    logRef.content += `=== LLM CONVERSATION END [${conversationId}] ===\n`;
+                    logRef.content += `End Timestamp: ${endTimestamp}\n`;
+                    logRef.content += `Total Steps: ${stepCount}\n`;
+                    // Log final token usage summary using accumulated totals
+                    const finalTotalTokens = totalInputTokens + totalOutputTokens;
+                    logRef.content += `[TOTAL TOKEN USAGE] Input: ${totalInputTokens}, Output: ${totalOutputTokens}, Total: ${finalTotalTokens}`;
+                    if (totalCachedTokens > 0) {
+                        logRef.content += `, Cached: ${totalCachedTokens}`;
+                    }
+                    logRef.content += `\n`;
+                    try {
+                        // Create TokenRequestData object from accumulated step data (not AI SDK usage which is only final step)
+                        const requestTokenData = {
+                            requestId: conversationId.toString(),
+                            inputTokens: totalInputTokens,
+                            outputTokens: totalOutputTokens,
+                            totalTokens: totalInputTokens + totalOutputTokens,
+                            cachedInputTokens: totalCachedTokens,
+                            model: model.modelId,
+                            startTime: startTimestamp,
+                            endTime: endTimestamp,
+                        };
+                        await clark.context.peer?.call.aiPushTokenUsage(requestTokenData);
+                    }
+                    catch (error) {
+                        // Token tracking is non-critical - log error but don't fail the AI request
+                        getLogger().warn("Failed to send token usage data", error instanceof Error ? error.message : String(error));
+                    }
+                    // Save the complete log using saveGeneratedArtifact
+                    try {
+                        const logArtifact = {
+                            type: "file",
+                            filePath: `llm-conversation-${conversationId}.log`,
+                            content: logRef.content,
+                        };
+                        const stepId = `llm-conversation-${conversationId}`;
+                        await services.appShell.saveGeneratedArtifact(logArtifact, stepId, runTimestamp);
+                        getLogger().debug("LLM conversation log saved");
+                    }
+                    catch (error) {
+                        getLogger().error("Failed to save LLM conversation log", {
+                            error: {
+                                kind: "SaveLogError",
+                                message: error instanceof Error ? error.message : String(error),
+                                stack: error instanceof Error ? error.stack : undefined,
+                            },
+                        });
+                    }
+                    // Check if there are local draft changes
+                    const hasLocalDraft = await services.draftInterface.hasLocalDraftChanges();
+                    // In PLAN mode, only proceed to draft state if:
+                    // 1. User has approved the plan (userApprovedPlan === true)
+                    // 2. There are actual file changes (hasLocalDraft === true)
+                    if (mode === AiMode.PLAN) {
+                        if (!userApprovedPlan || !hasLocalDraft) {
+                            void transitionTo({
+                                type: APP_RUNTIME_UPDATED_WITHOUT_EDITS,
                             });
-                            services.clarkProfiler.getProfiler().endFrame(textTrack);
-                            textSpanActive = false;
-                            currentStepTextDeltas = [];
-                        }
-                        // Record text messages AFTER tool calls and reasoning
-                        if (step.text) {
-                            logRef.content += `[ASSISTANT TEXT] ${step.text}\n`;
-                            try {
-                                await services.chatSessionStore.recordAssistant({
-                                    type: "text",
-                                    text: step.text,
-                                    ...(apiGroup && { group: apiGroup }),
-                                });
-                            }
-                            catch (error) {
-                                getLogger().error("Failed to record message from LLM", getErrorMeta(error));
-                            }
+                            return;
                         }
-                        logRef.content += `\n`;
-                    },
-                    onFinish: (result) => {
-                        context.endTurn(result.totalUsage);
-                    },
-                }, clark.tracer, clark.logger);
-                for await (const chunk of build.fullStream) {
-                    await processStreamChunk(chunk, clark, logRef);
-                }
-                if (firstTokenReceived) {
-                    services.clarkProfiler.endFrame();
-                }
-                services.clarkProfiler.endFrame();
-                if (thinkingSpanActive) {
-                    services.clarkProfiler.getProfiler().endFrame("thinking");
-                }
-                if (textSpanActive) {
-                    services.clarkProfiler.getProfiler().endFrame("text_generation");
-                }
-                const endTimestamp = new Date().toISOString();
-                logRef.content += `=== LLM CONVERSATION END [${conversationId}] ===\n`;
-                logRef.content += `End Timestamp: ${endTimestamp}\n`;
-                logRef.content += `Total Steps: ${stepCount}\n`;
-                // Log final token usage summary using accumulated totals
-                const finalTotalTokens = totalInputTokens + totalOutputTokens;
-                logRef.content += `[TOTAL TOKEN USAGE] Input: ${totalInputTokens}, Output: ${totalOutputTokens}, Total: ${finalTotalTokens}`;
-                if (totalCachedTokens > 0) {
-                    logRef.content += `, Cached: ${totalCachedTokens}`;
-                }
-                logRef.content += `\n`;
-                try {
-                    // Create TokenRequestData object from accumulated step data (not AI SDK usage which is only final step)
-                    const requestTokenData = {
-                        requestId: conversationId.toString(),
-                        inputTokens: totalInputTokens,
-                        outputTokens: totalOutputTokens,
-                        totalTokens: totalInputTokens + totalOutputTokens,
-                        cachedInputTokens: totalCachedTokens,
-                        model: model.modelId,
-                        startTime: startTimestamp,
-                        endTime: endTimestamp,
-                    };
-                    await clark.context.peer?.call.aiPushTokenUsage(requestTokenData);
-                }
-                catch (error) {
-                    // Token tracking is non-critical - log error but don't fail the AI request
-                    getLogger().warn("Failed to send token usage data", error instanceof Error ? error.message : String(error));
-                }
-                // Save the complete log using saveGeneratedArtifact
-                try {
-                    const logArtifact = {
-                        type: "file",
-                        filePath: `llm-conversation-${conversationId}.log`,
-                        content: logRef.content,
-                    };
-                    const stepId = `llm-conversation-${conversationId}`;
-                    await services.appShell.saveGeneratedArtifact(logArtifact, stepId, runTimestamp);
-                    getLogger().debug("LLM conversation log saved");
-                }
-                catch (error) {
-                    getLogger().error("Failed to save LLM conversation log", {
-                        error: {
-                            kind: "SaveLogError",
-                            message: error instanceof Error ? error.message : String(error),
-                            stack: error instanceof Error ? error.stack : undefined,
-                        },
-                    });
-                }
-                // Check if there are local draft changes
-                const hasLocalDraft = await services.draftInterface.hasLocalDraftChanges();
-                // In PLAN mode, only proceed to draft state if:
-                // 1. User has approved the plan (userApprovedPlan === true)
-                // 2. There are actual file changes (hasLocalDraft === true)
-                if (mode === AiMode.PLAN) {
-                    if (!userApprovedPlan || !hasLocalDraft) {
+                        // If plan is approved and there are changes, fall through to create draft
+                    }
+                    if (hasLocalDraft) {
+                        void transitionTo({
+                            type: V3_AGENT_FINISHED,
+                        });
+                    }
+                    else {
                         void transitionTo({
                             type: APP_RUNTIME_UPDATED_WITHOUT_EDITS,
                         });
-                        return;
                     }
-                    // If plan is approved and there are changes, fall through to create draft
                 }
-                if (hasLocalDraft) {
-                    void transitionTo({
-                        type: V3_AGENT_FINISHED,
-                    });
-                }
-                else {
-                    void transitionTo({
-                        type: APP_RUNTIME_UPDATED_WITHOUT_EDITS,
-                    });
+                finally {
+                    // Always release the context lock when done
+                    contextHandle.release();
                 }
             }
         }