refactor: improve tool part schema type safety and remove dead code

1. Refactor MuxToolPartSchema using extend pattern for DRY code:
   - Base schema shares common fields (type, toolCallId, toolName, input, timestamp)
   - Pending variant: state="input-available", no output field
   - Available variant: state="output-available", output required

2. Remove dead ReasoningStartEventSchema:
   - Schema was defined but never emitted by backend
   - Not in WorkspaceChatMessageSchema union
   - No type guard existed for it

3. Update MuxToolPart type to use Zod inference:
   - Type now properly discriminates based on state
   - Accessing output requires narrowing to output-available state

Author: ThomasK33

src/browser/utils/messages/StreamingMessageAggregator.ts

@@ -419,12 +419,11 @@ export class StreamingMessageAggregator {
         if (data.parts) {
           // Sync up the tool results from the backend's parts array
           for (const backendPart of data.parts) {
-            if (backendPart.type === "dynamic-tool") {
+            if (backendPart.type === "dynamic-tool" && backendPart.state === "output-available") {
               // Find and update existing tool part
               const toolPart = message.parts.find(
                 (part): part is DynamicToolPart =>
-                  part.type === "dynamic-tool" &&
-                  (part as DynamicToolPart).toolCallId === backendPart.toolCallId
+                  part.type === "dynamic-tool" && part.toolCallId === backendPart.toolCallId
               );
               if (toolPart) {
                 // Update with result from backend
@@ -510,7 +509,7 @@ export class StreamingMessageAggregator {
     // Check if this tool call already exists to prevent duplicates
     const existingToolPart = message.parts.find(
       (part): part is DynamicToolPart =>
-        part.type === "dynamic-tool" && (part as DynamicToolPart).toolCallId === data.toolCallId
+        part.type === "dynamic-tool" && part.toolCallId === data.toolCallId
     );
 
     if (existingToolPart) {

src/common/orpc/schemas.ts

@@ -71,7 +71,6 @@ export {
   QueuedMessageChangedEventSchema,
   ReasoningDeltaEventSchema,
   ReasoningEndEventSchema,
-  ReasoningStartEventSchema,
   RestoreToInputEventSchema,
   SendMessageOptionsSchema,
   StreamAbortEventSchema,

src/common/orpc/schemas/message.ts

@@ -19,41 +19,32 @@ export const MuxReasoningPartSchema = z.object({
   timestamp: z.number().optional(),
 });
 
-// Discriminated tool part schemas for proper type inference
-export const DynamicToolPartAvailableSchema = z.object({
+// Base schema for tool parts - shared fields
+const MuxToolPartBase = z.object({
   type: z.literal("dynamic-tool"),
   toolCallId: z.string(),
   toolName: z.string(),
-  state: z.literal("output-available"),
   input: z.unknown(),
-  output: z.unknown(),
   timestamp: z.number().optional(),
 });
 
-export const DynamicToolPartPendingSchema = z.object({
-  type: z.literal("dynamic-tool"),
-  toolCallId: z.string(),
-  toolName: z.string(),
+// Discriminated tool part schemas - output required only when state is "output-available"
+export const DynamicToolPartPendingSchema = MuxToolPartBase.extend({
   state: z.literal("input-available"),
-  input: z.unknown(),
-  timestamp: z.number().optional(),
+});
+
+export const DynamicToolPartAvailableSchema = MuxToolPartBase.extend({
+  state: z.literal("output-available"),
+  output: z.unknown(),
 });
 
 export const DynamicToolPartSchema = z.discriminatedUnion("state", [
   DynamicToolPartAvailableSchema,
   DynamicToolPartPendingSchema,
 ]);
 
-// Alias for backward compatibility - used in message schemas
-export const MuxToolPartSchema = z.object({
-  type: z.literal("dynamic-tool"),
-  toolCallId: z.string(),
-  toolName: z.string(),
-  state: z.enum(["input-available", "output-available"]),
-  input: z.unknown(),
-  output: z.unknown().optional(),
-  timestamp: z.number().optional(),
-});
+// Alias for message schemas
+export const MuxToolPartSchema = DynamicToolPartSchema;
 
 export const MuxImagePartSchema = z.object({
   type: z.literal("file"),

src/common/orpc/schemas/stream.ts

@@ -154,12 +154,6 @@ export const ToolCallEndEventSchema = z.object({
   timestamp: z.number().meta({ description: "When tool call completed (Date.now())" }),
 });
 
-export const ReasoningStartEventSchema = z.object({
-  type: z.literal("reasoning-start"),
-  workspaceId: z.string(),
-  messageId: z.string(),
-});
-
 export const ReasoningDeltaEventSchema = z.object({
   type: z.literal("reasoning-delta"),
   workspaceId: z.string(),

src/common/types/message.ts

@@ -3,7 +3,8 @@ import type { LanguageModelV2Usage } from "@ai-sdk/provider";
 import type { StreamErrorType } from "./errors";
 import type { ToolPolicy } from "@/common/utils/tools/toolPolicy";
 import type { ChatUsageDisplay } from "@/common/utils/tokens/usageAggregator";
-import type { ImagePart } from "@/common/orpc/schemas";
+import type { ImagePart, MuxToolPartSchema } from "@/common/orpc/schemas";
+import type { z } from "zod";
 
 // Message to continue with after compaction
 export interface ContinueMessage {
@@ -54,15 +55,8 @@ export interface MuxMetadata {
 
 // Extended tool part type that supports interrupted tool calls (input-available state)
 // Standard AI SDK ToolUIPart only supports output-available (completed tools)
-export interface MuxToolPart {
-  type: "dynamic-tool";
-  toolCallId: string;
-  toolName: string;
-  state: "input-available" | "output-available";
-  input: unknown;
-  output?: unknown;
-  timestamp?: number; // When the tool call was emitted
-}
+// Uses discriminated union: output is required when state is "output-available", absent when "input-available"
+export type MuxToolPart = z.infer<typeof MuxToolPartSchema>;
 
 // Text part type
 export interface MuxTextPart {

src/common/types/stream.ts

@@ -8,7 +8,6 @@ import type {
   ErrorEventSchema,
   ReasoningDeltaEventSchema,
   ReasoningEndEventSchema,
-  ReasoningStartEventSchema,
   StreamAbortEventSchema,
   StreamDeltaEventSchema,
   StreamEndEventSchema,
@@ -36,7 +35,6 @@ export type ToolCallStartEvent = z.infer<typeof ToolCallStartEventSchema>;
 export type ToolCallDeltaEvent = z.infer<typeof ToolCallDeltaEventSchema>;
 export type ToolCallEndEvent = z.infer<typeof ToolCallEndEventSchema>;
 
-export type ReasoningStartEvent = z.infer<typeof ReasoningStartEventSchema>;
 export type ReasoningDeltaEvent = z.infer<typeof ReasoningDeltaEventSchema>;
 export type ReasoningEndEvent = z.infer<typeof ReasoningEndEventSchema>;
 
@@ -55,7 +53,6 @@ export type AIServiceEvent =
   | ToolCallStartEvent
   | ToolCallDeltaEvent
   | ToolCallEndEvent
-  | ReasoningStartEvent
   | ReasoningDeltaEvent
   | ReasoningEndEvent
   | UsageDeltaEvent;