From 012b5bae8046f208639f54f15631e8a24f53a419 Mon Sep 17 00:00:00 2001 From: Dominic Elm Date: Thu, 18 Jul 2024 11:10:12 +0200 Subject: [PATCH] feat: improve prompt, add ability to abort streaming, improve message parser --- .../bolt/app/components/chat/Artifact.tsx | 28 +- .../bolt/app/components/chat/BaseChat.tsx | 19 +- .../bolt/app/components/chat/Chat.client.tsx | 13 +- .../app/components/chat/SendButton.client.tsx | 7 +- .../bolt/app/components/workbench/Preview.tsx | 9 +- packages/bolt/app/lib/.server/llm/prompts.ts | 23 +- .../bolt/app/lib/hooks/useMessageParser.ts | 16 +- .../__snapshots__/message-parser.spec.ts.snap | 220 +++++++++++++++ .../app/lib/runtime/message-parser.spec.ts | 257 +++++++++++++----- .../bolt/app/lib/runtime/message-parser.ts | 86 +++--- packages/bolt/app/lib/stores/chat.ts | 5 + packages/bolt/app/lib/stores/workbench.ts | 110 ++++++-- 12 files changed, 633 insertions(+), 160 deletions(-) create mode 100644 packages/bolt/app/lib/runtime/__snapshots__/message-parser.spec.ts.snap create mode 100644 packages/bolt/app/lib/stores/chat.ts diff --git a/packages/bolt/app/components/chat/Artifact.tsx b/packages/bolt/app/components/chat/Artifact.tsx index 781b1df..130b184 100644 --- a/packages/bolt/app/components/chat/Artifact.tsx +++ b/packages/bolt/app/components/chat/Artifact.tsx @@ -1,8 +1,9 @@ import { useStore } from '@nanostores/react'; import { AnimatePresence, motion } from 'framer-motion'; import { computed } from 'nanostores'; -import { useState } from 'react'; +import { memo, useEffect, useRef, useState } from 'react'; import { createHighlighter, type BundledLanguage, type BundledTheme, type HighlighterGeneric } from 'shiki'; +import { chatStore } from '../../lib/stores/chat'; import { getArtifactKey, workbenchStore, type ActionState } from '../../lib/stores/workbench'; import { classNames } from '../../utils/classNames'; import { cubicEasingFn } from '../../utils/easings'; @@ -25,9 +26,11 @@ interface ArtifactProps { messageId: string; } -export function Artifact({ artifactId, messageId }: ArtifactProps) { +export const Artifact = memo(({ artifactId, messageId }: ArtifactProps) => { + const userToggledActions = useRef(false); const [showActions, setShowActions] = useState(false); + const chat = useStore(chatStore); const artifacts = useStore(workbenchStore.artifacts); const artifact = artifacts[getArtifactKey(artifactId, messageId)]; @@ -37,6 +40,17 @@ export function Artifact({ artifactId, messageId }: ArtifactProps) { }), ); + const toggleActions = () => { + userToggledActions.current = true; + setShowActions(!showActions); + }; + + useEffect(() => { + if (actions.length && !showActions && !userToggledActions.current) { + setShowActions(true); + } + }, [actions]); + return (
@@ -48,7 +62,7 @@ export function Artifact({ artifactId, messageId }: ArtifactProps) { }} >
- {!artifact?.closed ? ( + {!artifact?.closed && !chat.aborted ? (
) : (
@@ -67,7 +81,7 @@ export function Artifact({ artifactId, messageId }: ArtifactProps) { exit={{ width: 0 }} transition={{ duration: 0.15, ease: cubicEasingFn }} className="hover:bg-gray-200" - onClick={() => setShowActions(!showActions)} + onClick={toggleActions} >
@@ -98,9 +112,9 @@ export function Artifact({ artifactId, messageId }: ArtifactProps) { const { status, type, content, abort } = action; return ( -
  • +
  • -
    +
    {status === 'running' ? (
    ) : status === 'pending' ? ( @@ -136,7 +150,7 @@ export function Artifact({ artifactId, messageId }: ArtifactProps) {
    ); -} +}); function getTextColor(status: ActionState['status']) { switch (status) { diff --git a/packages/bolt/app/components/chat/BaseChat.tsx b/packages/bolt/app/components/chat/BaseChat.tsx index e61770d..2177aa7 100644 --- a/packages/bolt/app/components/chat/BaseChat.tsx +++ b/packages/bolt/app/components/chat/BaseChat.tsx @@ -16,6 +16,7 @@ interface BaseChatProps { enhancingPrompt?: boolean; promptEnhanced?: boolean; input?: string; + handleStop?: () => void; sendMessage?: () => void; handleInputChange?: (event: React.ChangeEvent) => void; enhancePrompt?: () => void; @@ -38,6 +39,7 @@ export const BaseChat = React.forwardRef( sendMessage, handleInputChange, enhancePrompt, + handleStop, }, ref, ) => { @@ -111,7 +113,22 @@ export const BaseChat = React.forwardRef( placeholder="How can Bolt help you today?" translate="no" /> - {() => 0} onClick={sendMessage} />} + + {() => ( + 0 || isStreaming} + isStreaming={isStreaming} + onClick={() => { + if (isStreaming) { + handleStop?.(); + return; + } + + sendMessage?.(); + }} + /> + )} +
    diff --git a/packages/bolt/app/components/chat/Chat.client.tsx b/packages/bolt/app/components/chat/Chat.client.tsx index 85168c5..854a1fd 100644 --- a/packages/bolt/app/components/chat/Chat.client.tsx +++ b/packages/bolt/app/components/chat/Chat.client.tsx @@ -2,6 +2,8 @@ import { useChat } from 'ai/react'; import { useAnimate } from 'framer-motion'; import { useEffect, useRef, useState } from 'react'; import { useMessageParser, usePromptEnhancer } from '../../lib/hooks'; +import { chatStore } from '../../lib/stores/chat'; +import { workbenchStore } from '../../lib/stores/workbench'; import { cubicEasingFn } from '../../utils/easings'; import { createScopedLogger } from '../../utils/logger'; import { BaseChat } from './BaseChat'; @@ -15,7 +17,7 @@ export function Chat() { const [animationScope, animate] = useAnimate(); - const { messages, isLoading, input, handleInputChange, setInput, handleSubmit } = useChat({ + const { messages, isLoading, input, handleInputChange, setInput, handleSubmit, stop } = useChat({ api: '/api/chat', onError: (error) => { logger.error(error); @@ -42,6 +44,12 @@ export function Chat() { } }; + const abort = () => { + stop(); + chatStore.setKey('aborted', true); + workbenchStore.abortAllActions(); + }; + useEffect(() => { const textarea = textareaRef.current; @@ -70,6 +78,8 @@ export function Chat() { return; } + chatStore.setKey('aborted', false); + runAnimation(); handleSubmit(); resetEnhancer(); @@ -88,6 +98,7 @@ export function Chat() { promptEnhanced={promptEnhanced} sendMessage={sendMessage} handleInputChange={handleInputChange} + handleStop={abort} messages={messages.map((message, i) => { if (message.role === 'user') { return message; diff --git a/packages/bolt/app/components/chat/SendButton.client.tsx b/packages/bolt/app/components/chat/SendButton.client.tsx index e2dcce7..cf77eb8 100644 --- a/packages/bolt/app/components/chat/SendButton.client.tsx +++ b/packages/bolt/app/components/chat/SendButton.client.tsx @@ -2,12 +2,13 @@ import { AnimatePresence, cubicBezier, motion } from 'framer-motion'; interface SendButtonProps { show: boolean; + isStreaming?: boolean; onClick?: VoidFunction; } const customEasingFn = cubicBezier(0.4, 0, 0.2, 1); -export function SendButton({ show, onClick }: SendButtonProps) { +export function SendButton({ show, isStreaming, onClick }: SendButtonProps) { return ( {show ? ( @@ -22,7 +23,9 @@ export function SendButton({ show, onClick }: SendButtonProps) { onClick?.(); }} > -
    +
    + {!isStreaming ?
    :
    } +
    ) : null}     diff --git a/packages/bolt/app/components/workbench/Preview.tsx b/packages/bolt/app/components/workbench/Preview.tsx index eef5f81..0b523c0 100644 --- a/packages/bolt/app/components/workbench/Preview.tsx +++ b/packages/bolt/app/components/workbench/Preview.tsx @@ -30,12 +30,13 @@ export const Preview = memo(() => { return (
    -
    -
    -
    +
    +
    + Preview +
    -
    +
    diff --git a/packages/bolt/app/lib/.server/llm/prompts.ts b/packages/bolt/app/lib/.server/llm/prompts.ts index 1f11043..bbe6afe 100644 --- a/packages/bolt/app/lib/.server/llm/prompts.ts +++ b/packages/bolt/app/lib/.server/llm/prompts.ts @@ -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 Use 2 spaces for code indentation - - 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. - - 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. diff --git a/packages/bolt/app/lib/hooks/useMessageParser.ts b/packages/bolt/app/lib/hooks/useMessageParser.ts index ac3e703..cef1601 100644 --- a/packages/bolt/app/lib/hooks/useMessageParser.ts +++ b/packages/bolt/app/lib/hooks/useMessageParser.ts @@ -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); }, diff --git a/packages/bolt/app/lib/runtime/__snapshots__/message-parser.spec.ts.snap b/packages/bolt/app/lib/runtime/__snapshots__/message-parser.spec.ts.snap new file mode 100644 index 0000000..1543770 --- /dev/null +++ b/packages/bolt/app/lib/runtime/__snapshots__/message-parser.spec.ts.snap @@ -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", +} +`; diff --git a/packages/bolt/app/lib/runtime/message-parser.spec.ts b/packages/bolt/app/lib/runtime/message-parser.spec.ts index 92d914f..bc8080f 100644 --- a/packages/bolt/app/lib/runtime/message-parser.spec.ts +++ b/packages/bolt/app/lib/runtime/message-parser.spec.ts @@ -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 world!')).toBe('Hello world!'); }); - it.each([ - ['Foo bar', 'Foo bar'], - ['Foo bar <', 'Foo bar '], - ['Foo bar foo Some more text', 'Some text before Some more text'], - [['Some text before foo Some more text'], 'Some text before Some more text'], - [['Some text before foo Some more text'], 'Some text before Some more text'], - [['Some text before fo', 'o Some more text'], 'Some text before Some more text'], - [ - ['Some text before fo', 'o', '<', '/boltArtifact> Some more text'], - 'Some text before Some more text', - ], - [ - ['Some text before fo', 'o<', '/boltArtifact> Some more text'], - 'Some text before Some more text', - ], - ['Before foo After', 'Before foo After'], - ['Before foo After', 'Before foo After'], - ['Before foo After', 'Before After'], - [ - 'Before npm install After', - 'Before After', - [{ type: 'shell', content: 'npm install' }], - ], - [ - 'Before npm installsome content 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 some text'], 'Foo bar some text'], + ])('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 ', 'Foo bar '], + ['Before foo After', 'Before foo After'], + ['Before foo After', 'Before foo 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 foo bar Some more text', + { + output: 'Some text before Some more text', + callbacks: { onArtifactOpen: 1, onArtifactClose: 1, onActionOpen: 0, onActionClose: 0 }, + }, + ], + [ + ['Some text before foo Some more text'], + { + output: 'Some text before Some more text', + callbacks: { onArtifactOpen: 1, onArtifactClose: 1, onActionOpen: 0, onActionClose: 0 }, + }, + ], + [ + [ + 'Some text before ', + 'foo Some more text', + ], + { + output: 'Some text before Some more text', + callbacks: { onArtifactOpen: 1, onArtifactClose: 1, onActionOpen: 0, onActionClose: 0 }, + }, + ], + [ + [ + 'Some text before fo', + 'o Some more text', + ], + { + output: 'Some text before Some more text', + callbacks: { onArtifactOpen: 1, onArtifactClose: 1, onActionOpen: 0, onActionClose: 0 }, + }, + ], + [ + [ + 'Some text before 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 fo', + 'o<', + '/boltArtifact> Some more text', + ], + { + output: 'Some text before Some more text', + callbacks: { onArtifactOpen: 1, onArtifactClose: 1, onActionOpen: 0, onActionClose: 0 }, + }, + ], + [ + 'Before foo 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 npm install After', + { + output: 'Before After', + callbacks: { onArtifactOpen: 1, onArtifactClose: 1, onActionOpen: 1, onActionClose: 1 }, + }, + ], + [ + 'Before npm installsome content 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((data) => { + expect(data).toMatchSnapshot('onArtifactOpen'); + }), + onArtifactClose: vi.fn((data) => { + expect(data).toMatchSnapshot('onArtifactClose'); + }), + onActionOpen: vi.fn((data) => { + expect(data).toMatchSnapshot('onActionOpen'); + }), + onActionClose: vi.fn((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); +} diff --git a/packages/bolt/app/lib/runtime/message-parser.ts b/packages/bolt/app/lib/runtime/message-parser.ts index bba4ec7..a607759 100644 --- a/packages/bolt/app/lib/runtime/message-parser.ts +++ b/packages/bolt/app/lib/runtime/message-parser.ts @@ -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; diff --git a/packages/bolt/app/lib/stores/chat.ts b/packages/bolt/app/lib/stores/chat.ts new file mode 100644 index 0000000..d5f3a37 --- /dev/null +++ b/packages/bolt/app/lib/stores/chat.ts @@ -0,0 +1,5 @@ +import { map } from 'nanostores'; + +export const chatStore = map({ + aborted: false, +}); diff --git a/packages/bolt/app/lib/stores/workbench.ts b/packages/bolt/app/lib/stores/workbench.ts index 0bdfaba..3b457f8 100644 --- a/packages/bolt/app/lib/stores/workbench.ts +++ b/packages/bolt/app/lib/stores/workbench.ts @@ -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 & { + status: 'failed'; + error: string; + }; -export type ActionState = RunningState | FailedState; +export type ActionState = BaseActionState | FailedActionState; + +type BaseActionUpdate = Partial>; export type ActionStateUpdate = - | { status: 'running' | 'complete' | 'pending' | 'aborted'; abort?: () => void } - | { status: 'failed'; error: string; abort?: () => void } - | { abort?: () => void }; + | BaseActionUpdate + | (Omit & { 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 }); } }); }