iBoss/bolt-diy
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

feat: improve prompt, add ability to abort streaming, improve message parser

Browse Source
This commit is contained in:
Dominic Elm
2024-07-18 11:10:12 +02:00
parent 637ad2b870
commit 012b5bae80
12 changed files with 633 additions and 160 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
packages/bolt/app/lib/.server/llm/prompts.ts
View File

@@ -10,24 +10,13 @@ You are Bolt, an expert AI assistant and exceptional senior software developer w
IMPORTANT: Git is NOT available.
Available shell commands: ['cat','chmod','cp','echo','hostname','kill','ln','ls','mkdir','mv','ps','pwd','rm','rmdir','xxd','alias','cd','clear','curl','env','false','getconf','head','sort','tail','touch','true','uptime','which','code','jq','loadenv','node','python3','wasm','xdg-open','command','exit','export','source']
Available shell commands: cat, chmod, cp, echo, hostname, kill, ln, ls, mkdir, mv, ps, pwd, rm, rmdir, xxd, alias, cd, clear, curl, env, false, getconf, head, sort, tail, touch, true, uptime, which, code, jq, loadenv, node, python3, wasm, xdg-open, command, exit, export, source
</system_constraints>
<code_formatting_info>
Use 2 spaces for code indentation
</code_formatting_info>
<best_practices>
Follow coding best practices:
- Ensure code is clean, readable, and maintainable.
- Adhere to proper naming conventions and consistent formatting.
Modularize functionality:
- Split functionality into smaller, reusable modules instead of placing everything in a single large file.
- Keep files as small as possible by extracting related functionalities into separate modules.
- Use imports to connect these modules together effectively.
</best_practices>
<artifact_info>
Bolt creates a SINGLE, comprehensive artifact for each project. The artifact contains all necessary steps and components, including:
@@ -67,6 +56,16 @@ You are Bolt, an expert AI assistant and exceptional senior software developer w
10. Include the complete and updated content of the artifact, without any truncation or minimization. Don't use "// rest of the code remains the same...".
11. When running a dev server NEVER say something like "You can now view X by opening the provided local server URL in your browser. The preview will be opened automatically or by the user manually!
12. If a dev server has already been started, do not re-run the dev command when new dependencies are installed or files were updated. Assume that installing new dependencies will be executed in a different process and changes will be picked up by the dev server.
13. ULTRA IMPORTANT: Use coding best practices and split functionality into smaller modules instead of putting everything in a single gigantic file. Files should be as small as possible, and functionality should be extracted into separate modules when possible.
- Ensure code is clean, readable, and maintainable.
- Adhere to proper naming conventions and consistent formatting.
- Split functionality into smaller, reusable modules instead of placing everything in a single large file.
- Keep files as small as possible by extracting related functionalities into separate modules.
- Use imports to connect these modules together effectively.
</artifact_instructions>
</artifact_info>

16
packages/bolt/app/lib/hooks/useMessageParser.ts
View File

@@ -19,8 +19,20 @@ const messageParser = new StreamingMessageParser({
workbenchStore.updateArtifact(data, { closed: true });
},
onAction: (data) => {
logger.debug('onAction', data);
onActionOpen: (data) => {
logger.debug('onActionOpen', data.action);
// we only add shell actions when when the close tag got parsed because only then we have the content
if (data.action.type !== 'shell') {
workbenchStore.addAction(data);
}
},
onActionClose: (data) => {
logger.debug('onActionClose', data.action);
if (data.action.type === 'shell') {
workbenchStore.addAction(data);
}
workbenchStore.runAction(data);
},

220
packages/bolt/app/lib/runtime/__snapshots__/message-parser.spec.ts.snap Normal file
View File

@@ -0,0 +1,220 @@
// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
exports[`StreamingMessageParser > valid artifacts with actions > should correctly parse chunks and strip out bolt artifacts (0) > onActionClose 1`] = `
{
"action": {
"content": "npm install",
"type": "shell",
},
"actionId": "0",
"artifactId": "artifact_1",
"messageId": "message_1",
}
`;
exports[`StreamingMessageParser > valid artifacts with actions > should correctly parse chunks and strip out bolt artifacts (0) > onActionOpen 1`] = `
{
"action": {
"content": "",
"type": "shell",
},
"actionId": "0",
"artifactId": "artifact_1",
"messageId": "message_1",
}
`;
exports[`StreamingMessageParser > valid artifacts with actions > should correctly parse chunks and strip out bolt artifacts (0) > onArtifactClose 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts with actions > should correctly parse chunks and strip out bolt artifacts (0) > onArtifactOpen 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts with actions > should correctly parse chunks and strip out bolt artifacts (1) > onActionClose 1`] = `
{
"action": {
"content": "npm install",
"type": "shell",
},
"actionId": "0",
"artifactId": "artifact_1",
"messageId": "message_1",
}
`;
exports[`StreamingMessageParser > valid artifacts with actions > should correctly parse chunks and strip out bolt artifacts (1) > onActionClose 2`] = `
{
"action": {
"content": "some content
",
"filePath": "index.js",
"type": "file",
},
"actionId": "1",
"artifactId": "artifact_1",
"messageId": "message_1",
}
`;
exports[`StreamingMessageParser > valid artifacts with actions > should correctly parse chunks and strip out bolt artifacts (1) > onActionOpen 1`] = `
{
"action": {
"content": "",
"type": "shell",
},
"actionId": "0",
"artifactId": "artifact_1",
"messageId": "message_1",
}
`;
exports[`StreamingMessageParser > valid artifacts with actions > should correctly parse chunks and strip out bolt artifacts (1) > onActionOpen 2`] = `
{
"action": {
"content": "",
"filePath": "index.js",
"type": "file",
},
"actionId": "1",
"artifactId": "artifact_1",
"messageId": "message_1",
}
`;
exports[`StreamingMessageParser > valid artifacts with actions > should correctly parse chunks and strip out bolt artifacts (1) > onArtifactClose 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts with actions > should correctly parse chunks and strip out bolt artifacts (1) > onArtifactOpen 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts without actions > should correctly parse chunks and strip out bolt artifacts (0) > onArtifactClose 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts without actions > should correctly parse chunks and strip out bolt artifacts (0) > onArtifactOpen 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts without actions > should correctly parse chunks and strip out bolt artifacts (1) > onArtifactClose 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts without actions > should correctly parse chunks and strip out bolt artifacts (1) > onArtifactOpen 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts without actions > should correctly parse chunks and strip out bolt artifacts (2) > onArtifactClose 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts without actions > should correctly parse chunks and strip out bolt artifacts (2) > onArtifactOpen 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts without actions > should correctly parse chunks and strip out bolt artifacts (3) > onArtifactClose 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts without actions > should correctly parse chunks and strip out bolt artifacts (3) > onArtifactOpen 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts without actions > should correctly parse chunks and strip out bolt artifacts (4) > onArtifactClose 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts without actions > should correctly parse chunks and strip out bolt artifacts (4) > onArtifactOpen 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts without actions > should correctly parse chunks and strip out bolt artifacts (5) > onArtifactClose 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts without actions > should correctly parse chunks and strip out bolt artifacts (5) > onArtifactOpen 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts without actions > should correctly parse chunks and strip out bolt artifacts (6) > onArtifactClose 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;
exports[`StreamingMessageParser > valid artifacts without actions > should correctly parse chunks and strip out bolt artifacts (6) > onArtifactOpen 1`] = `
{
"id": "artifact_1",
"messageId": "message_1",
"title": "Some title",
}
`;

257
packages/bolt/app/lib/runtime/message-parser.spec.ts
View File

@@ -1,5 +1,15 @@
import { describe, expect, it } from 'vitest';
import { StreamingMessageParser } from './message-parser';
import { describe, expect, it, vi } from 'vitest';
import { StreamingMessageParser, type ActionCallback, type ArtifactCallback } from './message-parser';
interface ExpectedResult {
output: string;
callbacks?: {
onArtifactOpen?: number;
onArtifactClose?: number;
onActionOpen?: number;
onActionClose?: number;
};
}
describe('StreamingMessageParser', () => {
it('should pass through normal text', () => {
@@ -12,75 +22,186 @@ describe('StreamingMessageParser', () => {
expect(parser.parse('test_id', 'Hello <strong>world</strong>!')).toBe('Hello <strong>world</strong>!');
});
it.each([
['Foo bar', 'Foo bar'],
['Foo bar <', 'Foo bar '],
['Foo bar <p', 'Foo bar <p'],
['Foo bar <b', 'Foo bar '],
['Foo bar <ba', 'Foo bar <ba'],
['Foo bar <bol', 'Foo bar '],
['Foo bar <bolt', 'Foo bar '],
['Foo bar <bolta', 'Foo bar <bolta'],
['Foo bar <boltA', 'Foo bar '],
['Some text before <boltArtifact>foo</boltArtifact> Some more text', 'Some text before Some more text'],
[['Some text before <boltArti', 'fact>foo</boltArtifact> Some more text'], 'Some text before Some more text'],
[['Some text before <boltArti', 'fac', 't>foo</boltArtifact> Some more text'], 'Some text before Some more text'],
[['Some text before <boltArti', 'fact>fo', 'o</boltArtifact> Some more text'], 'Some text before Some more text'],
[
['Some text before <boltArti', 'fact>fo', 'o', '<', '/boltArtifact> Some more text'],
'Some text before Some more text',
],
[
['Some text before <boltArti', 'fact>fo', 'o<', '/boltArtifact> Some more text'],
'Some text before Some more text',
],
['Before <oltArtfiact>foo</boltArtifact> After', 'Before <oltArtfiact>foo</boltArtifact> After'],
['Before <boltArtifactt>foo</boltArtifact> After', 'Before <boltArtifactt>foo</boltArtifact> After'],
['Before <boltArtifact title="Some title">foo</boltArtifact> After', 'Before After'],
[
'Before <boltArtifact title="Some title" id="artifact_1"><boltAction type="shell">npm install</boltAction></boltArtifact> After',
'Before After',
[{ type: 'shell', content: 'npm install' }],
],
[
'Before <boltArtifact title="Some title" id="artifact_1"><boltAction type="shell">npm install</boltAction><boltAction type="file" filePath="index.js">some content</boltAction></boltArtifact> After',
'Before After',
[
{ type: 'shell', content: 'npm install' },
{ type: 'file', filePath: 'index.js', content: 'some content\n' },
],
],
])('should correctly parse chunks and strip out bolt artifacts', (input, expected, expectedActions = []) => {
let actionCounter = 0;
const expectedArtifactId = 'artifact_1';
const expectedMessageId = 'message_1';
const parser = new StreamingMessageParser({
artifactElement: '',
callbacks: {
onAction: ({ artifactId, messageId, action }) => {
expect(artifactId).toBe(expectedArtifactId);
expect(messageId).toBe(expectedMessageId);
expect(action).toEqual(expectedActions[actionCounter]);
actionCounter++;
},
},
describe('no artifacts', () => {
it.each<[string | string[], ExpectedResult | string]>([
['Foo bar', 'Foo bar'],
['Foo bar <', 'Foo bar '],
['Foo bar <p', 'Foo bar <p'],
[['Foo bar <', 's', 'p', 'an>some text</span>'], 'Foo bar <span>some text</span>'],
])('should correctly parse chunks and strip out bolt artifacts (%#)', (input, expected) => {
runTest(input, expected);
});
});
let message = '';
describe('invalid or incomplete artifacts', () => {
it.each<[string | string[], ExpectedResult | string]>([
['Foo bar <b', 'Foo bar '],
['Foo bar <ba', 'Foo bar <ba'],
['Foo bar <bol', 'Foo bar '],
['Foo bar <bolt', 'Foo bar '],
['Foo bar <bolta', 'Foo bar <bolta'],
['Foo bar <boltA', 'Foo bar '],
['Foo bar <boltArtifacs></boltArtifact>', 'Foo bar <boltArtifacs></boltArtifact>'],
['Before <oltArtfiact>foo</boltArtifact> After', 'Before <oltArtfiact>foo</boltArtifact> After'],
['Before <boltArtifactt>foo</boltArtifact> After', 'Before <boltArtifactt>foo</boltArtifact> After'],
])('should correctly parse chunks and strip out bolt artifacts (%#)', (input, expected) => {
runTest(input, expected);
});
});
let result = '';
describe('valid artifacts without actions', () => {
it.each<[string | string[], ExpectedResult | string]>([
[
'Some text before <boltArtifact title="Some title" id="artifact_1">foo bar</boltArtifact> Some more text',
{
output: 'Some text before Some more text',
callbacks: { onArtifactOpen: 1, onArtifactClose: 1, onActionOpen: 0, onActionClose: 0 },
},
],
[
['Some text before <boltArti', 'fact', ' title="Some title" id="artifact_1">foo</boltArtifact> Some more text'],
{
output: 'Some text before Some more text',
callbacks: { onArtifactOpen: 1, onArtifactClose: 1, onActionOpen: 0, onActionClose: 0 },
},
],
[
[
'Some text before <boltArti',
'fac',
't title="Some title" id="artifact_1"',
' ',
'>',
'foo</boltArtifact> Some more text',
],
{
output: 'Some text before Some more text',
callbacks: { onArtifactOpen: 1, onArtifactClose: 1, onActionOpen: 0, onActionClose: 0 },
},
],
[
[
'Some text before <boltArti',
'fact',
' title="Some title" id="artifact_1"',
' >fo',
'o</boltArtifact> Some more text',
],
{
output: 'Some text before Some more text',
callbacks: { onArtifactOpen: 1, onArtifactClose: 1, onActionOpen: 0, onActionClose: 0 },
},
],
[
[
'Some text before <boltArti',
'fact tit',
'le="Some ',
'title" id="artifact_1">fo',
'o',
'<',
'/boltArtifact> Some more text',
],
{
output: 'Some text before Some more text',
callbacks: { onArtifactOpen: 1, onArtifactClose: 1, onActionOpen: 0, onActionClose: 0 },
},
],
[
[
'Some text before <boltArti',
'fact title="Some title" id="artif',
'act_1">fo',
'o<',
'/boltArtifact> Some more text',
],
{
output: 'Some text before Some more text',
callbacks: { onArtifactOpen: 1, onArtifactClose: 1, onActionOpen: 0, onActionClose: 0 },
},
],
[
'Before <boltArtifact title="Some title" id="artifact_1">foo</boltArtifact> After',
{
output: 'Before After',
callbacks: { onArtifactOpen: 1, onArtifactClose: 1, onActionOpen: 0, onActionClose: 0 },
},
],
])('should correctly parse chunks and strip out bolt artifacts (%#)', (input, expected) => {
runTest(input, expected);
});
});
const chunks = Array.isArray(input) ? input : input.split('');
for (const chunk of chunks) {
message += chunk;
result += parser.parse(expectedMessageId, message);
}
expect(actionCounter).toBe(expectedActions.length);
expect(result).toEqual(expected);
describe('valid artifacts with actions', () => {
it.each<[string | string[], ExpectedResult | string]>([
[
'Before <boltArtifact title="Some title" id="artifact_1"><boltAction type="shell">npm install</boltAction></boltArtifact> After',
{
output: 'Before After',
callbacks: { onArtifactOpen: 1, onArtifactClose: 1, onActionOpen: 1, onActionClose: 1 },
},
],
[
'Before <boltArtifact title="Some title" id="artifact_1"><boltAction type="shell">npm install</boltAction><boltAction type="file" filePath="index.js">some content</boltAction></boltArtifact> After',
{
output: 'Before After',
callbacks: { onArtifactOpen: 1, onArtifactClose: 1, onActionOpen: 2, onActionClose: 2 },
},
],
])('should correctly parse chunks and strip out bolt artifacts (%#)', (input, expected) => {
runTest(input, expected);
});
});
});
function runTest(input: string | string[], outputOrExpectedResult: string | ExpectedResult) {
let expected: ExpectedResult;
if (typeof outputOrExpectedResult === 'string') {
expected = { output: outputOrExpectedResult };
} else {
expected = outputOrExpectedResult;
}
const callbacks = {
onArtifactOpen: vi.fn<ArtifactCallback>((data) => {
expect(data).toMatchSnapshot('onArtifactOpen');
}),
onArtifactClose: vi.fn<ArtifactCallback>((data) => {
expect(data).toMatchSnapshot('onArtifactClose');
}),
onActionOpen: vi.fn<ActionCallback>((data) => {
expect(data).toMatchSnapshot('onActionOpen');
}),
onActionClose: vi.fn<ActionCallback>((data) => {
expect(data).toMatchSnapshot('onActionClose');
}),
};
const parser = new StreamingMessageParser({
artifactElement: '',
callbacks,
});
let message = '';
let result = '';
const chunks = Array.isArray(input) ? input : input.split('');
for (const chunk of chunks) {
message += chunk;
result += parser.parse('message_1', message);
}
for (const name in expected.callbacks) {
const callbackName = name;
expect(callbacks[callbackName as keyof typeof callbacks]).toHaveBeenCalledTimes(
expected.callbacks[callbackName as keyof typeof expected.callbacks] ?? 0,
);
}
expect(result).toEqual(expected.output);
}

86
packages/bolt/app/lib/runtime/message-parser.ts
View File

@@ -21,20 +21,20 @@ export interface ActionCallbackData {
action: BoltAction;
}
type ArtifactOpenCallback = (data: ArtifactCallbackData) => void;
type ArtifactCloseCallback = (data: ArtifactCallbackData) => void;
type ActionCallback = (data: ActionCallbackData) => void;
export type ArtifactCallback = (data: ArtifactCallbackData) => void;
export type ActionCallback = (data: ActionCallbackData) => void;
interface Callbacks {
onArtifactOpen?: ArtifactOpenCallback;
onArtifactClose?: ArtifactCloseCallback;
onAction?: ActionCallback;
export interface ParserCallbacks {
onArtifactOpen?: ArtifactCallback;
onArtifactClose?: ArtifactCallback;
onActionOpen?: ActionCallback;
onActionClose?: ActionCallback;
}
type ElementFactory = () => string;
interface StreamingMessageParserOptions {
callbacks?: Callbacks;
export interface StreamingMessageParserOptions {
callbacks?: ParserCallbacks;
artifactElement?: string | ElementFactory;
}
@@ -95,10 +95,17 @@ export class StreamingMessageParser {
currentAction.content = content;
this._options.callbacks?.onAction?.({
this._options.callbacks?.onActionClose?.({
artifactId: currentArtifact.id,
messageId,
actionId: String(state.actionId++),
/**
* We decrement the id because it's been incremented already
* when `onActionOpen` was emitted to make sure the ids are
* the same.
*/
actionId: String(state.actionId - 1),
action: currentAction as BoltAction,
});
@@ -117,31 +124,17 @@ export class StreamingMessageParser {
const actionEndIndex = input.indexOf('>', actionOpenIndex);
if (actionEndIndex !== -1) {
const actionTag = input.slice(actionOpenIndex, actionEndIndex + 1);
const actionType = this.#extractAttribute(actionTag, 'type') as ActionType;
const actionAttributes = {
type: actionType,
content: '',
};
if (actionType === 'file') {
const filePath = this.#extractAttribute(actionTag, 'filePath') as string;
if (!filePath) {
logger.debug('File path not specified');
}
(actionAttributes as FileAction).filePath = filePath;
} else if (actionType !== 'shell') {
logger.warn(`Unknown action type '${actionType}'`);
}
state.currentAction = actionAttributes as FileAction | ShellAction;
state.insideAction = true;
state.currentAction = this.#parseActionTag(input, actionOpenIndex, actionEndIndex);
this._options.callbacks?.onActionOpen?.({
artifactId: currentArtifact.id,
messageId,
actionId: String(state.actionId++),
action: state.currentAction as BoltAction,
});
i = actionEndIndex + 1;
} else {
break;
@@ -241,6 +234,31 @@ export class StreamingMessageParser {
this.#messages.clear();
}
#parseActionTag(input: string, actionOpenIndex: number, actionEndIndex: number) {
const actionTag = input.slice(actionOpenIndex, actionEndIndex + 1);
const actionType = this.#extractAttribute(actionTag, 'type') as ActionType;
const actionAttributes = {
type: actionType,
content: '',
};
if (actionType === 'file') {
const filePath = this.#extractAttribute(actionTag, 'filePath') as string;
if (!filePath) {
logger.debug('File path not specified');
}
(actionAttributes as FileAction).filePath = filePath;
} else if (actionType !== 'shell') {
logger.warn(`Unknown action type '${actionType}'`);
}
return actionAttributes as FileAction | ShellAction;
}
#extractAttribute(tag: string, attributeName: string): string | undefined {
const match = tag.match(new RegExp(`${attributeName}="([^"]*)"`, 'i'));
return match ? match[1] : undefined;

5
packages/bolt/app/lib/stores/chat.ts Normal file
View File

@@ -0,0 +1,5 @@
import { map } from 'nanostores';
export const chatStore = map({
aborted: false,
});

110
packages/bolt/app/lib/stores/workbench.ts
View File

@@ -4,25 +4,30 @@ import { unreachable } from '../../utils/unreachable';
import { ActionRunner } from '../runtime/action-runner';
import type { ActionCallbackData, ArtifactCallbackData } from '../runtime/message-parser';
import { webcontainer } from '../webcontainer';
import { chatStore } from './chat';
import { PreviewsStore } from './previews';
export type RunningState = BoltAction & {
const MIN_SPINNER_TIME = 200;
export type BaseActionState = BoltAction & {
status: 'running' | 'complete' | 'pending' | 'aborted';
executing: boolean;
abort?: () => void;
};
export type FailedState = BoltAction & {
status: 'failed';
error: string;
abort?: () => void;
};
export type FailedActionState = BoltAction &
Omit<BaseActionState, 'status'> & {
status: 'failed';
error: string;
};
export type ActionState = RunningState | FailedState;
export type ActionState = BaseActionState | FailedActionState;
type BaseActionUpdate = Partial<Pick<BaseActionState, 'status' | 'executing' | 'abort'>>;
export type ActionStateUpdate =
| { status: 'running' | 'complete' | 'pending' | 'aborted'; abort?: () => void }
| { status: 'failed'; error: string; abort?: () => void }
| { abort?: () => void };
| BaseActionUpdate
| (Omit<BaseActionUpdate, 'status'> & { status: 'failed'; error: string });
export interface ArtifactState {
title: string;
@@ -49,6 +54,16 @@ export class WorkbenchStore {
this.showWorkbench.set(show);
}
abortAllActions() {
for (const [, artifact] of Object.entries(this.artifacts.get())) {
for (const [, action] of Object.entries(artifact.actions.get())) {
if (action.status === 'running') {
action.abort?.();
}
}
}
}
addArtifact({ id, messageId, title }: ArtifactCallbackData) {
const artifacts = this.artifacts.get();
const artifactKey = getArtifactKey(id, messageId);
@@ -78,7 +93,7 @@ export class WorkbenchStore {
this.artifacts.setKey(key, { ...artifact, ...state });
}
async runAction(data: ActionCallbackData) {
async addAction(data: ActionCallbackData) {
const { artifactId, messageId, actionId } = data;
const artifacts = this.artifacts.get();
@@ -96,33 +111,70 @@ export class WorkbenchStore {
return;
}
artifact.actions.setKey(actionId, { ...data.action, status: 'pending' });
artifact.actions.setKey(actionId, { ...data.action, status: 'pending', executing: false });
artifact.currentActionPromise.then(() => {
if (chatStore.get().aborted) {
return;
}
this.#updateAction(key, actionId, { status: 'running' });
});
}
async runAction(data: ActionCallbackData) {
const { artifactId, messageId, actionId } = data;
const artifacts = this.artifacts.get();
const key = getArtifactKey(artifactId, messageId);
const artifact = artifacts[key];
if (!artifact) {
unreachable('Artifact not found');
}
const actions = artifact.actions.get();
const action = actions[actionId];
if (!action) {
unreachable('Expected action to exist');
}
if (action.executing || action.status === 'complete' || action.status === 'failed' || action.status === 'aborted') {
return;
}
artifact.currentActionPromise = artifact.currentActionPromise.then(async () => {
if (chatStore.get().aborted) {
return;
}
const abortController = new AbortController();
this.#updateAction(key, actionId, {
status: 'running',
executing: true,
abort: () => {
abortController.abort();
this.#updateAction(key, actionId, { status: 'aborted' });
},
});
try {
let abortController: AbortController | undefined;
await Promise.all([
this.#actionRunner.runAction(data, abortController.signal),
new Promise((resolve) => setTimeout(resolve, MIN_SPINNER_TIME)),
]);
if (data.action.type === 'shell') {
abortController = new AbortController();
if (!abortController.signal.aborted) {
this.#updateAction(key, actionId, { status: 'complete' });
}
let aborted = false;
this.#updateAction(key, actionId, {
status: 'running',
abort: () => {
aborted = true;
abortController?.abort();
},
});
await this.#actionRunner.runAction(data, abortController?.signal);
this.#updateAction(key, actionId, { status: aborted ? 'aborted' : 'complete' });
} catch (error) {
this.#updateAction(key, actionId, { status: 'failed', error: 'Action failed' });
throw error;
} finally {
this.#updateAction(key, actionId, { executing: false });
}
});
}