Fix draft agent handoff continuity

Move optimistic user message reconciliation into the stream reducers and keep draft-created agents in a handoff lifecycle until authoritative history lands.
This commit is contained in:
Mohamed Boudra
2026-05-25 20:53:49 +07:00
parent e5658654c0
commit 53faba64ca
21 changed files with 677 additions and 458 deletions
@@ -57,13 +57,12 @@ function createAgentWithStatus({ id, status }: { id: string; status: Agent["stat
function createBaseInput(): AgentScreenMachineInput {
return {
agent: null,
placeholderAgent: null,
continuity: { kind: "none" },
missingAgentState: { kind: "idle" },
isConnected: true,
isArchivingCurrentAgent: false,
isHistorySyncing: false,
needsAuthoritativeSync: false,
shouldUseOptimisticStream: false,
hasHydratedHistoryBefore: false,
};
}
@@ -298,8 +297,7 @@ describe("deriveAgentScreenViewState", () => {
const memory = createBaseMemory();
const input: AgentScreenMachineInput = {
...createBaseInput(),
placeholderAgent: createAgent("draft-agent"),
shouldUseOptimisticStream: true,
continuity: { kind: "optimistic-create", agent: createAgent("draft-agent") },
};
const result = deriveAgentScreenViewState({ input, memory });
@@ -335,8 +333,7 @@ describe("deriveAgentScreenViewState", () => {
const input: AgentScreenMachineInput = {
...createBaseInput(),
agent: createAgentWithStatus({ id: "agent-1", status: "idle" }),
placeholderAgent: createAgent("agent-1"),
shouldUseOptimisticStream: true,
continuity: { kind: "optimistic-create", agent: createAgent("agent-1") },
needsAuthoritativeSync: true,
isHistorySyncing: true,
hasHydratedHistoryBefore: false,
@@ -347,14 +344,14 @@ describe("deriveAgentScreenViewState", () => {
expect(ready.source).toBe("optimistic");
expect(ready.agent.status).toBe("running");
expect(ready.sync).toEqual({ status: "catching_up", ui: "silent" });
});
it("keeps optimistic flow non-blocking while transitioning to authoritative stream", () => {
const initialMemory = createBaseMemory();
const optimisticInput: AgentScreenMachineInput = {
...createBaseInput(),
placeholderAgent: createAgent("draft-agent"),
shouldUseOptimisticStream: true,
continuity: { kind: "optimistic-create", agent: createAgent("draft-agent") },
};
const optimistic = deriveAgentScreenViewState({
@@ -382,8 +379,7 @@ describe("deriveAgentScreenViewState", () => {
const input: AgentScreenMachineInput = {
...createBaseInput(),
agent: createAgentWithStatus({ id: "agent-1", status: "idle" }),
placeholderAgent: createAgent("agent-1"),
shouldUseOptimisticStream: true,
continuity: { kind: "optimistic-create", agent: createAgent("agent-1") },
};
const result = deriveAgentScreenViewState({ input, memory });
@@ -402,14 +398,12 @@ describe("deriveAgentScreenViewState", () => {
{
...createBaseInput(),
agent: createAgentWithStatus({ id: "agent-1", status: "idle" }),
placeholderAgent: createAgent("agent-1"),
shouldUseOptimisticStream: true,
continuity: { kind: "optimistic-create", agent: createAgent("agent-1") },
},
{
...createBaseInput(),
agent: createAgentWithStatus({ id: "agent-1", status: "running" }),
placeholderAgent: createAgent("agent-1"),
shouldUseOptimisticStream: true,
continuity: { kind: "optimistic-create", agent: createAgent("agent-1") },
},
{
...createBaseInput(),
@@ -431,8 +425,7 @@ describe("deriveAgentScreenViewState", () => {
const input: AgentScreenMachineInput = {
...createBaseInput(),
agent: createAgentWithStatus({ id: "agent-1", status: "initializing" }),
placeholderAgent: createAgent("agent-1"),
shouldUseOptimisticStream: true,
continuity: { kind: "optimistic-create", agent: createAgent("agent-1") },
};
const result = deriveAgentScreenViewState({ input, memory });
@@ -447,8 +440,7 @@ describe("deriveAgentScreenViewState", () => {
const input: AgentScreenMachineInput = {
...createBaseInput(),
agent: createAgentWithStatus({ id: "agent-1", status: "running" }),
placeholderAgent: createAgent("agent-1"),
shouldUseOptimisticStream: true,
continuity: { kind: "optimistic-create", agent: createAgent("agent-1") },
};
const result = deriveAgentScreenViewState({ input, memory });
@@ -463,8 +455,7 @@ describe("deriveAgentScreenViewState", () => {
const input: AgentScreenMachineInput = {
...createBaseInput(),
agent: createAgentWithStatus({ id: "agent-1", status: "error" }),
placeholderAgent: createAgent("agent-1"),
shouldUseOptimisticStream: true,
continuity: { kind: "optimistic-create", agent: createAgent("agent-1") },
};
const result = deriveAgentScreenViewState({ input, memory });
@@ -479,8 +470,7 @@ describe("deriveAgentScreenViewState", () => {
const input: AgentScreenMachineInput = {
...createBaseInput(),
agent: createAgentWithStatus({ id: "agent-1", status: "closed" }),
placeholderAgent: createAgent("agent-1"),
shouldUseOptimisticStream: true,
continuity: { kind: "optimistic-create", agent: createAgent("agent-1") },
};
const result = deriveAgentScreenViewState({ input, memory });
@@ -24,19 +24,26 @@ export type AgentScreenMissingState =
export interface AgentScreenMachineInput {
agent: AgentScreenAgent | null;
placeholderAgent: AgentScreenAgent | null;
missingAgentState: AgentScreenMissingState;
isConnected: boolean;
isArchivingCurrentAgent: boolean;
isHistorySyncing: boolean;
needsAuthoritativeSync: boolean;
shouldUseOptimisticStream: boolean;
continuity: AgentScreenContinuity;
hasHydratedHistoryBefore: boolean;
}
export type AgentScreenContinuity =
| { kind: "none" }
| { kind: "optimistic-create"; agent: AgentScreenAgent };
function hasOptimisticCreateContinuity(input: AgentScreenMachineInput): boolean {
return input.continuity.kind === "optimistic-create";
}
function shouldBlockInitialAuthoritativeReadyState(input: AgentScreenMachineInput): boolean {
return (
!input.shouldUseOptimisticStream &&
!hasOptimisticCreateContinuity(input) &&
!input.hasHydratedHistoryBefore &&
(input.needsAuthoritativeSync || input.isHistorySyncing)
);
@@ -93,8 +100,7 @@ function updateInitialSyncFailureMemory(args: {
function shouldUseOptimisticCreateFlowAgent(input: AgentScreenMachineInput): boolean {
return (
input.shouldUseOptimisticStream &&
Boolean(input.placeholderAgent) &&
input.continuity.kind === "optimistic-create" &&
(!input.agent || input.agent.status === "initializing" || input.agent.status === "idle")
);
}
@@ -104,29 +110,31 @@ function resolveCandidateAgent(args: {
useOptimisticCreateFlowAgent: boolean;
}): AgentScreenAgent | null {
const { input, useOptimisticCreateFlowAgent } = args;
if (input.agent && useOptimisticCreateFlowAgent && input.placeholderAgent) {
return { ...input.agent, status: input.placeholderAgent.status };
const continuityAgent =
input.continuity.kind === "optimistic-create" ? input.continuity.agent : null;
if (input.agent && useOptimisticCreateFlowAgent && continuityAgent) {
return { ...input.agent, status: continuityAgent.status };
}
return input.agent ?? input.placeholderAgent;
return input.agent ?? continuityAgent;
}
function resolveAgentScreenSource(args: {
useOptimisticCreateFlowAgent: boolean;
hasAgent: boolean;
shouldUseOptimisticStream: boolean;
hasOptimisticCreateContinuity: boolean;
}): "authoritative" | "optimistic" | "stale" {
if (args.useOptimisticCreateFlowAgent) return "optimistic";
if (args.hasAgent) return "authoritative";
if (args.shouldUseOptimisticStream) return "optimistic";
if (args.hasOptimisticCreateContinuity) return "optimistic";
return "stale";
}
function resolveCatchingUpUi(args: {
shouldUseOptimisticStream: boolean;
hasOptimisticCreateContinuity: boolean;
hasHydratedHistoryBefore: boolean;
hadInitialSyncFailure: boolean;
}): "overlay" | "silent" {
if (args.shouldUseOptimisticStream) return "silent";
if (args.hasOptimisticCreateContinuity) return "silent";
if (args.hasHydratedHistoryBefore) return "silent";
if (args.hadInitialSyncFailure) return "silent";
return "overlay";
@@ -147,7 +155,7 @@ function resolveAgentScreenSync(args: {
return {
status: "catching_up",
ui: resolveCatchingUpUi({
shouldUseOptimisticStream: input.shouldUseOptimisticStream,
hasOptimisticCreateContinuity: hasOptimisticCreateContinuity(input),
hasHydratedHistoryBefore: input.hasHydratedHistoryBefore,
hadInitialSyncFailure,
}),
@@ -227,7 +235,7 @@ export function deriveAgentScreenViewState({
const source = resolveAgentScreenSource({
useOptimisticCreateFlowAgent,
hasAgent: Boolean(input.agent),
shouldUseOptimisticStream: input.shouldUseOptimisticStream,
hasOptimisticCreateContinuity: hasOptimisticCreateContinuity(input),
});
const sync = resolveAgentScreenSync({
@@ -2,7 +2,9 @@ import { useCallback, useMemo, useReducer } from "react";
import type { ComposerAttachment } from "@/attachments/types";
import { splitComposerAttachmentsForSubmit } from "@/components/composer-attachments";
import { useCreateFlowStore } from "@/stores/create-flow-store";
import { useSessionStore } from "@/stores/session-store";
import {
buildOptimisticUserMessage,
generateMessageId,
type StreamItem,
type UserMessageImageAttachment,
@@ -110,6 +112,9 @@ export function useDraftAgentCreateFlow<TDraftAgent, TCreateResult>({
const updatePendingAgentId = useCreateFlowStore((state) => state.updateAgentId);
const markPendingCreateLifecycle = useCreateFlowStore((state) => state.markLifecycle);
const clearPendingCreateAttempt = useCreateFlowStore((state) => state.clear);
const appendOptimisticUserMessageToAgentStream = useSessionStore(
(state) => state.appendOptimisticUserMessageToAgentStream,
);
const formErrorMessage = machine.tag === "draft" ? machine.errorMessage : "";
const isSubmitting = machine.tag === "creating";
@@ -128,19 +133,13 @@ export function useDraftAgentCreateFlow<TDraftAgent, TCreateResult>({
}
return [
{
kind: "user_message",
buildOptimisticUserMessage({
id: machine.attempt.clientMessageId,
text: machine.attempt.text,
timestamp: machine.attempt.timestamp,
optimistic: true,
...(machine.attempt.images && machine.attempt.images.length > 0
? { images: machine.attempt.images }
: {}),
...(machine.attempt.attachments && machine.attempt.attachments.length > 0
? { attachments: machine.attempt.attachments }
: {}),
},
images: machine.attempt.images,
attachments: machine.attempt.attachments,
}),
];
}, [machine]);
@@ -228,6 +227,19 @@ export function useDraftAgentCreateFlow<TDraftAgent, TCreateResult>({
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 });
@@ -241,6 +253,7 @@ export function useDraftAgentCreateFlow<TDraftAgent, TCreateResult>({
}
},
[
appendOptimisticUserMessageToAgentStream,
clearPendingCreateAttempt,
createRequest,
draftId,