ai-design-system 0.1.91 → 0.1.93

package/components/blocks/AIConversation/AIConversation.tsx

@@ -56,6 +56,7 @@ interface AIMessage {
   isLoading?: boolean;
   blocks?: AIMessageBlock[];
   checkpointId?: string;
+  restoreLabel?: string;
   attachments?: UserMessageAttachment[];
 }
 
@@ -98,6 +99,17 @@ export interface AIConversationProps
    * Callback fired when user restores a conversation checkpoint
    */
   onRestoreCheckpoint?: (messageId: string, checkpointId: string) => void
+  /**
+   * Save the workspace as it stands after an orchestrator reply.
+   *
+   * Paired with `onRestoreCheckpoint`: undo sits on what the user asked for,
+   * save sits on what the agent produced. Omit it and no reply offers to save.
+   */
+  onSaveWorkspace?: () => void
+  /** True when the current version is already kept — every Save reads as saved. */
+  isCurrentVersionSaved?: boolean
+  /** A save is in flight — the control says so and refuses a second one. */
+  isSavingWorkspace?: boolean
 }
 
 /**
@@ -113,6 +125,9 @@ export const AIConversation = React.memo<AIConversationProps>(
     onToolAction,
     renderSystemMessage,
     onRestoreCheckpoint,
+    onSaveWorkspace,
+    isCurrentVersionSaved,
+    isSavingWorkspace,
     ...conversationProps
   }) => {
     const isEmpty = React.useMemo(
@@ -242,6 +257,7 @@ export const AIConversation = React.memo<AIConversationProps>(
               subAgents: msg.subAgents ? [...msg.subAgents] : undefined,
               blocks: msg.blocks ? [...msg.blocks] : undefined,
               checkpointId: msg.checkpointId,
+              restoreLabel: msg.restoreLabel,
             };
             result.push(currentGroup);
           }
@@ -257,6 +273,27 @@ export const AIConversation = React.memo<AIConversationProps>(
       return result;
     }, [messages]);
 
+  /**
+   * The last orchestrator reply, and the ONLY one that may offer to save.
+   *
+   * Saving anchors to the most recent request, so it captures the workspace as
+   * it is NOW — not as it was when some earlier reply was written. Offering the
+   * control on every reply implies each one saves its own moment, which is
+   * false and unrecoverable once acted on: the user would believe they had kept
+   * a point they never kept.
+   *
+   * `isCurrentVersionSaved` is likewise one fact about the workspace, not a
+   * property of each message. Rendered on every reply it read as "all of these
+   * are saved"; it belongs on the single reply the save would come from.
+   */
+  const lastReplyIndex = React.useMemo(() => {
+    for (let i = groupedMessages.length - 1; i >= 0; i--) {
+      const m = groupedMessages[i]
+      if (m.role === "orchestrator" || m.role === "assistant" || (!m.role && m.type === "ai")) return i
+    }
+    return -1
+  }, [groupedMessages])
+
     const renderedMessages = React.useMemo(
       () =>
         groupedMessages.map((message, index) => {
@@ -291,6 +328,7 @@ export const AIConversation = React.memo<AIConversationProps>(
                   avatarName: message.avatarName,
                   attachments: message.attachments,
                   checkpointId: message.checkpointId,
+                  restoreLabel: message.restoreLabel,
                 }}
                 showAvatar={showAvatars}
                 onRestore={onRestoreCheckpoint}
@@ -335,6 +373,16 @@ export const AIConversation = React.memo<AIConversationProps>(
                     isLoading: message.isLoading && !hasReasoning,
                   }}
                   showAvatar={showAvatars && hasDisplayContent}
+                  // Only on a settled reply with something to read: a save
+                  // offered beside a streaming or empty message would anchor a
+                  // snapshot to a turn that has not produced anything yet.
+                  onSave={
+                    index === lastReplyIndex && hasDisplayContent && !isStreaming
+                      ? onSaveWorkspace
+                      : undefined
+                  }
+                  isSaved={index === lastReplyIndex && isCurrentVersionSaved}
+                  isSaving={index === lastReplyIndex && isSavingWorkspace}
                 >
                   {/* Render reasoning-section for hidden tool results and subagents */}
                   {hasReasoning && (
@@ -425,7 +473,7 @@ export const AIConversation = React.memo<AIConversationProps>(
             </React.Fragment>
           )
         }),
-      [groupedMessages, showAvatars, onToolAction, onRestoreCheckpoint]
+      [groupedMessages, showAvatars, onToolAction, onRestoreCheckpoint, onSaveWorkspace, isCurrentVersionSaved, isSavingWorkspace, lastReplyIndex]
     )
 
     return (

package/components/blocks/WorkflowCanvas/WorkflowCanvas.tsx

@@ -6,8 +6,6 @@ import {
   type Connection,
   type Node,
   useReactFlow,
-  useStoreApi,
-  useUpdateNodeInternals,
 } from "@xyflow/react";
 import { useCallback, useEffect } from "react";
 import { Canvas } from "@/components/ai-elements/canvas";
@@ -24,22 +22,6 @@ import { AgentAnnotation } from "@/components/composites/AgentAnnotation";
 import type { WorkflowCanvasProps, WorkflowEdge } from "@/lib/workflow/interfaces";
 import "@xyflow/react/dist/style.css";
 
-// Bounded retry for node measurement — see the effect in WorkflowCanvasInner.
-//
-// ~30s, which sounds long and is not. The loop stops the instant every node
-// reports handleBounds, so in the ordinary case it ends after one pass; the cap
-// only bounds the pathological case where a node never appears at all.
-//
-// It was 3s, and that was too short for the case this exists to fix. The grid's
-// device nodes each carry an iframe, and on a cold sandbox they are not in the
-// DOM until Metro has served a bundle. `useUpdateNodeInternals` silently drops
-// every id it cannot resolve, so a window that closes before the iframes mount
-// spends all its attempts on nothing and gives up — leaving exactly the blank
-// canvas of edgeless nodes it was added to prevent. Verified live: all five
-// nodes resolvable, all five still `handleBounds: NULL`.
-const MEASURE_RETRY_MS = 250;
-const MEASURE_MAX_ATTEMPTS = 120;
-
 const edgeTypes = {
   straight: Edge.Strict,
   animated: Edge.Animated,
@@ -76,67 +58,6 @@ function WorkflowCanvasInner({
   className,
 }: WorkflowCanvasProps) {
   const { fitView } = useReactFlow();
-  const updateNodeInternals = useUpdateNodeInternals();
-  const store = useStoreApi();
-
-  /**
-   * Force xyflow to measure the nodes it is showing.
-   *
-   * xyflow will not draw an edge until BOTH endpoints have `handleBounds` —
-   * the measured position of each handle within its node — and those come from
-   * one place only: the per-node ResizeObserver. Publishing `initialWidth` /
-   * `initialHeight` (which the preview nodes do) satisfies the separate
-   * `nodeHasDimensions` check that controls VISIBILITY, and does nothing for
-   * `handleBounds`. The two are easy to conflate and fail very differently.
-   *
-   * When that observer does not deliver — it fires as an unsynchronized
-   * callback (see the fitView effect below, which works around the same race),
-   * and reports nothing for a node that is re-created but never changes size —
-   * the result is not a missing node. It is a canvas of perfectly normal nodes
-   * with EVERY edge silently absent: `edgeLookup` holds them all, each with
-   * valid endpoints and handle ids, and not one is rendered. Observed live with
-   * all five nodes at `measured: {}` and `handleBounds: null` while their DOM
-   * elements had real size, and xyflow logged no warning.
-   *
-   * Keyed on the node ids rather than the `nodes` array, which is a new
-   * reference on nearly every render.
-   *
-   * RETRIED, not fired once, because the request is silently dropped when it
-   * is early. `useUpdateNodeInternals` resolves each id through
-   * `domNode.querySelector('.react-flow__node[data-id=...]')` and keeps only
-   * the ones it finds — a node that has not mounted yet is skipped with no
-   * error and no retry, and since the id list has not changed, the effect
-   * never runs again. One shot at 60ms worked for the single-device view and
-   * missed the grid every time, where five nodes carrying iframes mount later.
-   * So it repeats until the store actually reports `handleBounds`, and stops
-   * the moment it does.
-   */
-  const nodeIdSignature = nodes.map((n) => n.id).join("|");
-  useEffect(() => {
-    if (!nodeIdSignature) return;
-    const ids = nodeIdSignature.split("|");
-    let attempts = 0;
-    let timer: ReturnType<typeof setTimeout>;
-
-    const measured = () => {
-      const lookup = store.getState().nodeLookup;
-      return ids.every((id) => !!lookup.get(id)?.internals?.handleBounds);
-    };
-
-    const pass = () => {
-      updateNodeInternals(ids);
-      attempts += 1;
-      // `updateNodeInternals` defers its own work to a rAF, so the result is
-      // only readable on a later tick — hence check on the next timer rather
-      // than immediately after the call.
-      timer = setTimeout(() => {
-        if (!measured() && attempts < MEASURE_MAX_ATTEMPTS) pass();
-      }, MEASURE_RETRY_MS);
-    };
-
-    timer = setTimeout(pass, 60);
-    return () => clearTimeout(timer);
-  }, [nodeIdSignature, updateNodeInternals, store]);
 
   // Opt-in re-fit, e.g. after the device-preview toolbar's Route select
   // switches to "All" and the node count/extent changes underneath the

package/components/composites/FilePreviewDialog/FilePreviewDialog.stories.tsx

@@ -1,4 +1,5 @@
 import type { Meta, StoryObj } from '@storybook/react'
+import { fn } from '@storybook/test'
 import { FilePreviewDialog } from './FilePreviewDialog'
 
 const meta = {
@@ -14,5 +15,26 @@ export default meta
 type Story = StoryObj<typeof meta>
 
 export const Default: Story = {
-  args: {} as any,
+  args: {
+    open: true,
+    onOpenChange: fn(),
+    title: 'Sprint checklist',
+    description: 'specs/sprint-plan.md',
+    content: [
+      '## Tasks',
+      '- [x] Draft API contract',
+      '- [ ] Land UI stubs',
+      '- [ ] Wire preview',
+    ].join('\n'),
+  },
+}
+
+export const Empty: Story = {
+  args: {
+    open: true,
+    onOpenChange: fn(),
+    title: 'Untitled spec',
+    description: 'specs/untitled.md',
+    content: undefined,
+  },
 }

package/components/composites/OrchestratorMessage/OrchestratorMessage.tsx

@@ -11,6 +11,7 @@ import {
   MessageTypingIndicator,
 } from "@/components/ai-elements/message"
 import { Response } from "@/components/composites/response"
+import { cn } from "@/lib/utils"
 
 /**
  * OrchestratorMessage Block
@@ -36,6 +37,23 @@ export interface OrchestratorMessageProps {
    * Whether to show avatar
    */
   showAvatar?: boolean
+  /**
+   * Save the workspace as it stands after this reply.
+   *
+   * Sits here and not only in the toolbar because a snapshot is a point in the
+   * CONVERSATION, not a global act: the reply is what tells the user something
+   * is worth keeping, so the control belongs where they read that. It pairs
+   * with the undo control on user messages — save beside what the agent
+   * produced, undo beside what you asked for.
+   *
+   * Absent when there is nothing to save to, so the control simply does not
+   * appear rather than offering an action that cannot work.
+   */
+  onSave?: () => void
+  /** True once this version is kept — the control says so instead of repeating. */
+  isSaved?: boolean
+  /** A save is in flight. Refuses a second press: two presses is two snapshots. */
+  isSaving?: boolean
   /**
    * Child components (SpecialistMessage, ToolCallDisplay)
    */
@@ -46,14 +64,14 @@ export interface OrchestratorMessageProps {
  * OrchestratorMessage component - displays coordinator messages with nested specialists and tools
  */
 export const OrchestratorMessage = React.memo<OrchestratorMessageProps>(
-  ({ message, showAvatar = true, children }) => {
+  ({ message, showAvatar = true, children, onSave, isSaved, isSaving }) => {
     const hasContent = React.useMemo(
       () => message.content && message.content.trim() !== "",
       [message.content]
     )
 
     return (
-      <Message from="assistant">
+      <Message from="assistant" className="group">
         {showAvatar && (message.avatarSrc
           ? <MessageAvatar
               src={message.avatarSrc}
@@ -80,6 +98,37 @@ export const OrchestratorMessage = React.memo<OrchestratorMessageProps>(
 
           {/* Child components (specialists and tool calls) */}
           {children && <div className="mt-4 space-y-4">{children}</div>}
+
+          {/* Revealed on hover or keyboard focus, like the undo control on a
+              user message: keeping a version is deliberate and occasional, and
+              a button standing on every reply reads as something you are
+              expected to press. Hidden entirely while streaming — there is no
+              settled version to keep yet. */}
+          {onSave && !message.isLoading && (
+            <button
+              type="button"
+              aria-label={isSaved ? "This version is already saved" : "Save this version"}
+              title={isSaved ? "This version is already saved" : "Save this version"}
+              disabled={isSaved || isSaving}
+              onClick={onSave}
+              className={cn(
+                "mt-2 inline-flex items-center gap-1.5 rounded-md px-2 py-1 text-xs",
+                "text-muted-foreground transition-opacity",
+                "hover:bg-accent hover:text-foreground",
+                "focus-visible:opacity-100 focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring",
+                isSaved || isSaving ? "opacity-100" : "opacity-0 group-hover:opacity-100",
+                "disabled:pointer-events-none"
+              )}
+            >
+              <Icon
+                name={isSaved ? "check" : isSaving ? "loader-2" : "save"}
+                size="sm"
+                className={isSaving ? "animate-spin" : undefined}
+                aria-hidden
+              />
+              {isSaved ? "Saved" : isSaving ? "Saving…" : "Save"}
+            </button>
+          )}
         </div>
       </Message>
     )

package/components/composites/ProjectSwitcher/ProjectSwitcher.stories.tsx

@@ -1,4 +1,5 @@
 import type { Meta, StoryObj } from '@storybook/react'
+import { fn } from '@storybook/test'
 import { ProjectSwitcher } from './ProjectSwitcher'
 
 const meta = {
@@ -14,5 +15,23 @@ export default meta
 type Story = StoryObj<typeof meta>
 
 export const Default: Story = {
-  args: {} as any,
+  args: {
+    projects: [
+      { id: '1', name: 'Atlas' },
+      { id: '2', name: 'Beacon' },
+      { id: '3', name: 'Cascade' },
+    ],
+    selectedProjectId: '1',
+    onSelectProject: fn(),
+    onCreateProject: fn(),
+  },
+}
+
+export const Empty: Story = {
+  args: {
+    projects: [],
+    selectedProjectId: null,
+    onSelectProject: fn(),
+    onCreateProject: fn(),
+  },
 }

package/components/composites/ProjectSwitcher/ProjectSwitcher.tsx

@@ -31,7 +31,13 @@ export interface ProjectSwitcherProps {
 }
 
 export const ProjectSwitcher = React.memo<ProjectSwitcherProps>(
-  ({ projects, selectedProjectId, onSelectProject, onCreateProject, className }) => {
+  ({
+    projects = [],
+    selectedProjectId,
+    onSelectProject = () => {},
+    onCreateProject = () => {},
+    className,
+  }) => {
     const [open, setOpen] = React.useState(false)
 
     const selectedProject = projects.find((p) => String(p.id) === String(selectedProjectId))

package/components/composites/PromptInput/PromptInput.stories.tsx

@@ -273,3 +273,66 @@ export const ToggleDialog: Story = {
     },
   },
 };
+
+/**
+ * With Context Window
+ * Displays context token consumption, progress ring, breakdown, and estimated cost
+ */
+export const WithContextWindow: Story = {
+  args: {
+    placeholder: "Ask a question about the codebase...",
+    context: {
+      usedTokens: 42500,
+      maxTokens: 128000,
+      modelId: "gpt-4o",
+      usage: {
+        inputTokens: 38000,
+        outputTokens: 3500,
+        reasoningTokens: 1000,
+        cachedInputTokens: 15000,
+      },
+    },
+    onSubmit: (message) => {
+      console.log("Submitted:", message);
+    },
+  },
+  parameters: {
+    docs: {
+      description: {
+        story:
+          "Displays context window usage via the AI Elements Context compound component in the toolbar. Shows token usage percentage, circular progress, input/output/reasoning/cache breakdown, and estimated USD cost on hover.",
+      },
+    },
+  },
+};
+
+/**
+ * With Near-Limit Context Window
+ * Demonstrates high context utilization nearing the maximum context window
+ */
+export const WithNearLimitContext: Story = {
+  args: {
+    placeholder: "Ask a question with high context utilization...",
+    context: {
+      usedTokens: 118500,
+      maxTokens: 128000,
+      modelId: "gpt-4o",
+      usage: {
+        inputTokens: 112000,
+        outputTokens: 6500,
+      },
+    },
+    onSubmit: (message) => {
+      console.log("Submitted:", message);
+    },
+  },
+  parameters: {
+    docs: {
+      description: {
+        story:
+          "Demonstrates context window indicator when the model is near its token capacity (92.6% utilized).",
+      },
+    },
+  },
+};
+

package/components/composites/PromptInput/PromptInput.tsx

@@ -2,6 +2,7 @@
 
 import * as React from "react";
 import type { ReactNode } from "react";
+import type { LanguageModelUsage } from "ai";
 import {
   PromptInput as AIPromptInput,
   PromptInputAttachment,
@@ -18,10 +19,29 @@ import {
   usePromptInputController,
   useOptionalPromptInputController,
 } from "@/components/ai-elements/prompt-input";
+import {
+  Context,
+  ContextTrigger,
+  ContextContent,
+  ContextContentHeader,
+  ContextContentBody,
+  ContextContentFooter,
+  ContextInputUsage,
+  ContextOutputUsage,
+  ContextReasoningUsage,
+  ContextCacheUsage,
+} from "@/components/ai-elements/context";
 import type { FormEvent } from "react";
 import { Button } from "@/components/primitives/Button";
 import { Icon } from "@/components/primitives/Icon";
 
+export interface PromptInputContextProps {
+  usedTokens: number;
+  maxTokens: number;
+  usage?: LanguageModelUsage;
+  modelId?: string;
+}
+
 export interface PromptInputBlockProps
   extends Omit<
     AIPromptInputProps,
@@ -38,6 +58,8 @@ export interface PromptInputBlockProps
   dialog?: ReactNode;
   loading?: boolean;
   onStop?: () => void;
+  context?: PromptInputContextProps | ReactNode;
+  tools?: ReactNode;
 }
 
 export const PromptInput = React.memo<PromptInputBlockProps>(
@@ -50,6 +72,8 @@ export const PromptInput = React.memo<PromptInputBlockProps>(
     dialog,
     loading = false,
     onStop,
+    context,
+    tools,
     accept = "image/*",
     multiple = true,
     maxFiles,
@@ -111,6 +135,13 @@ export const PromptInput = React.memo<PromptInputBlockProps>(
         <PromptInputFooter>
           <PromptInputTools>
             <AttachButton disabled={disabled || loading} />
+            {tools}
+            {context ? (
+              <PromptInputContextIndicator
+                context={context}
+                disabled={disabled || loading}
+              />
+            ) : null}
           </PromptInputTools>
           <PromptInputSubmit
             disabled={disabled || (loading && !onStop)}
@@ -187,4 +218,50 @@ function ExternalValueSync({ value }: { value?: string }) {
   return null;
 }
 
+/**
+ * Renders the context window usage indicator.
+ * Supports either a PromptInputContextProps object (auto-rendered using Context compound elements)
+ * or a custom ReactNode.
+ */
+function PromptInputContextIndicator({
+  context,
+  disabled,
+}: {
+  context: PromptInputContextProps | ReactNode;
+  disabled?: boolean;
+}) {
+  if (React.isValidElement(context)) {
+    return context;
+  }
+
+  const ctx = context as PromptInputContextProps;
+  if (!ctx || typeof ctx.usedTokens !== "number" || typeof ctx.maxTokens !== "number") {
+    return null;
+  }
+
+  return (
+    <Context
+      usedTokens={ctx.usedTokens}
+      maxTokens={ctx.maxTokens}
+      usage={ctx.usage}
+      modelId={ctx.modelId}
+    >
+      <ContextTrigger
+        className="h-8 px-2 text-xs font-normal gap-1.5"
+        disabled={disabled}
+      />
+      <ContextContent>
+        <ContextContentHeader />
+        <ContextContentBody className="space-y-1.5">
+          <ContextInputUsage />
+          <ContextOutputUsage />
+          <ContextReasoningUsage />
+          <ContextCacheUsage />
+        </ContextContentBody>
+        {ctx.modelId && <ContextContentFooter />}
+      </ContextContent>
+    </Context>
+  );
+}
+
 PromptInput.displayName = "PromptInput";

package/components/composites/PromptInput/index.ts

@@ -1,5 +1,5 @@
 export { PromptInput } from "./PromptInput";
-export type { PromptInputBlockProps } from "./PromptInput";
+export type { PromptInputBlockProps, PromptInputContextProps } from "./PromptInput";
 
 export {
   PromptInputProvider,
@@ -15,3 +15,29 @@ export type {
   PromptInputProviderProps,
 } from "@/components/ai-elements/prompt-input";
 
+export {
+  Context,
+  ContextTrigger,
+  ContextContent,
+  ContextContentHeader,
+  ContextContentBody,
+  ContextContentFooter,
+  ContextInputUsage,
+  ContextOutputUsage,
+  ContextReasoningUsage,
+  ContextCacheUsage,
+} from "@/components/ai-elements/context";
+
+export type {
+  ContextProps,
+  ContextTriggerProps,
+  ContextContentProps,
+  ContextContentHeaderProps,
+  ContextContentBodyProps,
+  ContextContentFooterProps,
+  ContextInputUsageProps,
+  ContextOutputUsageProps,
+  ContextReasoningUsageProps,
+  ContextCacheUsageProps,
+} from "@/components/ai-elements/context";
+