🤖 feat: well-typed extension payload with discriminated union

- Created ToolUsePayload discriminated union by toolName
- Each tool has specific args and result types (BashToolArgs, BashToolResult, etc.)
- PostToolUseHookPayload = ToolUsePayload & { runtime: Runtime }
- TypeScript narrows args/result types based on toolName check
- Updated docs with type-safe example showing discrimination
- Removed unknown types in favor of specific tool types
- Tests pass, types are fully type-safe

Author: ammar-agent

docs/extensions.md

@@ -58,28 +58,71 @@ interface Extension {
   onPostToolUse?: (payload: PostToolUseHookPayload) => Promise<unknown> | unknown;
 }
 
-interface PostToolUseHookPayload {
-  /** Tool name (e.g., "bash", "file_edit_replace_string") */
-  toolName: string;
-  
-  /** Unique ID for this tool invocation */
-  toolCallId: string;
-  
-  /** Tool-specific arguments (structure varies by tool) */
-  args: unknown;
-  
-  /** Tool result (structure varies by tool) - can be modified and returned */
-  result: unknown;
-  
-  /** Workspace identifier */
-  workspaceId: string;
-  
-  /** Unix timestamp in milliseconds */
-  timestamp: number;
-  
-  /** Full workspace runtime access (see Runtime API below) */
-  runtime: Runtime;
-}
+// PostToolUseHookPayload is a discriminated union by toolName
+// Each tool has specific arg and result types:
+
+type PostToolUseHookPayload =
+  | {
+      toolName: "bash";
+      args: { script: string; timeout_secs?: number };
+      result: { success: true; output: string; exitCode: 0; wall_duration_ms: number }
+             | { success: false; output?: string; exitCode: number; error: string; wall_duration_ms: number };
+      toolCallId: string;
+      workspaceId: string;
+      timestamp: number;
+      runtime: Runtime;
+    }
+  | {
+      toolName: "file_read";
+      args: { filePath: string; offset?: number; limit?: number };
+      result: { success: true; file_size: number; modifiedTime: string; lines_read: number; content: string }
+             | { success: false; error: string };
+      toolCallId: string;
+      workspaceId: string;
+      timestamp: number;
+      runtime: Runtime;
+    }
+  | {
+      toolName: "file_edit_replace_string";
+      args: { file_path: string; old_string: string; new_string: string; replace_count?: number };
+      result: { success: true; diff: string; edits_applied: number }
+             | { success: false; error: string };
+      toolCallId: string;
+      workspaceId: string;
+      timestamp: number;
+      runtime: Runtime;
+    }
+  // ... other tools (file_edit_insert, propose_plan, todo_write, status_set, etc.)
+  | {
+      // Catch-all for unknown tools
+      toolName: string;
+      args: unknown;
+      result: unknown;
+      toolCallId: string;
+      workspaceId: string;
+      timestamp: number;
+      runtime: Runtime;
+    };
+```
+
+**Type safety**: When you check `payload.toolName`, TypeScript narrows the `args` and `result` types automatically:
+
+```typescript
+const extension: Extension = {
+  async onPostToolUse(payload) {
+    if (payload.toolName === "bash") {
+      // TypeScript knows: payload.args is { script: string; timeout_secs?: number }
+      // TypeScript knows: payload.result has { success, output?, error?, exitCode, wall_duration_ms }
+      const command = payload.args.script;
+      
+      if (!payload.result.success) {
+        const errorMsg = payload.result.error; // Type-safe access
+      }
+    }
+    
+    return payload.result;
+  }
+};
 ```
 
 ## Runtime API

src/services/extensions/extensionHost.ts

@@ -15,7 +15,7 @@ import type {
   Extension,
   ExtensionInfo,
   ExtensionHostApi,
-  PostToolUseHookPayload,
+  ToolUsePayload,
 } from "../../types/extensions";
 import { NodeIpcProcessTransport } from "./nodeIpcTransport";
 
@@ -101,7 +101,7 @@ class ExtensionHostImpl extends RpcTarget implements ExtensionHostApi {
    * Dispatch post-tool-use hook to the extension
    * @returns The (possibly modified) tool result, or undefined if unchanged
    */
-  async onPostToolUse(payload: Omit<PostToolUseHookPayload, "runtime">): Promise<unknown> {
+  async onPostToolUse(payload: ToolUsePayload): Promise<unknown> {
     if (!this.extensionModule || !this.extensionModule.onPostToolUse) {
       // Extension doesn't have this hook - return result unchanged
       return payload.result;

src/services/extensions/extensionManager.ts

@@ -17,7 +17,7 @@ import { promises as fs } from "fs";
 import type { WorkspaceMetadata } from "@/types/workspace";
 import type { RuntimeConfig } from "@/types/runtime";
 import type {
-  PostToolUseHookPayload,
+  ToolUsePayload,
   ExtensionInfo,
   ExtensionHostApi,
 } from "@/types/extensions";
@@ -338,7 +338,7 @@ export class ExtensionManager {
    */
   async postToolUse(
     workspaceId: string,
-    payload: Omit<PostToolUseHookPayload, "runtime">
+    payload: ToolUsePayload
   ): Promise<unknown> {
     if (this.hosts.size === 0) {
       // No extensions loaded - return original result

src/types/extensions.ts

@@ -1,6 +1,24 @@
 import type { Runtime } from "@/runtime/Runtime";
 import type { RuntimeConfig } from "./runtime";
 import type { RpcTarget } from "capnweb";
+import type {
+  BashToolArgs,
+  BashToolResult,
+  FileReadToolArgs,
+  FileReadToolResult,
+  FileEditReplaceStringToolArgs,
+  FileEditReplaceStringToolResult,
+  FileEditReplaceLinesToolArgs,
+  FileEditReplaceLinesToolResult,
+  FileEditInsertToolArgs,
+  FileEditInsertToolResult,
+  ProposePlanToolArgs,
+  ProposePlanToolResult,
+  TodoWriteToolArgs,
+  TodoWriteToolResult,
+  StatusSetToolArgs,
+  StatusSetToolResult,
+} from "./tools";
 
 /**
  * Extension manifest structure (manifest.json)
@@ -10,17 +28,90 @@ export interface ExtensionManifest {
 }
 
 /**
- * Hook payload for post-tool-use hook
+ * Tool execution payload - discriminated union by tool name
  */
-export interface PostToolUseHookPayload {
-  toolName: string;
-  toolCallId: string;
-  args: unknown;
-  result: unknown;
-  workspaceId: string;
-  timestamp: number;
+export type ToolUsePayload =
+  | {
+      toolName: "bash";
+      toolCallId: string;
+      args: BashToolArgs;
+      result: BashToolResult;
+      workspaceId: string;
+      timestamp: number;
+    }
+  | {
+      toolName: "file_read";
+      toolCallId: string;
+      args: FileReadToolArgs;
+      result: FileReadToolResult;
+      workspaceId: string;
+      timestamp: number;
+    }
+  | {
+      toolName: "file_edit_replace_string";
+      toolCallId: string;
+      args: FileEditReplaceStringToolArgs;
+      result: FileEditReplaceStringToolResult;
+      workspaceId: string;
+      timestamp: number;
+    }
+  | {
+      toolName: "file_edit_replace_lines";
+      toolCallId: string;
+      args: FileEditReplaceLinesToolArgs;
+      result: FileEditReplaceLinesToolResult;
+      workspaceId: string;
+      timestamp: number;
+    }
+  | {
+      toolName: "file_edit_insert";
+      toolCallId: string;
+      args: FileEditInsertToolArgs;
+      result: FileEditInsertToolResult;
+      workspaceId: string;
+      timestamp: number;
+    }
+  | {
+      toolName: "propose_plan";
+      toolCallId: string;
+      args: ProposePlanToolArgs;
+      result: ProposePlanToolResult;
+      workspaceId: string;
+      timestamp: number;
+    }
+  | {
+      toolName: "todo_write";
+      toolCallId: string;
+      args: TodoWriteToolArgs;
+      result: TodoWriteToolResult;
+      workspaceId: string;
+      timestamp: number;
+    }
+  | {
+      toolName: "status_set";
+      toolCallId: string;
+      args: StatusSetToolArgs;
+      result: StatusSetToolResult;
+      workspaceId: string;
+      timestamp: number;
+    }
+  | {
+      // Catch-all for unknown tools
+      toolName: string;
+      toolCallId: string;
+      args: unknown;
+      result: unknown;
+      workspaceId: string;
+      timestamp: number;
+    };
+
+/**
+ * Hook payload for post-tool-use hook with Runtime access
+ * This adds the runtime field to each variant of ToolUsePayload
+ */
+export type PostToolUseHookPayload = ToolUsePayload & {
   runtime: Runtime; // Extensions get full workspace access via Runtime
-}
+};
 
 /**
  * Extension export interface - what extensions must export as default
@@ -31,7 +122,7 @@ export interface Extension {
    * Extensions can monitor, log, or modify the tool result.
    * 
    * @param payload - Tool execution context with full Runtime access
-   * @returns The tool result (can be modified) or undefined to leave unchanged
+   * @returns The tool result (modified or unmodified). Return undefined to leave unchanged.
    */
   onPostToolUse?: (payload: PostToolUseHookPayload) => Promise<unknown> | unknown;
 }
@@ -81,7 +172,7 @@ export interface ExtensionHostApi extends RpcTarget {
    * @param payload Hook payload (runtime will be added by host)
    * @returns The (possibly modified) tool result, or undefined if unchanged
    */
-  onPostToolUse(payload: Omit<PostToolUseHookPayload, "runtime">): Promise<unknown>;
+  onPostToolUse(payload: ToolUsePayload): Promise<unknown>;
 
   /**
    * Gracefully shutdown the extension host