Show optimistic chat for new workspaces

This commit is contained in:
Mohamed Boudra
2026-05-25 21:26:31 +07:00
parent d271597ee8
commit 30a8bffc24
11 changed files with 452 additions and 140 deletions
@@ -0,0 +1,90 @@
/**
* @vitest-environment jsdom
*/
import { act, renderHook } from "@testing-library/react";
import { beforeEach, describe, expect, it, vi } from "vitest";
import { useCreateFlowStore } from "@/stores/create-flow-store";
import type { UserMessageImageAttachment } from "@/types/stream";
import type { AgentAttachment } from "@server/shared/messages";
import { useDraftAgentCreateFlow, type DraftCreateAttempt } from "./use-draft-agent-create-flow";
describe("useDraftAgentCreateFlow", () => {
beforeEach(() => {
useCreateFlowStore.setState({ pendingByDraftId: {} });
});
it("renders a prepared new-workspace create attempt as optimistic chat before continuing it", async () => {
const image: UserMessageImageAttachment = {
id: "image-1",
mimeType: "image/png",
storageType: "web-indexeddb",
storageKey: "image-key",
createdAt: 123,
};
const attachment = {
type: "review",
cwd: "/repo",
summary: "Review",
} as unknown as AgentAttachment;
const attempt: DraftCreateAttempt = {
clientMessageId: "msg-prepared",
text: "build this",
timestamp: new Date("2026-05-25T00:00:00.000Z"),
images: [image],
attachments: [attachment],
};
const createRequest = vi.fn(
async (ctx: {
attempt: DraftCreateAttempt;
text: string;
images?: UserMessageImageAttachment[];
attachments?: AgentAttachment[];
cwd: string;
}) => ({
agentId: "agent-1",
result: { id: "agent-1", ctx },
}),
);
const onCreateSuccess = vi.fn();
const { result } = renderHook(() =>
useDraftAgentCreateFlow({
draftId: "draft-1",
getPendingServerId: () => "server-1",
initialAttempt: attempt,
buildDraftAgent: (currentAttempt) => ({ currentAttempt }),
createRequest,
onCreateSuccess,
}),
);
expect(result.current.isSubmitting).toBe(true);
expect(result.current.draftAgent).toEqual({ currentAttempt: attempt });
expect(result.current.optimisticStreamItems).toEqual([
{
kind: "user_message",
id: "msg-prepared",
text: "build this",
timestamp: attempt.timestamp,
optimistic: true,
images: [image],
attachments: [attachment],
},
]);
expect(createRequest).not.toHaveBeenCalled();
await act(async () => {
await result.current.continueCreateFromAttempt({ attempt, cwd: "/repo" });
});
expect(createRequest).toHaveBeenCalledTimes(1);
expect(createRequest).toHaveBeenCalledWith({
attempt,
text: "build this",
images: [image],
attachments: [attachment],
cwd: "/repo",
});
expect(onCreateSuccess).toHaveBeenCalledTimes(1);
});
});
@@ -81,6 +81,7 @@ interface CreateRequestContext {
interface UseDraftAgentCreateFlowOptions<TDraftAgent, TCreateResult> {
draftId: string;
getPendingServerId: () => string | null;
initialAttempt?: CreateAttempt | null;
allowEmptyText?: boolean;
validateBeforeSubmit?: (ctx: SubmitContext) => string | null;
onBeforeSubmit?: (ctx: CreateRequestContext) => void;
@@ -94,6 +95,7 @@ interface UseDraftAgentCreateFlowOptions<TDraftAgent, TCreateResult> {
export function useDraftAgentCreateFlow<TDraftAgent, TCreateResult>({
draftId,
getPendingServerId,
initialAttempt = null,
allowEmptyText = false,
validateBeforeSubmit,
onBeforeSubmit,
@@ -103,10 +105,17 @@ export function useDraftAgentCreateFlow<TDraftAgent, TCreateResult>({
onCreateSuccess,
onCreateError,
}: UseDraftAgentCreateFlowOptions<TDraftAgent, TCreateResult>) {
const [machine, dispatch] = useReducer(reducer, {
tag: "draft",
errorMessage: "",
} as DraftAgentMachineState);
const [machine, dispatch] = useReducer(
reducer,
initialAttempt,
(attempt): DraftAgentMachineState =>
attempt
? { tag: "creating", attempt }
: {
tag: "draft",
errorMessage: "",
},
);
const setPendingCreateAttempt = useCreateFlowStore((state) => state.setPending);
const updatePendingAgentId = useCreateFlowStore((state) => state.updateAgentId);
@@ -150,6 +159,73 @@ export function useDraftAgentCreateFlow<TDraftAgent, TCreateResult>({
return buildDraftAgent(machine.attempt);
}, [buildDraftAgent, machine]);
const runCreateAttempt = useCallback(
async ({ attempt, cwd }: { attempt: CreateAttempt; cwd: string }) => {
const pendingServerId = getPendingServerId();
if (!pendingServerId) {
const error = new Error("No host selected");
dispatch({ type: "DRAFT_SET_ERROR", message: error.message });
throw error;
}
onBeforeSubmit?.({
attempt,
text: attempt.text,
images: attempt.images,
attachments: attempt.attachments,
cwd,
});
try {
const createResult = await createRequest({
attempt,
text: attempt.text,
images: attempt.images,
attachments: attempt.attachments,
cwd,
});
if (createResult.agentId) {
updatePendingAgentId({ draftId, agentId: createResult.agentId });
appendOptimisticUserMessageToAgentStream(
pendingServerId,
createResult.agentId,
buildOptimisticUserMessage({
id: attempt.clientMessageId,
text: attempt.text,
timestamp: attempt.timestamp,
images: attempt.images,
attachments: attempt.attachments,
}),
{ placement: "tail", skipIfUserMessageExists: true },
);
markPendingCreateLifecycle({ draftId, lifecycle: "sent" });
}
await onCreateSuccess({ result: createResult.result, attempt });
} catch (error) {
const resolved = error instanceof Error ? error : new Error("Failed to create agent");
dispatch({ type: "CREATE_FAILED", message: resolved.message });
markPendingCreateLifecycle({ draftId, lifecycle: "abandoned" });
clearPendingCreateAttempt({ draftId });
onCreateError?.(resolved);
throw error;
}
},
[
appendOptimisticUserMessageToAgentStream,
clearPendingCreateAttempt,
createRequest,
draftId,
getPendingServerId,
markPendingCreateLifecycle,
onBeforeSubmit,
onCreateError,
onCreateSuccess,
updatePendingAgentId,
],
);
const handleCreateFromInput = useCallback(
async ({ text, attachments, cwd }: SubmitContext) => {
if (isSubmitting) {
@@ -206,71 +282,32 @@ export function useDraftAgentCreateFlow<TDraftAgent, TCreateResult>({
: {}),
});
onBeforeSubmit?.({
attempt,
text: trimmedPrompt,
images,
attachments: wirePayload.attachments,
cwd,
});
dispatch({ type: "SUBMIT", attempt });
onCreateStart?.();
try {
const createResult = await createRequest({
attempt,
text: trimmedPrompt,
images,
attachments: wirePayload.attachments,
cwd,
});
if (createResult.agentId) {
updatePendingAgentId({ draftId, agentId: createResult.agentId });
appendOptimisticUserMessageToAgentStream(
pendingServerId,
createResult.agentId,
buildOptimisticUserMessage({
id: attempt.clientMessageId,
text: attempt.text,
timestamp: attempt.timestamp,
images: attempt.images,
attachments: attempt.attachments,
}),
{ placement: "tail", skipIfUserMessageExists: true },
);
markPendingCreateLifecycle({ draftId, lifecycle: "sent" });
}
await onCreateSuccess({ result: createResult.result, attempt });
} catch (error) {
const resolved = error instanceof Error ? error : new Error("Failed to create agent");
dispatch({ type: "CREATE_FAILED", message: resolved.message });
markPendingCreateLifecycle({ draftId, lifecycle: "abandoned" });
clearPendingCreateAttempt({ draftId });
onCreateError?.(resolved);
throw error;
}
await runCreateAttempt({ attempt, cwd });
},
[
appendOptimisticUserMessageToAgentStream,
clearPendingCreateAttempt,
createRequest,
allowEmptyText,
draftId,
getPendingServerId,
isSubmitting,
markPendingCreateLifecycle,
onBeforeSubmit,
onCreateError,
onCreateStart,
onCreateSuccess,
runCreateAttempt,
setPendingCreateAttempt,
updatePendingAgentId,
validateBeforeSubmit,
allowEmptyText,
],
);
const continueCreateFromAttempt = useCallback(
async ({ attempt, cwd }: { attempt: CreateAttempt; cwd: string }) => {
if (!isSubmitting) {
dispatch({ type: "SUBMIT", attempt });
}
await runCreateAttempt({ attempt, cwd });
},
[isSubmitting, runCreateAttempt],
);
return {
machine,
formErrorMessage,
@@ -278,6 +315,7 @@ export function useDraftAgentCreateFlow<TDraftAgent, TCreateResult>({
optimisticStreamItems,
draftAgent,
handleCreateFromInput,
continueCreateFromAttempt,
};
}