@superblocksteam/vite-plugin-file-sync 2.0.43-next.13 → 2.0.43-next.15

package/dist/ai-service/agent/subagents/apis/generate-api-source.js

@@ -1,198 +1,35 @@
 import fs from "node:fs/promises";
 import path from "node:path";
-import { hasToolCall, smoothStream, tool, } from "ai";
+import { smoothStream } from "ai";
 import z from "zod";
-import { getErrorMeta, getLogger } from "../../../../util/logger.js";
+import { getErrorMeta, getLogger, getPrefixedLogger, } from "../../../../util/logger.js";
 import { datasourceSdkClassByType, Paths } from "../../../const.js";
 import { tracedStreamText } from "../../../llmobs/helpers.js";
 import { getContextId } from "../../../state-machine/helpers/context-id.js";
-import { ApiBuilderTypeScriptError } from "../../../transform/api-builder/shared.js";
 import { YamlToApiBuilderTransformer } from "../../../transform/api-builder/to-sdk-transformer.js";
-import { ApiBuilderToYamlTransformer } from "../../../transform/api-builder/to-yaml-transformer.js";
 import { applyFileTransformations, renderPath, } from "../../../transform/shared.js";
 import { processLLMConfig } from "../../../util/llm-config-utils.js";
+import { safeJsonStringify } from "../../../util/safe-stringify.js";
+import { hasToolSuccess } from "../../../util/stop-condition.js";
 import { getToolCallArguments } from "../../tool-message-utils.js";
-import { executeRequestToIntegrationToolFactory, searchIntegrationsToolFactory, } from "../../tools/index.js";
-import { ensureScopeFileDraft, getPageEntitiesFiltered, processStreamChunk, } from "../../utils.js";
+import { buildApiToolFactory } from "../../tools/apis/build-api.js";
+import { finalizeApiToolFactory } from "../../tools/apis/finalize-api.js";
+import { executeRequestToIntegrationToolFactory, searchIntegrationsToolFactory, readIntegrationMetadataToolFactory, } from "../../tools/index.js";
+import { getPageEntitiesFiltered, processStreamChunk } from "../../utils.js";
 import { createToolFactory } from "./../../tools2/types.js";
+import { ApiGenerationContext } from "./context.js";
 import { selectExamplesByTags, formatSelectedExamples, } from "./example-selector.js";
 import { VALID_TAGS } from "./examples.js";
 import { buildApiPrompt } from "./prompt-builder.js";
-import { apiStaticAnalysis } from "./static-analysis.js";
+import { ApiGenerationStateMap } from "./state.js";
 import systemPrompt from "./system-prompt.js";
-export const buildApiTools = (input, isNewApi, promptContext, clark, services, scopeUpdateQueue, 
-// TODO(colin): find a better way to do this
-onFinalized) => {
-    const { apiName, pageName, prompt } = input;
-    const integrationById = new Map();
-    let apiYaml;
-    let buildApiAttempts = 0;
-    for (const integration of promptContext?.integrations ?? []) {
-        integrationById.set(integration.id, integration);
-    }
-    const readIntegrationMetadata = tool({
-        description: "Read the metadata for a given integration configuration ID. REQUIRED before using database integrations (PostgreSQL, MySQL, Databricks, Snowflake, etc.) to get the correct schema/table structure. Call this first when working with any integration to avoid making up column names or table structures.",
-        inputSchema: z.object({ integrationId: z.string() }),
-        execute: async ({ integrationId }) => {
-            const integration = integrationById.get(integrationId);
-            if (!integration) {
-                throw new Error(`Integration with ID ${integrationId} not found`);
-            }
-            const integrationWithOptimizedMetadata = await services.integrationStore.getIntegrationForPromptContext(integration, prompt, clark, services);
-            return integrationWithOptimizedMetadata?.metadata ?? {};
-        },
-    });
-    const fetchIntegration = tool({
-        description: "Fetch the integration configuration for a given integration ID",
-        inputSchema: z.object({ integrationId: z.string() }),
-        execute: async ({ integrationId }) => {
-            const integration = await clark.context.peer?.call.aiGetIntegration({
-                integrationId,
-            });
-            if (integration) {
-                const existing = integrationById.get(integrationId);
-                const mergedMetadata = {
-                    ...(integration?.metadata ?? {}),
-                    ...(existing?.metadata ?? {}),
-                };
-                integrationById.set(integrationId, {
-                    ...integration,
-                    metadata: mergedMetadata,
-                });
-            }
-            return integration;
-        },
-    });
-    const buildApi = tool({
-        description: "Compile the API source code.",
-        inputSchema: z.object({
-            apiSource: z.string().describe("The source code of the API"),
-            apiName: z.string().describe("The name of the API"),
-        }),
-        execute: async ({ apiSource }) => {
-            buildApiAttempts++;
-            if (buildApiAttempts > 5) {
-                throw new Error("Fatal error: Too many attempts to build the API");
-            }
-            const generatedArtifact = {
-                type: "file",
-                filePath: renderPath(Paths.GeneratedApis, {
-                    pageName,
-                    apiName,
-                }) + ".ts",
-                content: apiSource,
-            };
-            try {
-                [apiYaml] = await applyFileTransformations([generatedArtifact], [
-                    new ApiBuilderToYamlTransformer({
-                        templateRenderer: services.templateRenderer,
-                        skipCleanup: true,
-                    }),
-                ]);
-                // At this point we know the apiSource is valid TypeScript.
-                // Now do some static analysis to ensure it's semantically valid
-                // We do this against the TS source instead of the YAML because
-                // it's easier to provide better errors to the LLM, which is dealing in TS
-                const analysisWarnings = apiStaticAnalysis(apiSource);
-                if (analysisWarnings.length > 0) {
-                    throw new Error(`Static analysis detected potential issues in the API:\n\n${analysisWarnings.join("\n\n")}`);
-                }
-            }
-            catch (error) {
-                if (error instanceof ApiBuilderTypeScriptError) {
-                    throw new Error(`There were issues compiling the apiSource to TypeScript. Correct the errors below and try again.
-
-${error.message.replace(/src\/to-yaml/g, "")}
-            `);
-                }
-                else {
-                    throw error;
-                }
-            }
-            return {
-                message: "API compiled successfully",
-                apiName,
-                pageName,
-            };
-        },
-    });
-    const finalizeApi = tool({
-        description: "Finalize. Create a summary of the work completed and a TypeScript interface describing the expected response from the API.",
-        inputSchema: z.object({
-            responseInterface: z
-                .string()
-                .describe("A TypeScript interface describing the expected response from the API"),
-            inputInterface: z
-                .string()
-                .describe("A TypeScript interface describing the expected input to the API"),
-            summary: z.string().describe("A summary of the API"),
-            apiName: z.string().describe("The name of the API"),
-            givingUpDueToFatalError: z
-                .boolean()
-                .optional()
-                .describe("Whether to give up due to a fatal error"),
-        }),
-        execute: async ({ responseInterface, inputInterface, summary, givingUpDueToFatalError, }) => {
-            if (givingUpDueToFatalError) {
-                const message = "Giving up due to a fatal error. No API was created.";
-                onFinalized(message);
-                return {
-                    message,
-                    apiName,
-                    pageName,
-                };
-            }
-            if (!apiYaml) {
-                throw new Error("Unexpected state: Was the buildApi tool called successfully?");
-            }
-            if (isNewApi) {
-                // if it's a new API, the scope will be touched
-                await scopeUpdateQueue.enqueue(async () => {
-                    await ensureScopeFileDraft({
-                        identifier: pageName,
-                        type: "name",
-                    }, services);
-                    await services.draftInterface.createDraftFile(apiYaml.filePath, apiYaml.content);
-                });
-            }
-            else {
-                await services.draftInterface.createDraftFile(apiYaml.filePath, apiYaml.content);
-            }
-            let message = `${summary}
-
-The API response corresponds to the following TypeScript interface:
-
-\`\`\`typescript
-${responseInterface}
-\`\`\`
-
-The API input corresponds to the following TypeScript interface:
-
-\`\`\`typescript
-${inputInterface}
-\`\`\`
-`;
-            if (isNewApi) {
-                message = `${message}
-
-The ${pageName} scope was updated to include the ${apiName} entity.
-`;
-            }
-            onFinalized(message);
-            return {
-                message: summary,
-                apiName,
-                pageName,
-            };
-        },
-    });
+export const buildApiTools = (clark, services, scopeUpdateQueue) => {
     return {
-        readIntegrationMetadata,
-        fetchIntegration,
+        readIntegrationMetadata: readIntegrationMetadataToolFactory.create(clark, services),
         searchIntegrations: searchIntegrationsToolFactory.create(clark, services),
         executeRequestToIntegration: executeRequestToIntegrationToolFactory.create(clark, services),
-        buildApi,
-        finalizeApi,
+        buildApi: buildApiToolFactory.create(clark, services),
+        finalizeApi: finalizeApiToolFactory.create(clark, services, scopeUpdateQueue),
     };
 };
 export const readApiSource = async (pageName, apiName, pageEntityNames, services) => {
@@ -225,6 +62,7 @@ export const readApiSource = async (pageName, apiName, pageEntityNames, services
     return sdkArtifact;
 };
 export const generateApiSourceSubagentToolFactory = createToolFactory("build_generateApiSource", (clark, services, scopeUpdateQueue, promptContext, logRef) => {
+    const logger = getPrefixedLogger("[api-subagent]");
     const model = services.llmProvider.modelForClassification("broad_edit");
     const llmConfig = clark.context.llmConfig;
     return {
@@ -233,8 +71,14 @@ export const generateApiSourceSubagentToolFactory = createToolFactory("build_gen
 When calling this tool, analyze the user's requirements and provide 2-4 relevant exampleHints tags.
 `,
         inputSchema: z.object({
-            apiName: z.string().describe("The name of the API"),
-            pageName: z.string().describe("The name of the page"), // this could be an enum from app context
+            id: z
+                .object({
+                apiName: z.string().describe("The name of the API"),
+                pageName: z
+                    .string()
+                    .describe("The name of the page the API belongs to"), // this could be an enum from app context
+            })
+                .describe("The reference ID of the API being finalized"),
             action: z.enum(["create", "edit"]).describe("The action to take"),
             prompt: z
                 .string()
@@ -245,330 +89,362 @@ When calling this tool, analyze the user's requirements and provide 2-4 relevant
                 .describe(`Optional tags to help select relevant examples to guide the API subagent. Provide 2-4 most relevant tags based on the task requirements.`),
         }),
         execute: async (input) => {
-            const { apiName, pageName, prompt, exampleHints } = input;
-            let eligibleIntegrations;
-            let missingIntegrationIds;
-            const pageEntities = getPageEntitiesFiltered(pageName, ["StateVar", "Component"], services);
-            const pageEntityNames = pageEntities.map((entity) => entity.name);
-            const existingSdkApiArtifact = await readApiSource(pageName, apiName, pageEntityNames, services);
-            if (existingSdkApiArtifact) {
-                const existingIntegrationIds = services.appContextStore.apiSemanticsByName[apiName]?.blocks
-                    .map((block) => block.id)
-                    .filter((id) => id !== undefined) ?? [];
-                eligibleIntegrations = existingIntegrationIds
-                    .map((id) => promptContext?.integrations?.find((integration) => integration.id === id))
-                    .filter((x) => x !== undefined);
-                const promptContextIntegrationIds = new Set(promptContext?.integrations?.map((integration) => integration.id) ??
-                    []);
-                missingIntegrationIds = existingIntegrationIds.filter((id) => !promptContextIntegrationIds.has(id));
-            }
-            else {
-                eligibleIntegrations = promptContext?.integrations;
-            }
-            let finalizedMessage = "";
-            const tools = buildApiTools(input, !existingSdkApiArtifact, promptContext, clark, services, scopeUpdateQueue, (message) => {
-                finalizedMessage = message;
-            });
-            // Build API prompt using shared helper
-            const apiPrompt = buildApiPrompt({
-                apiName,
-                userPrompt: prompt,
-                pageEntities: pageEntities.map((e) => ({
-                    name: e.name,
-                    type: e.type,
-                    componentType: e.componentType,
-                })),
-                availableIntegrations: eligibleIntegrations?.map((integration) => ({
-                    name: integration.name,
-                    id: integration.id,
-                    type: datasourceSdkClassByType[integration.type] || integration.type,
-                })),
-                existingApiSource: existingSdkApiArtifact?.content,
-                missingIntegrationIds,
-                isEdit: !!existingSdkApiArtifact,
-            });
-            // Select relevant examples based on hints provided by upstream LLM
-            const selectedExampleMetadata = selectExamplesByTags(exampleHints, {
-                maxExamples: 6,
-                includeNegative: true,
-                maxTokenBudget: 7000,
-            });
-            const examplesContent = formatSelectedExamples(selectedExampleMetadata);
-            // Log selected examples for debugging
-            if (exampleHints && exampleHints.length > 0) {
-                logRef.content += `[EXAMPLE SELECTION] Hints: ${exampleHints.join(", ")}\n`;
-                logRef.content += `[EXAMPLE SELECTION] Selected: ${selectedExampleMetadata.map((e) => e.id).join(", ")}\n\n`;
-            }
-            // Set up context for the API subagent
-            const contextId = getContextId(clark, services);
+            const { id, prompt, exampleHints } = input;
+            const { pageName, apiName } = id;
+            // Set up context for the API subagent with thread-safe ownership
+            const subagentName = `api-subagent-${pageName}-${apiName}`;
+            const contextId = getContextId(clark, services, subagentName);
             const contextOptions = clark.context.llmConfig?.contextOptions;
-            const context = await services.contextManager.getContext(contextId, contextOptions);
-            const mainSystemPrompt = context.getSystemPrompt();
-            // Combine system prompt with examples
-            const fullSystemPrompt = examplesContent
-                ? `${systemPrompt}\n\n${examplesContent}`
-                : systemPrompt;
-            context.setSystemPrompt({
-                role: "system",
-                content: fullSystemPrompt,
-            });
-            context.startTurn({
-                role: "user",
-                content: apiPrompt,
+            // Generate a unique owner ID for this subagent execution
+            const ownerId = `${subagentName}-${Date.now()}-${Math.random().toString(36).slice(2, 9)}`;
+            // Acquire exclusive context access with retries
+            const contextHandle = await services.contextManager.acquireContext(contextId, ownerId, {
+                contextOptions,
+                acquireOptions: {
+                    waitTimeoutMs: 10000, // Wait up to 10 seconds for context
+                    lockTimeoutMs: 300000, // Hold lock for up to 5 minutes
+                    retryIntervalMs: 100,
+                },
             });
-            // Process LLM configuration
-            const { providerOptions, headers } = processLLMConfig(llmConfig, 5000, // default budget tokens
-            `API Subagent (model=${model.modelId})`);
-            const apiSubagentTrack = "api_subagent";
-            services.clarkProfiler
-                .getProfiler()
-                .createTrack(apiSubagentTrack, "API Subagent", "tool");
-            return services.clarkProfiler.getProfiler().timeAsync(`generateApiSource:${apiName}`, async () => {
-                let stepCount = 0;
-                let thinkingSpanActive = false;
-                let textSpanActive = false;
-                const allReasoningDeltas = [];
-                let currentStepReasoningDeltas = [];
-                let currentStepTextDeltas = [];
-                // Step accumulators for accurate token counting
-                const startTime = new Date().toISOString();
-                let totalInputTokens = 0;
-                let totalOutputTokens = 0;
-                let totalCachedTokens = 0;
-                let finalizeApiData = {
+            try {
+                const context = contextHandle.context;
+                const mainSystemPrompt = context.getSystemPrompt();
+                let eligibleIntegrations;
+                let missingIntegrationIds;
+                const pageEntities = getPageEntitiesFiltered(pageName, ["StateVar", "Component"], services);
+                const pageEntityNames = pageEntities.map((entity) => entity.name);
+                const existingSdkApiArtifact = await readApiSource(pageName, apiName, pageEntityNames, services);
+                if (existingSdkApiArtifact) {
+                    const existingIntegrationIds = services.appContextStore.apiSemanticsByName[apiName]?.blocks
+                        .map((block) => block.id)
+                        .filter((id) => id !== undefined) ?? [];
+                    eligibleIntegrations = existingIntegrationIds
+                        .map((id) => promptContext?.integrations?.find((integration) => integration.id === id))
+                        .filter((x) => x !== undefined);
+                    const promptContextIntegrationIds = new Set(promptContext?.integrations?.map((integration) => integration.id) ?? []);
+                    missingIntegrationIds = existingIntegrationIds.filter((id) => !promptContextIntegrationIds.has(id));
+                }
+                else {
+                    eligibleIntegrations = promptContext?.integrations;
+                }
+                if (!clark.context.apiGenStateMap) {
+                    logger.info("Initializing API generation state map");
+                    clark.context.apiGenStateMap = new ApiGenerationStateMap();
+                }
+                const apiGenStateMap = clark.context.apiGenStateMap;
+                apiGenStateMap.set(id, {
+                    status: "ready",
+                    isNewApi: !existingSdkApiArtifact,
+                });
+                logger.info(`API generation started for ${pageName}/${apiName} (existing: ${existingSdkApiArtifact ? "yes" : "no"})`);
+                const tools = buildApiTools(clark, services, scopeUpdateQueue);
+                // Build API prompt using shared helper
+                const apiPrompt = buildApiPrompt({
                     apiName,
-                    responseInterface: "",
-                    inputInterface: "",
-                    summary: "",
-                };
-                const response = tracedStreamText({
-                    model,
-                    messages: context.getMessages(),
-                    abortSignal: clark.context.abortController?.signal,
-                    providerOptions,
-                    headers,
-                    experimental_transform: [smoothStream({ chunking: "line" })],
-                    tools,
-                    prepareStep: async (step) => {
-                        context.startStep();
-                        const messages = context.getMessages();
-                        return { ...step, messages };
-                    },
-                    stopWhen: hasToolCall("finalizeApi"),
-                    onChunk: (chunkData) => {
-                        if (chunkData.chunk.type === "reasoning-delta") {
-                            allReasoningDeltas.push(chunkData.chunk.text);
-                            currentStepReasoningDeltas.push(chunkData.chunk.text);
-                            if (!thinkingSpanActive) {
-                                thinkingSpanActive = true;
+                    userPrompt: prompt,
+                    pageEntities: pageEntities.map((e) => ({
+                        name: e.name,
+                        type: e.type,
+                        componentType: e.componentType,
+                    })),
+                    availableIntegrations: eligibleIntegrations?.map((integration) => ({
+                        name: integration.name,
+                        id: integration.id,
+                        type: datasourceSdkClassByType[integration.type] || integration.type,
+                    })),
+                    existingApiSource: existingSdkApiArtifact?.content,
+                    missingIntegrationIds,
+                    isEdit: !!existingSdkApiArtifact,
+                });
+                // Select relevant examples based on hints provided by upstream LLM
+                const selectedExampleMetadata = selectExamplesByTags(exampleHints, {
+                    maxExamples: 6,
+                    includeNegative: true,
+                    maxTokenBudget: 7000,
+                });
+                const examplesContent = formatSelectedExamples(selectedExampleMetadata);
+                // Log selected examples for debugging
+                if (exampleHints && exampleHints.length > 0) {
+                    logRef.content += `[EXAMPLE SELECTION] Hints: ${exampleHints.join(", ")}\n`;
+                    logRef.content += `[EXAMPLE SELECTION] Selected: ${selectedExampleMetadata.map((e) => e.id).join(", ")}\n\n`;
+                }
+                // Combine system prompt with examples
+                const fullSystemPrompt = examplesContent
+                    ? `${systemPrompt}\n\n${examplesContent}`
+                    : systemPrompt;
+                context.setSystemPrompt({
+                    role: "system",
+                    content: fullSystemPrompt,
+                });
+                context.startTurn({
+                    role: "user",
+                    content: apiPrompt,
+                });
+                // Process LLM configuration
+                const { providerOptions, headers } = processLLMConfig(llmConfig, 5000, // default budget tokens
+                `API Subagent (model=${model.modelId})`);
+                const apiSubagentTrack = "api_subagent";
+                services.clarkProfiler
+                    .getProfiler()
+                    .createTrack(apiSubagentTrack, "API Subagent", "tool");
+                return await services.clarkProfiler.getProfiler().timeAsync(`generateApiSource:${apiName}`, async () => {
+                    let stepCount = 0;
+                    let thinkingSpanActive = false;
+                    let textSpanActive = false;
+                    const allReasoningDeltas = [];
+                    let currentStepReasoningDeltas = [];
+                    let currentStepTextDeltas = [];
+                    // Step accumulators for accurate token counting
+                    const startTime = new Date().toISOString();
+                    let totalInputTokens = 0;
+                    let totalOutputTokens = 0;
+                    let totalCachedTokens = 0;
+                    const response = tracedStreamText({
+                        model,
+                        messages: context.getMessages(),
+                        abortSignal: clark.context.abortController?.signal,
+                        providerOptions,
+                        headers,
+                        experimental_transform: [smoothStream({ chunking: "line" })],
+                        tools,
+                        prepareStep: async (step) => {
+                            context.startStep();
+                            const messages = context.getMessages();
+                            return { ...step, messages };
+                        },
+                        stopWhen: hasToolSuccess("finalizeApi"),
+                        onChunk: (chunkData) => {
+                            if (chunkData.chunk.type === "reasoning-delta") {
+                                allReasoningDeltas.push(chunkData.chunk.text);
+                                currentStepReasoningDeltas.push(chunkData.chunk.text);
+                                if (!thinkingSpanActive) {
+                                    thinkingSpanActive = true;
+                                    services.clarkProfiler
+                                        .getProfiler()
+                                        .startFrame(`API Thinking Step ${stepCount + 1}`, apiSubagentTrack, {
+                                        stepNumber: stepCount + 1,
+                                        chunkType: chunkData.chunk.type,
+                                        apiName,
+                                    });
+                                }
+                            }
+                            if (chunkData.chunk.type === "text-delta") {
+                                currentStepTextDeltas.push(chunkData.chunk.text);
+                                if (!textSpanActive) {
+                                    textSpanActive = true;
+                                    services.clarkProfiler
+                                        .getProfiler()
+                                        .startFrame(`API Text Generation Step ${stepCount + 1}`, apiSubagentTrack, {
+                                        stepNumber: stepCount + 1,
+                                        firstTextDelta: chunkData.chunk.text.slice(0, 50) +
+                                            (chunkData.chunk.text.length > 50 ? "..." : ""),
+                                        apiName,
+                                    });
+                                }
+                            }
+                        },
+                        // Clark often gets the id wrong across tool calls, so we put it in context
+                        experimental_context: new ApiGenerationContext(id),
+                        onStepFinish: async (step) => {
+                            stepCount++;
+                            const stepTimestamp = new Date().toISOString();
+                            context.endStep(step.response.messages, step.usage);
+                            logRef.content += `--- OUTPUT STEP (API subagent) [${stepTimestamp}] ---\n`;
+                            totalInputTokens += step.usage?.inputTokens ?? 0;
+                            if (step.reasoning && thinkingSpanActive) {
                                 services.clarkProfiler
                                     .getProfiler()
-                                    .startFrame(`API Thinking Step ${stepCount + 1}`, apiSubagentTrack, {
-                                    stepNumber: stepCount + 1,
-                                    chunkType: chunkData.chunk.type,
+                                    .updateActiveFrameArgs(apiSubagentTrack, {
+                                    completeReasoningText: currentStepReasoningDeltas.join(" "),
+                                    reasoningLength: step.reasoning.length,
+                                    stepComplete: true,
                                     apiName,
                                 });
+                                services.clarkProfiler
+                                    .getProfiler()
+                                    .endFrame(apiSubagentTrack);
+                                thinkingSpanActive = false;
+                                currentStepReasoningDeltas = [];
                             }
-                        }
-                        if (chunkData.chunk.type === "text-delta") {
-                            currentStepTextDeltas.push(chunkData.chunk.text);
-                            if (!textSpanActive) {
-                                textSpanActive = true;
+                            if (step.text && textSpanActive) {
                                 services.clarkProfiler
                                     .getProfiler()
-                                    .startFrame(`API Text Generation Step ${stepCount + 1}`, apiSubagentTrack, {
-                                    stepNumber: stepCount + 1,
-                                    firstTextDelta: chunkData.chunk.text.slice(0, 50) +
-                                        (chunkData.chunk.text.length > 50 ? "..." : ""),
+                                    .updateActiveFrameArgs(apiSubagentTrack, {
+                                    completeTextContent: currentStepTextDeltas.join(""),
+                                    finalText: step.text,
+                                    textLength: step.text.length,
+                                    stepComplete: true,
                                     apiName,
                                 });
+                                services.clarkProfiler
+                                    .getProfiler()
+                                    .endFrame(apiSubagentTrack);
+                                textSpanActive = false;
+                                currentStepTextDeltas = [];
                             }
-                        }
-                    },
-                    onStepFinish: async (step) => {
-                        stepCount++;
-                        const stepTimestamp = new Date().toISOString();
-                        context.endStep(step.response.messages, step.usage);
-                        logRef.content += `--- OUTPUT STEP (API subagent) [${stepTimestamp}] ---\n`;
-                        totalInputTokens += step.usage?.inputTokens ?? 0;
-                        if (step.reasoning && thinkingSpanActive) {
-                            services.clarkProfiler
-                                .getProfiler()
-                                .updateActiveFrameArgs(apiSubagentTrack, {
-                                completeReasoningText: currentStepReasoningDeltas.join(" "),
-                                reasoningLength: step.reasoning.length,
-                                stepComplete: true,
-                                apiName,
-                            });
-                            services.clarkProfiler
-                                .getProfiler()
-                                .endFrame(apiSubagentTrack);
-                            thinkingSpanActive = false;
-                            currentStepReasoningDeltas = [];
-                        }
-                        if (step.text && textSpanActive) {
-                            services.clarkProfiler
-                                .getProfiler()
-                                .updateActiveFrameArgs(apiSubagentTrack, {
-                                completeTextContent: currentStepTextDeltas.join(""),
-                                finalText: step.text,
-                                textLength: step.text.length,
-                                stepComplete: true,
-                                apiName,
-                            });
-                            services.clarkProfiler
-                                .getProfiler()
-                                .endFrame(apiSubagentTrack);
-                            textSpanActive = false;
-                            currentStepTextDeltas = [];
-                        }
-                        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.trim()) {
-                                void services.chatSessionStore.recordAssistant({
-                                    type: "reasoning",
-                                    text: reasoningText,
-                                    group: `api-${apiName}`,
-                                });
-                            }
-                        }
-                        if (step.toolCalls.length > 0) {
-                            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,
+                            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.trim()) {
+                                    void services.chatSessionStore.recordAssistant({
+                                        type: "reasoning",
+                                        text: reasoningText,
                                         group: `api-${apiName}`,
                                     });
-                                }));
-                            }
-                            catch (error) {
-                                getLogger().error("Failed to record subagent tool calls", getErrorMeta(error));
-                            }
-                        }
-                        // Record text messages AFTER tool calls
-                        if (step.text) {
-                            logRef.content += `[ASSISTANT TEXT] ${step.text}\n`;
-                            void services.chatSessionStore.recordAssistant({
-                                type: "text",
-                                text: step.text,
-                                group: `api-${apiName}`,
-                            });
-                        }
-                        // Accumulate token usage for this step
-                        if (step.usage) {
-                            const stepInputTokens = step.usage.inputTokens ?? 0;
-                            const stepOutputTokens = step.usage.outputTokens ?? 0;
-                            const stepCachedTokens = step.usage.cachedInputTokens ?? 0;
-                            // Accumulate tokens across all steps
-                            totalInputTokens += stepInputTokens;
-                            totalOutputTokens += stepOutputTokens;
-                            totalCachedTokens += stepCachedTokens;
-                        }
-                        const finalizeApiResult = step.toolResults.find((r) => r.toolName === "finalizeApi");
-                        if (finalizeApiResult) {
-                            finalizeApiData = finalizeApiResult.input;
-                        }
-                        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`;
-                            });
-                            toolsCalled.forEach((tool, idx) => {
-                                let parsedInput, parsedOutput;
-                                try {
-                                    parsedInput = JSON.parse(tool.input);
-                                }
-                                catch {
-                                    parsedInput = tool.input;
                                 }
+                            }
+                            if (step.toolCalls.length > 0) {
                                 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,
+                                            group: `api-${apiName}`,
+                                        });
+                                    }));
                                 }
-                                catch {
-                                    parsedOutput = tool.output;
+                                catch (error) {
+                                    getLogger().error("Failed to record subagent tool calls", getErrorMeta(error));
                                 }
-                                services.clarkProfiler
-                                    .getProfiler()
-                                    .addInstantEvent(`API Tool Call: ${tool.toolName}`, apiSubagentTrack, {
-                                    step: stepCount,
-                                    toolIndex: idx + 1,
-                                    toolName: tool.toolName,
-                                    input: parsedInput,
-                                    output: parsedOutput,
-                                    inputSize: tool.input.length,
-                                    outputSize: tool.output.length,
-                                    apiName,
+                            }
+                            // Record text messages AFTER tool calls
+                            if (step.text) {
+                                logRef.content += `[ASSISTANT TEXT] ${step.text}\n`;
+                                void services.chatSessionStore.recordAssistant({
+                                    type: "text",
+                                    text: step.text,
+                                    group: `api-${apiName}`,
+                                });
+                            }
+                            // Accumulate token usage for this step
+                            if (step.usage) {
+                                const stepInputTokens = step.usage.inputTokens ?? 0;
+                                const stepOutputTokens = step.usage.outputTokens ?? 0;
+                                const stepCachedTokens = step.usage.cachedInputTokens ?? 0;
+                                // Accumulate tokens across all steps
+                                totalInputTokens += stepInputTokens;
+                                totalOutputTokens += stepOutputTokens;
+                                totalCachedTokens += stepCachedTokens;
+                            }
+                            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`;
+                                });
+                                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(`API Tool Call: ${tool.toolName}`, apiSubagentTrack, {
+                                        step: stepCount,
+                                        toolIndex: idx + 1,
+                                        toolName: tool.toolName,
+                                        input: parsedInput,
+                                        output: parsedOutput,
+                                        inputSize: tool.input.length,
+                                        outputSize: tool.output.length,
+                                        apiName,
+                                    });
                                 });
-                            });
+                            }
+                            logRef.content += `\n`;
+                        },
+                        onFinish: (result) => {
+                            context.endTurn(result.totalUsage);
+                            if (mainSystemPrompt) {
+                                // restore main system prompt
+                                context.setSystemPrompt(mainSystemPrompt);
+                            }
+                        },
+                    }, clark.tracer, clark.logger);
+                    for await (const chunk of response.fullStream) {
+                        await processStreamChunk(chunk, clark, logRef, `api-${apiName}`);
+                    }
+                    if (thinkingSpanActive) {
+                        services.clarkProfiler.getProfiler().endFrame(apiSubagentTrack);
+                    }
+                    if (textSpanActive) {
+                        services.clarkProfiler.getProfiler().endFrame(apiSubagentTrack);
+                    }
+                    try {
+                        const requestTokenData = {
+                            requestId: `api-subagent-${apiName}-${Date.now()}`,
+                            inputTokens: totalInputTokens,
+                            outputTokens: totalOutputTokens,
+                            totalTokens: totalInputTokens + totalOutputTokens,
+                            cachedInputTokens: totalCachedTokens,
+                            model: model.modelId,
+                            startTime: startTime,
+                            endTime: new Date().toISOString(),
+                        };
+                        await clark.context.peer?.call.aiPushTokenUsage(requestTokenData);
+                    }
+                    catch (error) {
+                        // Token tracking is non-critical - log error but don't fail the API generation
+                        getLogger().warn("Failed to send token usage data for API generation", error instanceof Error ? error.message : String(error));
+                    }
+                    const stateMap = clark.context.apiGenStateMap;
+                    if (!stateMap) {
+                        throw new Error("No API generation state map found in context");
+                    }
+                    const finalState = stateMap.get(id);
+                    if (!finalState) {
+                        throw new Error("No state found for API ${pageName}/${apiName} after generation");
+                    }
+                    try {
+                        if (finalState.status === "finalized") {
+                            logger.info(`API generation succeeded for ${pageName}/${apiName}`);
+                            return finalState;
                         }
-                        logRef.content += `\n`;
-                    },
-                    onFinish: (result) => {
-                        context.endTurn(result.totalUsage);
-                        if (mainSystemPrompt) {
-                            // restore main system prompt
-                            context.setSystemPrompt(mainSystemPrompt);
+                        else {
+                            logger.error(`API generation failed for ${pageName}/${apiName}: final status: ${finalState.status}`);
+                            throw new Error(`API generation failed, final state: ${safeJsonStringify(finalState)}`);
                         }
-                    },
-                }, clark.tracer, clark.logger);
-                for await (const chunk of response.fullStream) {
-                    await processStreamChunk(chunk, clark, logRef, `api-${apiName}`);
-                }
-                if (thinkingSpanActive) {
-                    services.clarkProfiler.getProfiler().endFrame(apiSubagentTrack);
-                }
-                if (textSpanActive) {
-                    services.clarkProfiler.getProfiler().endFrame(apiSubagentTrack);
-                }
-                try {
-                    const requestTokenData = {
-                        requestId: `api-subagent-${apiName}-${Date.now()}`,
-                        inputTokens: totalInputTokens,
-                        outputTokens: totalOutputTokens,
-                        totalTokens: totalInputTokens + totalOutputTokens,
-                        cachedInputTokens: totalCachedTokens,
-                        model: model.modelId,
-                        startTime: startTime,
-                        endTime: new Date().toISOString(),
-                    };
-                    await clark.context.peer?.call.aiPushTokenUsage(requestTokenData);
-                }
-                catch (error) {
-                    // Token tracking is non-critical - log error but don't fail the API generation
-                    getLogger().warn("Failed to send token usage data for API generation", error instanceof Error ? error.message : String(error));
-                }
-                return {
-                    message: finalizedMessage,
+                    }
+                    finally {
+                        stateMap.delete(id);
+                    }
+                }, apiSubagentTrack, {
                     apiName,
                     pageName,
-                    apiData: finalizeApiData,
-                };
-            }, apiSubagentTrack, {
-                apiName,
-                pageName,
-                promptLength: apiPrompt.length,
-                model: model.modelId,
-            });
+                    promptLength: apiPrompt.length,
+                    model: model.modelId,
+                });
+            }
+            finally {
+                // Always release the context lock when done
+                contextHandle.release();
+            }
         },
     };
 });