@m6d/cortex-server 1.8.0 → 2.0.1

package/src/lib/ai/index.ts

@@ -1,324 +1,370 @@
 import {
+    chat,
+    chatParamsFromRequest,
+    convertMessagesToModelMessages,
+    generateMessageId,
+    maxIterations,
+    mergeAgentTools,
+    normalizeToUIMessage,
+    resumeServerSentEventsResponse,
+    toServerSentEventsResponse,
     type UIMessage,
-    type ToolSet,
-    convertToModelMessages,
-    generateId,
-    generateText,
-    safeValidateUIMessages,
-    stepCountIs,
-    streamText,
-    wrapLanguageModel,
-} from "ai";
+} from "@tanstack/ai";
 import { HTTPException } from "hono/http-exception";
-import type { ResolvedCortexAgentConfig } from "../config.ts";
-import type { MessageMetadata, Thread, TokenUsage } from "../types.ts";
-import { createModel, createEmbeddingModel } from "./helpers.ts";
-import { buildSystemPrompt, resolveSession } from "./prompt.ts";
-import { createQueryGraphTool } from "./tools/query-graph.tool.ts";
-import { createExecuteCodeTool } from "./tools/execute-code.tool.ts";
-import { captureFilesTool } from "./tools/capture-files.tool.ts";
-import { createRequestInterceptor } from "./interceptors/request-interceptor.ts";
-import { createNeo4jClient } from "../graph/neo4j.ts";
-import { resolveFromGraph } from "../graph/resolver.ts";
-import { notify } from "../ws/index.ts";
-import { buildContextMessages } from "./context/builder.ts";
-import { optimizeThreadContext, estimateTokens, trimMessagesToFit } from "./context/index.ts";
+import type { ResolvedCortexAgentConfig, ToolSet } from "@/config";
+import type { ChatMessage, MessageMetadata, Thread } from "@/types";
+import { createModel } from "./helpers";
+import { buildSystemPrompt, resolveSession } from "./prompt";
+import { ControlCenterClient } from "@/cc/client";
+import { getControlCenterConfig } from "@/cc/config-cache";
+import { buildPromptVariables } from "@/cc/format";
+import { getRecentCcTools } from "@/cc/registry";
+import { createEmbedder, createNeo4jClient } from "@cortex/contracts/graph";
+import { resolveFromGraph } from "@/graph/resolver";
+import { notify } from "@/ws/connections";
+import { buildContextMessages, trimMessagesToFit } from "./context/builder";
+import { estimateTokens } from "./context/token-estimator";
+import { createCompressionMiddleware } from "./context/compression-middleware";
+import { createCcRuntime } from "./cc-runtime";
 import {
-    compressIntraTurnToolResults,
-    estimateToolResultTokens,
-    summarizeOldStepResults,
-} from "./context/intra-turn-compressor.ts";
-import {
-    registerStream,
-    attachSseStream,
-    removeStream,
-    isStreamRunning,
-} from "./active-streams.ts";
-import { toThreadSummary } from "../types.ts";
-
-export async function stream(
-    messages: unknown[],
+    buildTurnTools,
+    createToolInstrumentation,
+    hasDefaultAttachmentInterceptor,
+} from "./turn-tools";
+import { createInspector } from "./inspector";
+import { commitBeforeTerminal } from "./commit-gate";
+import { finishTurn } from "./finish-turn";
+import { endRun, isRunning, joinLog, producerLog, startRun } from "./active-runs";
+import { toThreadSummary } from "@/types";
+import { buildAttachmentListing, prepareAttachmentsForTurn } from "./attachments";
+
+/** The AG-UI request body, as the route parsed it. */
+export type TurnParams = Awaited<ReturnType<typeof chatParamsFromRequest>>;
+
+/** Cortex's agentic loop cap; the SDK's own default is 5. */
+const MAX_ITERATIONS = 25;
+
+export async function startTurn(
+    params: TurnParams,
     thread: Thread,
     userId: string,
     token: string,
     requestContext: Record<string, unknown>,
     config: ResolvedCortexAgentConfig,
 ) {
-    const abortController = new AbortController();
-
-    const validationResult = await safeValidateUIMessages({ messages });
-    if (!validationResult.success) {
-        throw new HTTPException(423, { message: "Invalid messages format" });
-    }
-
-    const validatedMessages = validationResult.data;
-    await config.db.messages.upsert(thread.id, validatedMessages);
-    const updatedThread = await config.db.threads.touch(thread.id);
+    const run = await startRun(thread.id);
+    try {
+        // The client owns its own messages, and the assistant turn it sends back
+        // once a client-side tool has answered — without that answer stored, the
+        // next turn replays an assistant message whose tool calls nothing ever
+        // resolved, and the provider rejects the thread from there on. Everything
+        // else the server generated stays the server's; the repository keeps the
+        // usage and model id only it ever knew.
+        const incoming = params.messages
+            .map((message) => toChatMessage(normalizeToUIMessage(message, generateMessageId)))
+            .filter((message) => message.role === "user" || answersToolCalls(message));
+        await config.db.messages.upsert(thread.id, incoming);
+
+        const lastUserMessage = incoming.findLast((message) => message.role === "user");
+        const prompt = textOf(lastUserMessage);
+        if (thread.title === null && prompt) {
+            void generateTitle(thread.id, prompt, userId, config);
+        }
 
-    const activeStream = registerStream(thread.id, abortController);
+        const updatedThread = await config.db.threads.touch(thread.id);
 
-    notify(userId, thread.agentId, {
-        type: "thread:run-started",
-        payload: { thread: toThreadSummary(updatedThread, true) },
-    });
+        // A stop can land while the claim exists but before the run is
+        // announced; announcing anyway would put run-started after the abort's
+        // run-finished and pin the thread as running on every other tab. No
+        // await sits between this check and the notify, so once it passes the
+        // abort's run-finished can only follow the announcement.
+        if (run.abortController.signal.aborted) {
+            releaseRun(thread.id, run.runId);
+            return null;
+        }
 
-    // Extract prompt from the just-upserted messages (last user message)
-    // so we can start graph resolution without waiting for history fetch
-    const prompt =
-        validatedMessages
-            .filter((x) => x.role === "user")
-            .at(-1)
-            ?.parts.find((x) => x.type === "text")?.text ?? "";
-
-    const model = createModel(config.model);
-    const embeddingModel = createEmbeddingModel(config.embedding);
-    const neo4j = config.neo4j ? createNeo4jClient(config.neo4j, embeddingModel) : undefined;
-
-    // Run independent operations in parallel
-    const [contextResult, resolved, session] = await Promise.all([
-        // Branch A: Load messages + build token-aware context window
-        buildContextMessages(userId, thread, config.db, config.context),
-        // Branch B: Resolve graph context (400-2000ms, the bottleneck)
-        neo4j
-            ? resolveFromGraph(prompt, {
-                  neo4j,
-                  embeddingModel,
-                  reranker: config.reranker,
+        notify(userId, thread.agentId, {
+            type: "thread:run-started",
+            payload: { thread: toThreadSummary(updatedThread, true) },
+        });
+
+        const threadAttachments = lastUserMessage
+            ? await prepareAttachmentsForTurn({
+                  config,
+                  threadId: thread.id,
+                  userId,
+                  message: lastUserMessage,
+                  abortSignal: run.abortController.signal,
               })
-            : Promise.resolve(null),
-        // Branch C: Resolve session data
-        resolveSession(config, thread, token),
-    ]);
+            : await config.db.attachments.listByThread(thread.id, userId);
 
-    const { messages: contextMessages, allMessages: originalMessages } = contextResult;
+        // Claiming rewrites the message the user just sent, and the tabs that did
+        // not send it are already showing that message — with the run underway,
+        // this is their only chance to learn which files it carries.
+        if (threadAttachments.some((attachment) => attachment.messageId === lastUserMessage?.id)) {
+            notify(userId, thread.agentId, {
+                type: "thread:messages-updated",
+                payload: { threadId: thread.id, thread: toThreadSummary(updatedThread, true) },
+            });
+        }
 
-    // Build tools
-    const builtInTools: ToolSet = {
-        captureFiles: captureFilesTool,
-    };
+        const turnKey = lastUserMessage?.id ?? "0";
+
+        const inspector = createInspector();
+        const { model, embed, neo4j, ccClient } = createTurnProviders(config, inspector.fetch);
+
+        // Run independent operations in parallel
+        const [contextResult, resolved, session, ccConfig, ccResolved] = await Promise.all([
+            // Branch A: Load messages + build token-aware context window
+            buildContextMessages(userId, thread, config.db, config.context),
+            // Branch B: Resolve graph context (400-2000ms, the bottleneck)
+            neo4j && embed
+                ? resolveFromGraph(prompt, { neo4j, embed, reranker: config.reranker })
+                : Promise.resolve(null),
+            // Branch C: Resolve session data
+            resolveSession(config, thread, token),
+            // Branch D: Control Center config (ETag-revalidated, last-known-good on outage)
+            ccClient
+                ? getControlCenterConfig(ccClient, config.agentId, run.abortController.signal)
+                : Promise.resolve(null),
+            // Branch E: Control Center /resolve — on failure the turn degrades to
+            // the search meta-tools, it never blocks
+            ccClient && prompt
+                ? ccClient.resolve(
+                      config.agentId,
+                      {
+                          query: prompt,
+                          hints: {
+                              threadId: thread.id,
+                              recentToolNames: getRecentCcTools(thread.id),
+                          },
+                      },
+                      run.abortController.signal,
+                  )
+                : Promise.resolve(null),
+        ]);
+
+        const cc = createCcRuntime({
+            ccClient,
+            ccConfig,
+            ccResolved,
+            config,
+            thread,
+            userId,
+            token,
+            turnKey,
+            abortSignal: run.abortController.signal,
+            getStepIndex: () => inspector.steps.length,
+            onToolProgress: (toolName, toolCallId, status) =>
+                emitToolProgress(userId, thread, toolName, toolCallId, status),
+        });
+
+        const { messages: contextMessages, allMessages } = contextResult;
+
+        const turnToolsOptions = {
+            config,
+            cc,
+            neo4j,
+            thread,
+            userId,
+            token,
+            session,
+            requestContext,
+            threadAttachments,
+        };
+        const tools = buildTurnTools(turnToolsOptions);
+        const instrumentation = createToolInstrumentation({
+            ...turnToolsOptions,
+            onToolStart: (toolName, toolCallId) =>
+                emitToolProgress(userId, thread, toolName, toolCallId, "started"),
+            onToolFinish: (toolName, toolCallId) =>
+                emitToolProgress(userId, thread, toolName, toolCallId, "finished"),
+        });
+
+        const systemPrompt = await buildSystemPrompt(
+            config,
+            resolved,
+            { session, requestContext },
+            cc
+                ? {
+                      config: cc.config,
+                      resolved: ccResolved,
+                      variables: buildPromptVariables(cc.config.promptVariables, {
+                          // Computed per request; consumer-supplied context
+                          // with the same key wins via the spreads below. No
+                          // such default for `timezone`: the server's zone is
+                          // not the session's, so guessing would feed the
+                          // model wrong user context — only the consumer can
+                          // supply it.
+                          utcTime: new Date().toISOString().replace(/\.\d{3}Z$/, "Z"),
+                          ...session,
+                          ...requestContext,
+                          locale: cc.locale,
+                      }),
+                  }
+                : undefined,
+            buildAttachmentListing(threadAttachments, lastUserMessage?.id, {
+                readAttachment: tools.some((tool) => tool.name === "readAttachment"),
+                attachmentSentinels: hasDefaultAttachmentInterceptor(config),
+            }),
+        );
+
+        const stream = chat({
+            adapter: model,
+            // Client tool declarations ride in on every request, so the server
+            // never re-declares them.
+            tools: mergeAgentTools(tools, params.tools),
+            messages: convertMessagesToModelMessages(
+                fitToContextWindow(contextMessages, systemPrompt, tools, config),
+            ),
+            systemPrompts: [systemPrompt],
+            agentLoopStrategy: maxIterations(MAX_ITERATIONS),
+            abortController: run.abortController,
+            middleware: [
+                createCompressionMiddleware(config),
+                ...(instrumentation ? [instrumentation] : []),
+                inspector.middleware,
+            ],
+            threadId: thread.id,
+            runId: run.runId,
+        });
+
+        const producer = toServerSentEventsResponse(
+            commitBeforeTerminal(stream, {
+                isAborted: () => run.abortController.signal.aborted,
+                commit: async function (produced) {
+                    const messages = produced.map(toChatMessage).map(withTurnMetadata);
+                    await finishTurn({
+                        messages,
+                        // Summarization and the consumer hook judge the thread, not
+                        // the turn — the run's own output is only its tail.
+                        threadMessages: [...allMessages, ...messages],
+                        isAborted: run.abortController.signal.aborted,
+                        assistantMessageId: messages.at(-1)?.id,
+                        thread,
+                        userId,
+                        capturedSteps: inspector.steps,
+                        config,
+                    });
+                },
+            }),
+            {
+                // Append per chunk rather than per batch of 32: a client joining
+                // the log is waiting on real tokens, not on a buffer filling up.
+                durability: { adapter: producerLog(run.runId), batch: 1 },
+                abortController: run.abortController,
+            },
+        );
+
+        function withTurnMetadata(message: ChatMessage, index: number, all: ChatMessage[]) {
+            if (index !== all.length - 1 || message.role !== "assistant") return message;
+
+            return {
+                ...message,
+                metadata: {
+                    modelId: inspector.modelId,
+                    isAborted: run.abortController.signal.aborted,
+                    tokenUsage: inspector.tokenUsage,
+                } satisfies MessageMetadata,
+            };
+        }
 
-    if (neo4j) {
-        builtInTools["queryGraph"] = createQueryGraphTool(neo4j);
+        // Generation outlives the request that started it, so a reload mid-turn
+        // rejoins the log instead of restarting the model. Draining into a sink
+        // that never cancels is what detaches the producer from this response;
+        // the client gets a reader on the same log.
+        void producer
+            .body!.pipeTo(new WritableStream())
+            .catch(() => undefined)
+            .finally(() => releaseRun(thread.id, run.runId));
+
+        return resumeServerSentEventsResponse({ adapter: joinLog(thread.id, run.runId) });
+    } catch (error) {
+        releaseRun(thread.id, run.runId);
+        throw error;
     }
+}
 
-    if (config.backendFetch) {
-        const backendFetchWithInterceptor = {
-            ...config.backendFetch,
-            transformRequestBody:
-                config.backendFetch.transformRequestBody ??
-                createRequestInterceptor(
-                    config.db,
-                    config.storage,
-                    config.backendFetch.interceptor,
-                ),
-        };
+/** Releasing the claim must never mask the turn's own error or reject unhandled. */
+function releaseRun(threadId: string, runId: string) {
+    endRun(threadId, runId).catch((error: unknown) =>
+        console.error("[cortex-server] Failed to release run claim:", error),
+    );
+}
 
-        builtInTools["executeCode"] = createExecuteCodeTool(backendFetchWithInterceptor, token);
-    }
+/** Drops the SDK's transport shape into the envelope cortex persists and serves. */
+function toChatMessage(message: UIMessage) {
+    return { id: message.id, role: message.role, parts: message.parts } satisfies ChatMessage;
+}
 
-    const resolvedUserTools =
-        typeof config.tools === "function"
-            ? config.tools({ thread, userId, token, session, requestContext })
-            : config.tools;
+/**
+ * An assistant turn the client is sending back because it resolved a tool call.
+ *
+ * The wire fans that turn out into a `role: "tool"` message per result, and each
+ * of those normalizes back into an assistant message carrying only the result,
+ * under a freshly generated id. Stored as rows of their own they would replay as
+ * assistant turns answering calls that no message makes, so the turn is picked
+ * out by holding both halves of the exchange.
+ */
+function answersToolCalls(message: ChatMessage) {
+    return (
+        message.role === "assistant" &&
+        message.parts.some((part) => part.type === "tool-call") &&
+        message.parts.some((part) => part.type === "tool-result")
+    );
+}
 
-    const tools = {
-        ...builtInTools,
-        ...resolvedUserTools,
-    } as ToolSet;
+function textOf(message: ChatMessage | undefined) {
+    const part = message?.parts.find((candidate) => candidate.type === "text");
+    return part?.type === "text" ? part.content : "";
+}
 
-    const systemPrompt = await buildSystemPrompt(config, resolved, {
-        session,
-        requestContext,
-    });
+/**
+ * Builds the clients one turn needs. Neo4j runs vector search through an embedding
+ * model, so configuring a graph without one is a misconfiguration we surface
+ * loudly rather than quietly serving every turn with no graph context.
+ */
+function createTurnProviders(config: ResolvedCortexAgentConfig, fetch: typeof globalThis.fetch) {
+    if (config.neo4j && !config.embedding) {
+        throw new HTTPException(500, {
+            message: "embedding config is required when neo4j is configured",
+        });
+    }
 
-    // The context builder reserved a static token budget for the system prompt + tools.
-    // Now that we have the actual values, verify the reserve was sufficient and trim
-    // the oldest messages if it wasn't.
+    const embed = config.embedding ? createEmbedder(config.embedding) : undefined;
+
+    return {
+        // The inspector's fetch is per-turn, so the model is too.
+        model: createModel(config.model, { fetch }),
+        embed,
+        neo4j: config.neo4j && embed ? createNeo4jClient(config.neo4j, embed) : undefined,
+        ccClient: config.controlCenter ? new ControlCenterClient(config.controlCenter) : null,
+    };
+}
+
+/**
+ * The context builder reserved a static token budget for the system prompt and
+ * tools before either existed. Now that both are real, drop the oldest messages
+ * if the reserve fell short.
+ */
+function fitToContextWindow(
+    contextMessages: ChatMessage[],
+    systemPrompt: string,
+    tools: ToolSet,
+    config: ResolvedCortexAgentConfig,
+) {
     const actualFixedCost = estimateTokens(systemPrompt) + estimateTokens(JSON.stringify(tools));
     const { reservedTokenBudget, maxContextTokens } = config.context;
-    const trimmedMessages =
-        actualFixedCost > reservedTokenBudget
-            ? trimMessagesToFit(contextMessages, maxContextTokens - actualFixedCost)
-            : contextMessages;
-
-    const recentMessages = await convertToModelMessages(trimmedMessages);
-
-    const { context: contextConfig } = config;
-
-    const capturedSteps: {
-        prompt: string;
-        output: string | null;
-        tokenUsage: TokenUsage | null;
-    }[] = [];
-
-    const result = streamText({
-        model: wrapLanguageModel({
-            model,
-            middleware: {
-                specificationVersion: "v3",
-                transformParams: async ({ params }) => {
-                    // Strategy 1: Compress older step tool results to one-liners
-                    let optimizedPrompt = compressIntraTurnToolResults(
-                        params.prompt,
-                        contextConfig.toolResultMaxTokens,
-                    );
-
-                    // Strategy 5: If still over threshold, summarize with LLM
-                    const summarizationModel = contextConfig.summarizationModel ?? config.model;
-                    const toolTokens = estimateToolResultTokens(optimizedPrompt);
-
-                    if (toolTokens > contextConfig.intraTurnSummarizationThresholdTokens) {
-                        try {
-                            optimizedPrompt = await summarizeOldStepResults(
-                                optimizedPrompt,
-                                contextConfig.intraTurnSummarizationThresholdTokens,
-                                summarizationModel,
-                            );
-                        } catch (err) {
-                            console.error("[cortex-server] Intra-turn summarization failed:", err);
-                        }
-                    }
-
-                    return { ...params, prompt: optimizedPrompt };
-                },
-            },
-        }),
-        system: systemPrompt,
-        tools,
-        messages: recentMessages,
-        abortSignal: abortController.signal,
-        stopWhen: stepCountIs(25),
-        onStepFinish: (step) => {
-            const usage = step.usage;
-            capturedSteps.push({
-                prompt: JSON.stringify(step.request.body),
-                output: JSON.stringify({
-                    text: step.text,
-                    content: step.content,
-                    finishReason: step.finishReason,
-                    response: {
-                        id: step.response.id,
-                        modelId: step.response.modelId,
-                        messages: step.response.messages,
-                    },
-                    providerMetadata: step.providerMetadata,
-                }),
-                tokenUsage: {
-                    input: {
-                        noCache: usage.inputTokenDetails.noCacheTokens ?? 0,
-                        cacheRead: usage.inputTokenDetails.cacheReadTokens ?? 0,
-                        cacheWrite: usage.inputTokenDetails.cacheWriteTokens ?? 0,
-                        total: usage.inputTokens ?? 0,
-                    },
-                    output: {
-                        reasoning: usage.outputTokenDetails.reasoningTokens ?? 0,
-                        text: usage.outputTokenDetails.textTokens ?? 0,
-                        total: usage.outputTokens ?? 0,
-                    },
-                    total: usage.totalTokens ?? 0,
-                },
-            });
-        },
-    });
+    if (actualFixedCost <= reservedTokenBudget) return contextMessages;
 
-    return result.toUIMessageStreamResponse<UIMessage<MessageMetadata>>({
-        originalMessages,
-        generateMessageId: generateId,
-        consumeSseStream: ({ stream: sseStream }) => {
-            attachSseStream(thread.id, sseStream);
-        },
-        onFinish: async ({ messages: finishedMessages, isAborted }) => {
-            if (isAborted) {
-                finalizeAbortedMessages(finishedMessages);
-            }
-
-            // Record token usage (result promises reject on abort, so skip)
-            const lastAssistantMessage = finishedMessages
-                .filter((x) => x.role === "assistant")
-                .at(-1);
-            if (lastAssistantMessage && !isAborted) {
-                const providerMetadata = await result.providerMetadata;
-                const response = await result.response;
-                const usage = await result.totalUsage;
-                let metadata: MessageMetadata = {
-                    isAborted,
-                    providerMetadata,
-                    modelId: response.modelId,
-                    tokenUsage: {
-                        input: {
-                            noCache: usage.inputTokenDetails.noCacheTokens ?? 0,
-                            cacheRead: usage.inputTokenDetails.cacheReadTokens ?? 0,
-                            cacheWrite: usage.inputTokenDetails.cacheWriteTokens ?? 0,
-                            total: usage.inputTokens ?? 0,
-                        },
-                        output: {
-                            reasoning: usage.outputTokenDetails.reasoningTokens ?? 0,
-                            text: usage.outputTokenDetails.textTokens ?? 0,
-                            total: usage.outputTokens ?? 0,
-                        },
-                        total: usage.totalTokens ?? 0,
-                    },
-                };
-
-                lastAssistantMessage.metadata = metadata;
-            } else if (lastAssistantMessage) {
-                lastAssistantMessage.metadata = {
-                    isAborted,
-                    modelId: "",
-                    providerMetadata: undefined,
-                };
-            }
-
-            const persistedThread = await config.db.threads.getById(userId, thread.id);
-            if (!persistedThread) {
-                removeStream(thread.id);
-                return;
-            }
-
-            await config.db.messages.upsert(thread.id, finishedMessages);
-
-            // Persist captured LLM request/response data per step
-            if (capturedSteps.length > 0 && lastAssistantMessage) {
-                const llmRequestInserts = capturedSteps.map((step, index) => ({
-                    messageId: lastAssistantMessage.id,
-                    stepNumber: index,
-                    prompt: step.prompt,
-                    output: step.output,
-                    tokenUsage: step.tokenUsage,
-                }));
-
-                try {
-                    await config.db.llmRequests.insert(llmRequestInserts);
-                } catch (err) {
-                    console.error("[cortex-server] Failed to persist AI requests:", err);
-                }
-            }
-
-            config.onStreamFinish?.({ messages: finishedMessages, isAborted });
-
-            // XXX: we need to notify the user so that the client can
-            // fetch new messages. The client can't fetch messages
-            // immediately because messages may not have been saved yet.
-            notify(userId, persistedThread.agentId, {
-                type: "thread:messages-updated",
-                payload: {
-                    threadId: thread.id,
-                    thread: toThreadSummary(persistedThread, false),
-                },
-            });
+    return trimMessagesToFit(contextMessages, maxContextTokens - actualFixedCost);
+}
 
-            setTimeout(() => removeStream(thread.id, activeStream.id), 10_000);
+const TITLE_SYSTEM_PROMPT = `You are an expert in generating titles for threads of chats
+given the first message in that thread.
 
-            // Fire-and-forget: optimize context for next request
-            // Runs after response is delivered — no perceived latency
-            try {
-                optimizeThreadContext(thread, finishedMessages, config);
-            } catch (err) {
-                console.error("[cortex-server] Context optimization failed:", err);
-            }
-        },
-    });
-}
+When asked, only respond with the title without saying you're
+going to do so or any other speech. Spit out only the title.`;
 
 export async function generateTitle(
     threadId: string,
@@ -326,53 +372,44 @@ export async function generateTitle(
     userId: string,
     config: ResolvedCortexAgentConfig,
 ) {
-    const model = createModel(config.model);
-
-    const { output } = await generateText({
-        model,
-        system: `You are an expert in generating titles for threads of chats
-given the first message in that thread.
-
-When asked, only respond with the title without saying you're
-going to do so or any other speech. Spit out only the title.`,
-        prompt: `Generate a title for this prompt: ${prompt}`,
+    const title = await chat({
+        adapter: createModel(config.model, { name: "title" }),
+        stream: false,
+        systemPrompts: [TITLE_SYSTEM_PROMPT],
+        messages: [
+            {
+                id: generateMessageId(),
+                role: "user",
+                parts: [{ type: "text", content: `Generate a title for this prompt: ${prompt}` }],
+            },
+        ],
     });
 
-    await config.db.threads.updateTitle(threadId, output ?? "");
+    // A failed generation resolves to empty text rather than throwing, and the
+    // only trigger for a retry is the title still being null — so writing "" here
+    // would leave the thread permanently unnamed.
+    if (!title) return;
+
+    await config.db.threads.updateTitle(threadId, title);
 
     const thread = await config.db.threads.getById(userId, threadId);
     if (!thread) return;
 
     notify(userId, thread.agentId, {
         type: "thread:title-updated",
-        payload: { thread: toThreadSummary(thread, isStreamRunning(thread.id)) },
+        payload: { thread: toThreadSummary(thread, await isRunning(thread.id)) },
     });
 }
 
-const TERMINAL_TOOL_STATES = new Set(["output-available", "output-error", "output-denied"]);
-
-function finalizeAbortedMessages(messages: UIMessage[]) {
-    const lastMessage = messages.at(-1);
-    if (!lastMessage || lastMessage.role !== "assistant") return;
-
-    lastMessage.parts = lastMessage.parts.map((part) => {
-        if ((part.type === "text" || part.type === "reasoning") && part.state === "streaming") {
-            return { ...part, state: "done" as const };
-        }
-
-        if ("toolCallId" in part && "state" in part) {
-            const toolState = part.state;
-            if (!TERMINAL_TOOL_STATES.has(toolState)) {
-                const { approval: _, ...rest } = part;
-                return {
-                    ...rest,
-                    state: "output-error" as const,
-                    errorText: "Generation was aborted",
-                    output: undefined,
-                };
-            }
-        }
-
-        return part;
+function emitToolProgress(
+    userId: string,
+    thread: Thread,
+    toolName: string,
+    toolCallId: string,
+    status: "started" | "finished",
+) {
+    notify(userId, thread.agentId, {
+        type: "thread:tool-progress",
+        payload: { threadId: thread.id, toolCallId, toolName, status },
     });
 }