style(toolbar): refactor UI and karton instantiation

Author: glenntws

agent/client/src/Agent.ts

@@ -1,3 +1,8 @@
 import { Agent } from './Agent.js';
 
 export { Agent };
+export type {
+  AgentCallbacks,
+  AgentProcedures,
+} from './types/agent-callbacks.js';
+export type { AgentEvent, AgentEventCallback } from './utils/event-utils.js';

agent/client/src/index.ts

@@ -0,0 +1,43 @@
+import type { KartonContract, ChatMessage } from '@stagewise/karton-contract';
+
+type KartonState = KartonContract['state'];
+type StateRecipe = (draft: KartonState) => void;
+
+export type AgentStateGetter = () => KartonState;
+export type AgentStateSetter = (recipe: StateRecipe) => KartonState;
+
+export interface AgentCallbacks {
+  getState: AgentStateGetter;
+  setState: AgentStateSetter;
+}
+
+export interface AgentProcedures {
+  undoToolCallsUntilUserMessage: (
+    userMessageId: string,
+    chatId: string,
+  ) => Promise<void>;
+  undoToolCallsUntilLatestUserMessage: (
+    chatId: string,
+  ) => Promise<ChatMessage | null>;
+  retrySendingUserMessage: () => Promise<void>;
+  refreshSubscription: () => Promise<void>;
+  abortAgentCall: () => Promise<void>;
+  approveToolCall: (
+    toolCallId: string,
+    callingClientId: string,
+  ) => Promise<void>;
+  rejectToolCall: (
+    toolCallId: string,
+    callingClientId: string,
+  ) => Promise<void>;
+  createChat: () => Promise<string>;
+  switchChat: (chatId: string, callingClientId: string) => Promise<void>;
+  deleteChat: (chatId: string, callingClientId: string) => Promise<void>;
+  sendUserMessage: (
+    message: ChatMessage,
+    callingClientId: string,
+  ) => Promise<void>;
+  assistantMadeCodeChangesUntilLatestUserMessage: (
+    chatId: string,
+  ) => Promise<boolean>;
+}

agent/client/src/types/agent-callbacks.ts

@@ -1,29 +1,27 @@
-import type { KartonContract, ChatMessage } from '@stagewise/karton-contract';
-import type { KartonServer } from '@stagewise/karton/server';
+import type { ChatMessage } from '@stagewise/karton-contract';
 import type { ToolCallProcessingResult } from './tool-call-utils.js';
 import type { InferUIMessageChunk, ToolUIPart } from 'ai';
+import type { AgentCallbacks } from '../types/agent-callbacks.js';
 
 /**
  * Checks if a message with the given ID exists in the active chat
- * @param karton - The Karton server instance containing chat state
+ * @param callbacks - The agent callbacks for state access
  * @param messageId - The unique identifier of the message to check
  * @returns True if the message exists in the active chat, false otherwise
  */
-function messageExists(
-  karton: KartonServer<KartonContract>,
-  messageId: string,
-): boolean {
-  return karton.state.chats[karton.state.activeChatId!]!.messages.some(
-    (m) => m.id === messageId,
+function messageExists(callbacks: AgentCallbacks, messageId: string): boolean {
+  const state = callbacks.getState();
+  return state.chats[state.activeChatId!]!.messages.some(
+    (m: ChatMessage) => m.id === messageId,
   );
 }
 
 /**
  * Creates a new chat with a timestamped title and sets it as the active chat
- * @param karton - The Karton server instance to modify
+ * @param callbacks - The agent callbacks for state modification
  * @returns The unique ID of the newly created chat
  */
-export function createAndActivateNewChat(karton: KartonServer<KartonContract>) {
+export function createAndActivateNewChat(callbacks: AgentCallbacks) {
   const chatId = crypto.randomUUID();
   const title = `New Chat - ${new Date().toLocaleString('en-US', {
     month: 'short',
@@ -32,7 +30,7 @@ export function createAndActivateNewChat(karton: KartonServer<KartonContract>) {
     minute: '2-digit',
     hour12: true,
   })}`;
-  karton.setState((draft) => {
+  callbacks.setState((draft) => {
     draft.chats[chatId] = {
       title,
       createdAt: new Date(),
@@ -46,21 +44,22 @@ export function createAndActivateNewChat(karton: KartonServer<KartonContract>) {
 /**
  * Appends text content to a message, creating the message if it doesn't exist
  * or creating/appending to a text part at the specified index
- * @param karton - The Karton server instance to modify
+ * @param callbacks - The agent callbacks for state modification
  * @param messageId - The unique identifier of the message to append to
  * @param delta - The text content to append
  * @param partIndex - The index of the message part to append to
  */
 export function appendTextDeltaToMessage(
-  karton: KartonServer<KartonContract>,
+  callbacks: AgentCallbacks,
   messageId: string,
   delta: string,
   partIndex: number,
 ) {
   // If the message doesn't exist, create it
-  if (!messageExists(karton, messageId)) {
-    karton.setState((draft) => {
-      draft.chats[karton.state.activeChatId!]!.messages.push({
+  if (!messageExists(callbacks, messageId)) {
+    callbacks.setState((draft) => {
+      const state = callbacks.getState();
+      draft.chats[state.activeChatId!]!.messages.push({
         role: 'assistant',
         id: messageId,
         parts: [{ type: 'text', text: delta }],
@@ -71,9 +70,10 @@ export function appendTextDeltaToMessage(
     });
   } else {
     // If the message exists, create a text part or append to the existing one
-    karton.setState((draft) => {
-      const message = draft.chats[karton.state.activeChatId!]!.messages.find(
-        (m) => m.id === messageId,
+    callbacks.setState((draft) => {
+      const state = callbacks.getState();
+      const message = draft.chats[state.activeChatId!]!.messages.find(
+        (m: ChatMessage) => m.id === messageId,
       )!;
 
       // Create a new part if it's a new one
@@ -96,28 +96,29 @@ export function appendTextDeltaToMessage(
 /**
  * Appends tool input information to a message, creating the message if it doesn't exist
  * or updating the tool part at the specified index
- * @param karton - The Karton server instance to modify
+ * @param callbacks - The agent callbacks for state modification
  * @param messageId - The unique identifier of the message to append to
  * @param chunk - The tool input chunk containing tool call details
  * @param partIndex - The index of the message part to update
  */
 export function appendToolInputToMessage(
-  karton: KartonServer<KartonContract>,
+  callbacks: AgentCallbacks,
   messageId: string,
   chunk: Extract<
     InferUIMessageChunk<ChatMessage>,
     { type: 'tool-input-available' }
   >,
   partIndex: number,
 ) {
-  karton.setState((draft) => {
-    const message = draft.chats[karton.state.activeChatId!]!.messages.find(
-      (m) => m.id === messageId,
+  callbacks.setState((draft) => {
+    const state = callbacks.getState();
+    const message = draft.chats[state.activeChatId!]!.messages.find(
+      (m: ChatMessage) => m.id === messageId,
     );
 
     if (!message) {
       // If the message doesn't exist, create it
-      draft.chats[karton.state.activeChatId!]!.messages.push({
+      draft.chats[state.activeChatId!]!.messages.push({
         role: 'assistant',
         id: messageId,
         parts: [
@@ -159,23 +160,24 @@ export function appendToolInputToMessage(
 /**
  * Attaches tool execution results to the corresponding tool parts in a message
  * Updates the tool part state to reflect success or error outcomes
- * @param karton - The Karton server instance to modify
+ * @param callbacks - The agent callbacks for state modification
  * @param toolResults - Array of tool execution results to attach
  * @param messageId - The unique identifier of the message containing the tool parts
  */
 export function attachToolOutputToMessage(
-  karton: KartonServer<KartonContract>,
+  callbacks: AgentCallbacks,
   toolResults: ToolCallProcessingResult[],
   messageId: string,
 ) {
-  karton.setState((draft) => {
-    const message = draft.chats[karton.state.activeChatId!]!.messages.find(
-      (m) => m.id === messageId,
+  callbacks.setState((draft) => {
+    const state = callbacks.getState();
+    const message = draft.chats[state.activeChatId!]!.messages.find(
+      (m: ChatMessage) => m.id === messageId,
     );
     if (!message) return;
     for (const result of toolResults) {
       const part = message.parts.find(
-        (p) => 'toolCallId' in p && p.toolCallId === result.toolCallId,
+        (p: any) => 'toolCallId' in p && p.toolCallId === result.toolCallId,
       );
       if (!part) continue;
       if (part.type !== 'dynamic-tool' && !part.type.startsWith('tool-'))
@@ -196,14 +198,16 @@ export function attachToolOutputToMessage(
 
 /**
  * Finds tool calls in the last assistant message that don't have corresponding results
+ * @param callbacks - The agent callbacks for state access
  * @param chatId - The chat ID to check
  * @returns Array of pending tool call IDs and their names
  */
 export function findPendingToolCalls(
-  karton: KartonServer<KartonContract>,
+  callbacks: AgentCallbacks,
   chatId: string,
 ): Array<{ toolCallId: string }> {
-  const chat = karton.state.chats[chatId];
+  const state = callbacks.getState();
+  const chat = state.chats[chatId];
   if (!chat) return [];
 
   const messages = chat.messages;

agent/client/src/utils/karton-helpers.ts

@@ -33,6 +33,7 @@
   },
   "devDependencies": {
     "@stagewise/karton": "workspace:*",
+    "@stagewise/karton-contract": "workspace:*",
     "@inquirer/prompts": "^7.0.0",
     "@stagewise-plugins/angular": "workspace:*",
     "@stagewise-plugins/react": "workspace:*",

apps/cli/package.json

@@ -1,11 +1,17 @@
 import { printInfoMessages } from '@/utils/print-info-messages.js';
 import { log } from '../utils/logger.js';
 import configResolver from '@/config/index.js';
-import { Agent } from '@stagewise/agent-client';
+import { Agent, type AgentCallbacks } from '@stagewise/agent-client';
 import { ClientRuntimeNode } from '@stagewise/agent-runtime-node';
 import { analyticsEvents } from '@/utils/telemetry.js';
+import {
+  createKartonServer,
+  type KartonServer,
+} from '@stagewise/karton/server';
+import type { KartonContract } from '@stagewise/karton-contract';
 
 let agentInstance: Agent | null = null;
+let kartonServer: KartonServer<KartonContract> | null = null;
 
 /**
  * Loads and initializes the agent server
@@ -31,7 +37,23 @@ export async function loadAndInitializeAgent(
       workingDirectory: config.dir,
     });
 
-    // Create agent instance
+    // Create a placeholder for the karton server
+    let tempKartonServer: KartonServer<KartonContract> | null = null;
+
+    // Create callbacks that will use the karton server (will be set later)
+    const callbacks: AgentCallbacks = {
+      getState: () => {
+        if (!tempKartonServer) throw new Error('Karton server not initialized');
+        return tempKartonServer.state;
+      },
+      setState: (recipe) => {
+        if (!tempKartonServer) throw new Error('Karton server not initialized');
+        // @ts-ignore we'll fix this whole temp instantiation shit later
+        return tempKartonServer.setState(recipe);
+      },
+    };
+
+    // Create agent instance with callbacks
     agentInstance = Agent.getInstance({
       clientRuntime,
       accessToken,
@@ -55,17 +77,36 @@ export async function loadAndInitializeAgent(
             break;
         }
       },
+      callbacks,
     });
 
-    // Initialize agent with Express integration
-    // This will automatically set up the Karton endpoint
-    const agentServer = await agentInstance.initialize();
+    // Now create the karton server with agent procedures
+    kartonServer = await createKartonServer<KartonContract>({
+      procedures: agentInstance.getAgentProcedures() as any,
+      initialState: {
+        workspaceInfo: {
+          path: clientRuntime.fileSystem.getCurrentWorkingDirectory(),
+          devAppPort: 0,
+          loadedPlugins: [],
+        },
+        activeChatId: null,
+        chats: {},
+        isWorking: false,
+        toolCallApprovalRequests: [],
+        subscription: undefined,
+      },
+    });
+
+    // Set the karton server reference in callbacks
+    tempKartonServer = kartonServer;
+
+    // Initialize the agent
+    await agentInstance.initialize();
 
-    // Return the WebSocket server instance if available
-    // The agent SDK may not return the WebSocket server in current versions
+    // Return the WebSocket server instance from karton server
     return {
       success: true,
-      wss: agentServer.wss,
+      wss: kartonServer.wss,
     };
   } catch (error) {
     log.error(
@@ -82,15 +123,31 @@ export function shutdownAgent(): void {
   if (agentInstance?.shutdown) {
     try {
       agentInstance.shutdown();
-      log.debug('Agent server shut down successfully');
+      log.debug('Agent shut down successfully');
     } catch (error) {
       log.error(
         `Error shutting down agent: ${error instanceof Error ? error.message : 'Unknown error'}`,
       );
     }
   }
-  // Clear the instance reference
+
+  if (kartonServer) {
+    try {
+      // Close WebSocket server
+      if (kartonServer.wss) {
+        kartonServer.wss.close();
+      }
+      log.debug('Karton server shut down successfully');
+    } catch (error) {
+      log.error(
+        `Error shutting down karton server: ${error instanceof Error ? error.message : 'Unknown error'}`,
+      );
+    }
+  }
+
+  // Clear the instance references
   agentInstance = null;
+  kartonServer = null;
 }
 
 export function getAgentInstance(): any {

apps/cli/src/server/agent-loader.ts

@@ -43,6 +43,7 @@ export type AgentError = {
 
 type AppState = {
   activeChatId: ChatId | null;
+  workspacePath: string | null;
   chats: Record<ChatId, Chat>;
   toolCallApprovalRequests: string[];
   isWorking: boolean;

packages/karton-contract/src/index.ts

@@ -13,10 +13,12 @@ export const buttonVariants = cva(
         ghost: 'bg-transparent font-medium text-foreground hover:bg-zinc-500/5',
       },
       size: {
+        xs: 'h-6 rounded-xl px-2 py-1 text-xs',
         sm: 'h-8 rounded-xl px-2 py-1 text-sm',
         md: 'h-10 rounded-xl px-4 py-2 text-sm',
         lg: 'h-12 rounded-xl px-6 py-3 text-base',
         xl: 'h-14 rounded-xl px-8 py-4 text-lg',
+        'icon-xs': 'size-6 rounded-full',
         'icon-sm': 'size-8 rounded-full',
         'icon-md': 'size-10 rounded-full',
       },

packages/stage-ui/src/components/button.tsx

@@ -67,7 +67,7 @@ export function MenuItem({ children, className, ...props }: MenuItemProps) {
     <MenuBase.Item
       {...props}
       className={cn(
-        'flex w-full min-w-24 cursor-default flex-row items-center justify-start gap-2 rounded-md py-1.5 pr-6 pl-2 text-sm transition-all duration-150 ease-out hover:bg-black/5 hover:pl-2.25 dark:hover:bg-white/5',
+        'flex w-full min-w-24 cursor-default flex-row items-center justify-start gap-2 rounded-md py-1.5 pr-6 pl-2 text-sm transition-all duration-150 ease-out hover:bg-black/5 hover:pr-5.75 hover:pl-2.25 dark:hover:bg-white/5',
         className,
       )}
     >
@@ -96,7 +96,7 @@ export function MenuSubmenuTrigger({
     <MenuBase.SubmenuTrigger
       {...props}
       className={cn(
-        'group flex w-full min-w-24 cursor-default flex-row items-center justify-start gap-2 rounded-md px-2 py-1.5 text-sm transition-all duration-150 ease-out hover:bg-black/5 hover:pl-2.25 dark:hover:bg-white/5',
+        'group flex w-full min-w-24 cursor-default flex-row items-center justify-start gap-2 rounded-md px-2 py-1.5 text-sm transition-all duration-150 ease-out hover:bg-black/5 hover:pr-1.75 hover:pl-2.25 dark:hover:bg-white/5',
         className,
       )}
     >