diff --git a/src/dispatch.ts b/src/dispatch.ts
index 468422e..a52b08e 100644
--- a/src/dispatch.ts
+++ b/src/dispatch.ts
@@ -2,7 +2,7 @@
 // Dispatch Messages received via WebSocket
 
 import { basename } from "path"
-import type { Message } from "./shared/types"
+import type { Message, InputMessage, SaveFileMessage } from "./shared/types"
 import { runCommand } from "./shell"
 import { send } from "./websocket"
 import { setState } from "./state"
@@ -11,10 +11,10 @@ export async function dispatchMessage(ws: any, msg: Message) {
   console.log("<- receive", msg)
   switch (msg.type) {
     case "input":
-      await inputMessage(ws, msg); break
+      await inputMessage(ws, msg as InputMessage); break
 
     case "save-file":
-      await saveFileMessage(ws, msg); break
+      await saveFileMessage(ws, msg as SaveFileMessage); break
 
     case "ui:mode":
       setState("ui:mode", msg.data); break
@@ -24,8 +24,8 @@ export async function dispatchMessage(ws: any, msg: Message) {
   }
 }
 
-async function inputMessage(ws: any, msg: Message) {
-  const result = await runCommand(msg.session || "", msg.id || "", msg.data as string, ws)
+async function inputMessage(ws: any, msg: InputMessage) {
+  const result = await runCommand(msg.session, msg.id, msg.data as string, ws)
 
   if (typeof result.output === "object" && "game" in result.output) {
     send(ws, { id: msg.id, type: "game:start", data: result.output.game })
@@ -34,7 +34,7 @@ async function inputMessage(ws: any, msg: Message) {
   }
 }
 
-async function saveFileMessage(ws: any, msg: Message) {
+async function saveFileMessage(ws: any, msg: SaveFileMessage) {
   if (msg.id && typeof msg.data === "string") {
     await Bun.write(msg.id.replace("..", ""), msg.data, { createPath: true })
     send(ws, { type: "output", data: { status: "ok", output: `saved ${basename(msg.id)}` } })
diff --git a/src/shared/types.ts b/src/shared/types.ts
index f0d0819..8a9eea9 100644
--- a/src/shared/types.ts
+++ b/src/shared/types.ts
@@ -3,7 +3,10 @@ export type Message = {
     id?: string
     type: MessageType
     data?: CommandResult | CommandOutput
-} | SessionUpdateMessage
+}
+    | InputMessage
+    | SaveFileMessage
+    | SessionUpdateMessage
 
 export type MessageType = "error" | "input" | "output" | "commands" | "save-file"
     | "game:start"
@@ -21,7 +24,21 @@ export type CommandResult = {
     output: CommandOutput
 }
 
-type SessionUpdateMessage = {
+export type InputMessage = {
+    type: "input"
+    id: string
+    session: string
+    data: CommandResult | CommandOutput
+}
+
+export type SaveFileMessage = {
+    type: "save-file"
+    id: string
+    session: string
+    data: CommandResult | CommandOutput
+}
+
+export type SessionUpdateMessage = {
     type: "session:update",
     data: Record<string, string>
 }
\ No newline at end of file