Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions extensions/github1s-ai/assets/chat.css
Original file line number Diff line number Diff line change
Expand Up @@ -1144,6 +1144,16 @@ summary:focus-visible,
font: inherit;
}

.message-retry {
gap: 4px;
margin-top: 6px;
}

.message-retry .icon {
width: 12px;
height: 12px;
}

.message-status {
margin: 5px 0 0;
color: var(--vscode-descriptionForeground);
Expand Down
16 changes: 16 additions & 0 deletions extensions/github1s-ai/src/common/conversation.ts
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,22 @@ export interface Conversation extends ConversationSummary {
messages: ConversationMessage[];
}

export const isInterruptedMessage = (message: ConversationMessage): boolean =>
message.metadata.status === 'failed' ||
message.metadata.status === 'aborted' ||
message.metadata.status === 'unknown';

export const getRetryableMessage = (messages: readonly ConversationMessage[]): ConversationMessage | undefined => {
const assistant = messages.at(-1);
const user = messages.at(-2);
return assistant?.role === 'assistant' &&
isInterruptedMessage(assistant) &&
user?.role === 'user' &&
user.metadata.turnId === assistant.metadata.turnId
? assistant
: undefined;
};

export const createUserMessage = (
turnId: string,
text: string,
Expand Down
2 changes: 2 additions & 0 deletions extensions/github1s-ai/src/common/protocol.ts
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ type AllViewEvents =
| { type: 'app.openSettings' }
| { type: 'app.openFile'; source: string }
| { type: 'chat.send'; text: string }
| { type: 'chat.retry'; id: string }
| { type: 'chat.runQuickAction'; action: ChatQuickAction }
| { type: 'chat.addContextAttachment'; action: ContextAttachmentAction }
| { type: 'chat.addContextAttachment'; action: 'descriptor'; descriptor: ContextAttachmentDescriptor }
Expand Down Expand Up @@ -96,6 +97,7 @@ export const parseViewEvent = (value: unknown): ViewEvent | undefined => {
? { type: event.type, enabled: event.enabled }
: undefined;

case 'chat.retry':
case 'chat.removeContextAttachment':
case 'history.selectConversation':
case 'history.deleteConversation':
Expand Down
5 changes: 5 additions & 0 deletions extensions/github1s-ai/src/controllers/chat.ts
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,11 @@ export class ChatController extends Controller {
await this.runner.send({ text: event.text });
}

@Controller.handler('chat.retry')
async handleRetry(event: ViewEvent<'chat.retry'>): Promise<void> {
await this.runner.send({ retryMessageId: event.id });
}

@Controller.handler('chat.cancel')
async handleCancel(_event: ViewEvent<'chat.cancel'>): Promise<void> {
await this.runner.cancelCurrent();
Expand Down
75 changes: 32 additions & 43 deletions extensions/github1s-ai/src/controllers/runner.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import { addLanguageModelUsage, createNullLanguageModelUsage } from 'ai/internal
import {
createAssistantMessage,
createUserMessage,
getRetryableMessage,
withMessageStatus,
type Conversation,
type ConversationMessage,
Expand All @@ -19,7 +20,7 @@ import { resolveContextAttachments } from './context';

type ActiveRequest = {
controller: AbortController;
turnId: string;
messageId: string;
};

export class ConversationRunner {
Expand All @@ -31,15 +32,16 @@ export class ConversationRunner {
private readonly publishState: () => Promise<void>,
) {}

async send(input: { text: string } | { action: ChatQuickAction }): Promise<void> {
async send(input: { text: string } | { action: ChatQuickAction } | { retryMessageId: string }): Promise<void> {
if (('text' in input && !input.text.trim()) || this.preparation) return;

const [promptsConfig, config] = await Promise.all([
this.stores.promptsConfig.get(),
this.stores.modelConfigs.getSelected(),
]);
const prompts = resolvePrompts(promptsConfig);
const text = 'text' in input ? input.text : prompts.quickActions[input.action];
const text = 'text' in input ? input.text : 'action' in input ? prompts.quickActions[input.action] : '';
const retry = 'retryMessageId' in input;

if (!config) {
await this.stores.runtime.setIn('chat.notice', {
Expand All @@ -50,36 +52,38 @@ export class ConversationRunner {
}
const runtime = await this.stores.runtime.get();
if (this.preparation) return;
const previous = runtime.chat.conversation;
if (retry && (!previous || getRetryableMessage(previous.messages)?.id !== input.retryMessageId)) return;
const conversationId = previous?.id ?? globalThis.crypto.randomUUID();
if (this.requests.has(conversationId)) return;
const descriptors = (runtime.chat.pendingAttachments ?? []).map((descriptor) => ({ ...descriptor }));
const preparation = Symbol('conversation-preparation');
this.preparation = preparation;
await this.stores.runtime.setIn('chat', { ...runtime.chat, preparing: true, notice: undefined });

let attachments;
let userMessage: ConversationMessage;
try {
await this.publishState();
attachments = await resolveContextAttachments(descriptors);
if (retry) {
userMessage = previous!.messages.at(-2)!;
} else {
const attachments = await resolveContextAttachments(descriptors);
const recentFiles = runtime.chat.includeRecentFiles === false ? [] : runtime.chat.recentFiles;
userMessage = createUserMessage(globalThis.crypto.randomUUID(), text, attachments, recentFiles);
}
} catch (error) {
const detail = error instanceof Error ? error.message : 'Unknown error.';
await this.finishPreparation(preparation, `Unable to read the selected context attachment. ${detail}`);
return;
}
if (this.preparation !== preparation) return;

const previous = runtime.chat.conversation;
const startedAt = Date.now();
const conversationId = previous?.id ?? globalThis.crypto.randomUUID();
if (this.requests.has(conversationId)) {
await this.finishPreparation(preparation);
return;
}

const turnId = globalThis.crypto.randomUUID();
const recentFiles = runtime.chat.includeRecentFiles === false ? [] : runtime.chat.recentFiles;
const userMessage = createUserMessage(turnId, text, attachments, recentFiles);
const turnId = userMessage.metadata.turnId;
const history = retry ? previous!.messages.slice(0, -2) : (previous?.messages ?? []);
let providerMessages;
try {
providerMessages = await buildModelMessages(previous?.messages ?? [], [userMessage]);
providerMessages = await buildModelMessages(history, [userMessage]);
} catch (error) {
const detail = error instanceof Error ? error.message : 'Unknown error.';
await this.finishPreparation(preparation, `Unable to prepare the conversation. ${detail}`);
Expand All @@ -91,7 +95,7 @@ export class ConversationRunner {
const conversation: Conversation = previous
? {
...previous,
messages: [...previous.messages, userMessage, assistantMessage],
messages: [...history, userMessage, assistantMessage],
updatedAt: startedAt,
}
: {
Expand All @@ -116,15 +120,15 @@ export class ConversationRunner {
return;
}
if (this.preparation !== preparation) {
await this.finishTurn(conversationId, turnId, 'aborted');
await this.setMessageStatus(conversationId, assistantMessage.id, 'aborted');
return;
}

this.preparation = undefined;
const abortController = new AbortController();
const request: ActiveRequest = {
controller: abortController,
turnId,
messageId: assistantMessage.id,
};
this.requests.set(conversationId, request);
let mcp: Awaited<ReturnType<typeof connectMcpTools>> | undefined;
Expand All @@ -134,7 +138,7 @@ export class ConversationRunner {
await this.stores.runtime.setIn('chat', {
...currentRuntime.chat,
conversation,
pendingAttachments: [],
pendingAttachments: retry ? currentRuntime.chat.pendingAttachments : [],
preparing: false,
});
await this.stores.conversations.select(conversationId);
Expand All @@ -154,16 +158,16 @@ export class ConversationRunner {
if (!abortController.signal.aborted && this.requests.get(conversationId) === request) {
const usage = await agentStream.usage;
if (!abortController.signal.aborted && this.requests.get(conversationId) === request) {
await this.completeTurn(conversationId, turnId, usage);
await this.setMessageStatus(conversationId, request.messageId, 'completed', undefined, usage);
}
}
} catch (error) {
if (this.requests.get(conversationId) !== request) return;
if (abortController.signal.aborted) {
await this.finishTurn(conversationId, turnId, 'aborted');
await this.setMessageStatus(conversationId, request.messageId, 'aborted');
} else {
const message = error instanceof Error ? error.message : 'Unknown error.';
await this.finishTurn(conversationId, turnId, 'failed', message);
await this.setMessageStatus(conversationId, request.messageId, 'failed', message);
}
} finally {
if (this.requests.get(conversationId) === request) this.requests.delete(conversationId);
Expand All @@ -185,7 +189,7 @@ export class ConversationRunner {
if (!request) return;
this.requests.delete(conversationId);
request.controller.abort();
await this.finishTurn(conversationId, request.turnId, 'aborted');
await this.setMessageStatus(conversationId, request.messageId, 'aborted');
}

isRequestActive(conversationId: string | undefined): boolean {
Expand Down Expand Up @@ -218,32 +222,17 @@ export class ConversationRunner {
await this.publishState();
}

private async completeTurn(conversationId: string, turnId: string, usage: LanguageModelUsage): Promise<void> {
await this.setTurnStatus(conversationId, turnId, 'completed', undefined, usage);
}

private async finishTurn(
conversationId: string,
turnId: string,
status: 'aborted' | 'failed',
error?: string,
): Promise<void> {
await this.setTurnStatus(conversationId, turnId, status, error);
}

private async setTurnStatus(
private async setMessageStatus(
conversationId: string,
turnId: string,
messageId: string,
status: 'completed' | 'aborted' | 'failed',
error?: string,
usage?: LanguageModelUsage,
): Promise<void> {
const conversation = await this.stores.conversations.get(conversationId);
if (!conversation) return;
if (!conversation?.messages.some((message) => message.id === messageId)) return;
const messages = conversation.messages.map((message) =>
message.role === 'assistant' && message.metadata.turnId === turnId
? withMessageStatus(message, status, error)
: message,
message.id === messageId ? withMessageStatus(message, status, error) : message,
);
const updatedAt = Date.now();
const nextUsage = usage === undefined ? undefined : addUsage(conversation.usage, usage);
Expand Down
39 changes: 36 additions & 3 deletions extensions/github1s-ai/src/llm/helpers.ts
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import { convertToModelMessages, type ModelMessage } from 'ai';
import { convertToModelMessages, getToolName, isToolUIPart, type ModelMessage } from 'ai';

import { contextReferencePath, type ContextAttachment, type ContextReference } from '@/common/context';
import type { ConversationMessage } from '@/common/conversation';
import { isInterruptedMessage, type ConversationMessage } from '@/common/conversation';

const MAX_PROMPT_HISTORY_CHARACTERS = 128 * 1024;

Expand All @@ -10,7 +10,7 @@ export const buildModelMessages = (
input: readonly ConversationMessage[],
): Promise<ModelMessage[]> =>
convertToModelMessages<ConversationMessage>(
[...selectWholeRecentTurns(history, MAX_PROMPT_HISTORY_CHARACTERS), ...input],
[...selectWholeRecentTurns(history, MAX_PROMPT_HISTORY_CHARACTERS).map(repairInterruptedMessage), ...input],
{
ignoreIncompleteToolCalls: true,
convertDataPart: (part) => {
Expand All @@ -24,6 +24,39 @@ export const buildModelMessages = (
},
);

// Repair only the model input. Keep the original partial response intact for display.
const repairInterruptedMessage = (message: ConversationMessage): ConversationMessage => {
if (message.role !== 'assistant' || !isInterruptedMessage(message)) return message;
let endsWithToolStep = false;
const parts = message.parts.flatMap<ConversationMessage['parts'][number]>((part) => {
// Interrupted provider items and reasoning signatures may not be replayable.
if (part.type === 'text') return part.text.trim() ? [{ type: 'text', text: part.text }] : [];
if (part.type === 'step-start') {
endsWithToolStep = false;
return [part];
}
if (!isToolUIPart(part) || part.state === 'input-streaming') return [];
endsWithToolStep = true;
const tool = {
type: 'dynamic-tool' as const,
toolName: getToolName(part),
toolCallId: part.toolCallId,
input: part.state === 'output-error' && 'rawInput' in part ? (part.input ?? part.rawInput) : part.input,
};
if (part.state === 'output-available' && !part.preliminary)
return [{ ...tool, state: 'output-available', output: part.output }];
if (part.state === 'output-error') return [{ ...tool, state: 'output-error', errorText: part.errorText }];
if (part.state === 'output-denied' || (part.state === 'approval-responded' && !part.approval.approved))
return [{ ...tool, state: 'output-error', errorText: part.approval.reason ?? 'Tool execution was denied.' }];
return [{ ...tool, state: 'output-error', errorText: 'Tool execution was interrupted; its result is unknown.' }];
});
// Tool results follow all text in their step, so close that step before the next user turn.
if (endsWithToolStep || parts.at(-1)?.type !== 'text') {
parts.push({ type: 'step-start' }, { type: 'text', text: '[The previous assistant response was interrupted.]' });
}
return { ...message, parts };
};

const selectWholeRecentTurns = (
history: readonly ConversationMessage[],
maxCharacters: number,
Expand Down
13 changes: 11 additions & 2 deletions extensions/github1s-ai/src/webview/components/ChatPage.ts
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import { html } from 'htm/preact';
import { useLayoutEffect, useRef } from 'preact/hooks';

import type { ConversationMessage } from '@/common/conversation';
import { getRetryableMessage, type ConversationMessage } from '@/common/conversation';
import type { ViewEvent, ViewState } from '@/common/protocol';
import { QUICK_ACTIONS } from '@/common/quick-actions';

Expand Down Expand Up @@ -47,6 +47,7 @@ interface TranscriptProps extends ChatContentProps {
}

const Transcript = ({ state, messages, busy, post }: TranscriptProps) => {
const retryableMessage = busy ? undefined : getRetryableMessage(messages);
const transcript = useRef<HTMLElement>(null);
const content = useRef<HTMLDivElement>(null);
const followOutput = useRef(true);
Expand Down Expand Up @@ -78,7 +79,15 @@ const Transcript = ({ state, messages, busy, post }: TranscriptProps) => {
<div ref=${content} class="transcript-content">
${messages.length === 0
? html`<${EmptyChat} state=${state} busy=${busy} post=${post} />`
: messages.map((message) => html`<${MessageView} key=${message.id} message=${message} post=${post} />`)}
: messages.map(
(message) =>
html`<${MessageView}
key=${message.id}
message=${message}
canRetry=${message === retryableMessage}
post=${post}
/>`,
)}
</div>
</section>`;
};
Expand Down
Loading
Loading