chore(studio): extend Assistant response time and handle timeouts (#50892)

## Problem

- Assistant responses were capped at 120 seconds and 10 steps, which is
too short for longer reasoning or multi-step tool work.
- When the hosting platform ended a request at that limit, the
connection just dropped. The user got no explanation, and "Thinking…"
and tool rows kept spinning.
- Studio's own tools ignored the request's abort signal, so a stop,
disconnect or deadline couldn't cancel their in-flight requests.
- Aborted responses never closed their Braintrust span. Under TanStack
Start, the remote MCP client was only released on `res.on('close')`,
which the adapter never emits.

## Solution

Uses AI SDK options instead of custom stream handling:

- `maxDuration` goes to 300s and the step limit to 20. `streamText({
timeout: { totalMs } })` stops the response at 270s, leaving time to
finish the stream before the platform cutoff.
- `toUIMessageStream({ messageMetadata })` marks an aborted response
`timedOut: true`. `Chat` ignores `abort` chunks, so the client reads
this flag instead and shows a timeout alert with Retry. The flag is
saved with the message, so the alert survives a reload.
- `toUIMessageStream({ onEnd })` aborts the request whenever the stream
ends, releasing the MCP client on both runtimes. `streamText({ onAbort
})` ends the Braintrust span.
- Studio tools pass the SDK's `abortSignal` to their fetches. MCP tools
already did.
- Reasoning and server-tool rows that never finished show "Response
interrupted" instead of a spinner or "Ran X ✓".

There's no per-tool timeout. Approved SQL and migrations can
legitimately run longer, and aborting the HTTP request doesn't stop the
query in Postgres.

## Review instructions

1. Run the unit tests: `cd apps/studio && pnpm vitest run
lib/api/generate-v4.test.ts lib/ai components/ui/AIAssistantPanel`
2. To see a timeout without waiting 4.5 minutes, temporarily set
`ASSISTANT_TIMEOUT_MS` in `apps/studio/lib/ai/assistant-timeout.ts` to
`15_000` and run `pnpm dev:studio`.
3. Ask the Assistant something that needs several tool calls or long
reasoning, for example "Audit my schema for missing indexes and RLS
gaps, then write the fixes."
4. After 15 seconds, check that:
- the response stops and a "Assistant response timed out" alert appears
with Retry
- any in-progress reasoning or tool row shows "Response interrupted"
instead of spinning
   - Retry starts a new response
   - reloading the page still shows the alert on that chat
5. Stop a response with the Stop button before the deadline. It should
stop without the timeout alert.
6. With the default 270s, confirm that a normal response completes as
before.

## Checklist

Check all before review:

- [ ] I have read
[CONTRIBUTING.md](https://github.com/supabase/supabase/blob/master/CONTRIBUTING.md)
- [ ] If I wrote a new docs topic or edited an existing topic, I used
the `/write-the-docs` or `/edit-the-docs` skill, which references
[WORD_LIST](https://github.com/supabase/supabase/blob/master/apps/docs/WORD_LIST.md)
and the docs
[CONTRIBUTING](https://github.com/supabase/supabase/blob/master/apps/docs/CONTRIBUTING.md)
guide


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Improvements**
* AI assistant responses can now run for up to five minutes, supporting
longer requests.
* When a response times out, the assistant displays a message suggesting
you retry or ask for a smaller change.
* Incomplete responses now show a “Response interrupted” notice, and
loading indicators stop when generation ends.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Saxon FletcherandClaude Opus 5.5 authored and GitHub committed 2026-09-28 12:44:39 +10:00
1 parent 0eb08cb9f0
commit 5e59b6047e
17 files changed
+374 -71

No files matched your search

@@ -38,7 +38,11 @@ import { useLocalStorageQuery } from '@/hooks/misc/useLocalStorage'
import { useOrgAiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi'
import { useSelectedOrganizationQuery } from '@/hooks/misc/useSelectedOrganization'
import { useSelectedProjectQuery } from '@/hooks/misc/useSelectedProject'
import type { AssistantMessageMetadata } from '@/lib/ai/assistant-message-metadata'
import {
isTimedOutMessage,
type AssistantMessageMetadata,
} from '@/lib/ai/assistant-message-metadata'
import { ASSISTANT_TIMEOUT_MESSAGE } from '@/lib/ai/assistant-timeout'
import { getParallelApprovalIdsToReject } from '@/lib/ai/message-utils'
import { IS_PLATFORM } from '@/lib/constants'
import { uuidv4 } from '@/lib/helpers'
@@ -289,6 +293,11 @@ export const AssistantChat = ({
(error.message?.includes('context_length_exceeded') ||
error.message?.includes('exceeds the context window'))
const isTimedOut = !error && !isChatLoading && isTimedOutMessage(chatMessages.at(-1))
let displayError = IS_PLATFORM ? ASSISTANT_ERRORS['default'] : error
if (isContextExceededError) displayError = ASSISTANT_ERRORS['context-exceeded']
if (isTimedOut) displayError = { message: ASSISTANT_TIMEOUT_MESSAGE }
const editedMessageIndex = editingMessageId
? chatMessages.findIndex((message) => message.id === editingMessageId)
: -1
@@ -577,18 +586,16 @@ export const AssistantChat = ({
<ConversationContent className="w-full py-8 mb-10">
{renderedMessages}
<div className="w-full max-w-3xl mx-auto">
{error && (
{(error || isTimedOut) && (
<AlertError
error={
isContextExceededError
? ASSISTANT_ERRORS['context-exceeded']
: IS_PLATFORM
? ASSISTANT_ERRORS['default']
: error
}
error={displayError}
showErrorPrefix={false}
showInstructions={false}
subject="Sorry, I'm having trouble responding right now."
subject={
isTimedOut
? 'Assistant response timed out'
: "Sorry, I'm having trouble responding right now."
}
additionalActions={
<div className="flex items-center gap-x-2 mr-auto">
{isContextExceededError ? (
@@ -1,11 +1,33 @@
import type { ToolUIPart } from 'ai'
import { type PropsWithChildren } from 'react'
import { describe, expect, it } from 'vitest'
import { MessageProvider } from './Message.Context'
import { MessagePartSwitcher } from './Message.Parts'
import { customRender } from '@/tests/lib/custom-render'
type MessagePart = Parameters<typeof MessagePartSwitcher>[0]['part']
function Provider({
children,
isLoading = false,
isLastMessage = true,
}: PropsWithChildren<{ isLoading?: boolean; isLastMessage?: boolean }>) {
return (
<MessageProvider
messageInfo={{ id: 'message-1', isLoading, isLastMessage, state: 'idle' }}
messageActions={{
onDelete: () => {},
onEdit: () => {},
onBranch: () => {},
onCancelEdit: () => {},
}}
>
{children}
</MessageProvider>
)
}
describe('MessagePartSwitcher', () => {
it('keeps consecutive generic tool parts as direct siblings', () => {
const reasoningPart = {
@@ -22,10 +44,10 @@ describe('MessagePartSwitcher', () => {
} satisfies ToolUIPart
const { container } = customRender(
<>
<Provider>
<MessagePartSwitcher part={reasoningPart} />
<MessagePartSwitcher part={toolPart} />
</>
</Provider>
)
const toolRows = container.querySelectorAll('.tool-item')
@@ -33,4 +55,65 @@ describe('MessagePartSwitcher', () => {
expect(toolRows[0].nextElementSibling).toBe(toolRows[1])
expect(toolRows[0]).toHaveClass('max-w-3xl')
})
it.each([
{ type: 'reasoning', state: 'streaming', text: 'Still thinking' },
{ type: 'tool-execute_sql', state: 'input-streaming', toolCallId: 'sql-1' },
{ type: 'tool-create_notebook', state: 'input-streaming', toolCallId: 'notebook-1' },
{ type: 'tool-update_notebook', state: 'input-streaming', toolCallId: 'notebook-2' },
{ type: 'tool-query_logs', state: 'input-available', toolCallId: 'logs-1', input: {} },
] satisfies MessagePart[])('stops the $type indicator when the request ends', (part) => {
const { container, getByText, rerender } = customRender(
<Provider isLoading>
<MessagePartSwitcher part={part} />
</Provider>
)
expect(container.querySelector('.animate-spin')).not.toBeNull()
rerender(
<Provider>
<MessagePartSwitcher part={part} />
</Provider>
)
expect(getByText('Response interrupted')).toBeInTheDocument()
expect(container.querySelector('.animate-spin')).toBeNull()
})
it.each([
{ type: 'tool-search_docs', state: 'input-available', toolCallId: 'docs-1', input: {} },
{
type: 'dynamic-tool',
toolName: 'list_tables',
state: 'input-available',
toolCallId: 'mcp-1',
input: {},
},
] satisfies MessagePart[])(
'marks a $type call that never returned as interrupted once the request ends',
(part) => {
const { queryByText, getByText, rerender } = customRender(
<Provider isLoading>
<MessagePartSwitcher part={part} />
</Provider>
)
expect(queryByText('Response interrupted')).toBeNull()
rerender(
<Provider>
<MessagePartSwitcher part={part} />
</Provider>
)
expect(getByText('Response interrupted')).toBeInTheDocument()
}
)
it('does not restart an interrupted indicator when another message is streaming', () => {
const { container, getByText } = customRender(
<Provider isLoading isLastMessage={false}>
<MessagePartSwitcher part={{ type: 'reasoning', state: 'streaming', text: '' }} />
</Provider>
)
expect(getByText('Response interrupted')).toBeInTheDocument()
expect(container.querySelector('.animate-spin')).toBeNull()
})
})
@@ -1,6 +1,12 @@
import { UIMessage as VercelMessage } from '@ai-sdk/react'
import { type DynamicToolUIPart, type ReasoningUIPart, type TextUIPart, type ToolUIPart } from 'ai'
import { BrainIcon, CheckIcon, Loader2 } from 'lucide-react'
import {
isToolUIPart,
type DynamicToolUIPart,
type ReasoningUIPart,
type TextUIPart,
type ToolUIPart,
} from 'ai'
import { BrainIcon, CheckIcon, CircleStop, Loader2 } from 'lucide-react'
import { memo, type ReactNode } from 'react'
import { cn } from 'ui'
@@ -336,6 +342,32 @@ const isCompactToolPart = (part: NonNullable<VercelMessage['parts']>[number]) =>
export const MessagePartSwitcher = memo(
function MessagePartSwitcher({ part }: { part: NonNullable<VercelMessage['parts']>[number] }) {
const { isLoading, isLastMessage } = useMessageInfoContext()
const isActiveMessage = isLoading && isLastMessage
// Compact rows and query_logs run on the server, so `input-available` means the tool never
// returned. Other tools wait in that state for the user to act.
const isServerToolAwaitingOutput =
isToolUIPart(part) &&
part.state === 'input-available' &&
(isCompactToolPart(part) ||
part.type === 'tool-query_logs' ||
(part.type === 'dynamic-tool' && part.toolName === 'query_logs'))
const isIncompletePart =
(part.type === 'reasoning' && part.state === 'streaming') ||
(isToolUIPart(part) && part.state === 'input-streaming') ||
isServerToolAwaitingOutput
if (!isActiveMessage && isIncompletePart) {
return (
<Tool
icon={<CircleStop strokeWidth={1.5} size={12} className="text-foreground-muted" />}
label="Response interrupted"
>
{part.type === 'reasoning' ? part.text : undefined}
</Tool>
)
}
const content = (() => {
switch (part.type) {
case 'dynamic-tool': {
@@ -95,6 +95,7 @@ function FeedMessage({
id={message.id}
message={message}
isLoading={isLoading}
isLastMessage
isAfterEditedMessage={false}
isBeingEdited={false}
addToolApprovalResponse={addToolApprovalResponse}