refactor: update implementation

Author: catrielmuller

packages/kilo-gateway/src/server/routes.ts

@@ -530,8 +530,24 @@ export function createKiloRoutes(deps: KiloRoutesDeps) {
               "application/json": {
                 schema: resolver(
                   z.object({
+                    // `recovering` and `restoring` are transitional states the
+                    // worker reports while it brings an instance back online
+                    // after an unexpected stop or a snapshot restore — see
+                    // cloud `services/kiloclaw/src/index.ts` and the
+                    // `PlatformStatusResponse` type in
+                    // cloud/apps/web/src/lib/kiloclaw/types.ts. Keeping them in
+                    // the enum so the SDK types stay accurate.
                     status: z
-                      .enum(["provisioned", "starting", "restarting", "running", "stopped", "destroying"])
+                      .enum([
+                        "provisioned",
+                        "starting",
+                        "restarting",
+                        "recovering",
+                        "running",
+                        "stopped",
+                        "destroying",
+                        "restoring",
+                      ])
                       .nullable(),
                     sandboxId: z.string().optional(),
                     flyRegion: z.string().optional(),

packages/kilo-vscode/src/kiloclaw/KiloClawProvider.ts

@@ -523,9 +523,15 @@ export class KiloClawProvider implements vscode.Disposable {
     this.chatSubs.push(
       events.on("conversation.created", (ctx, e: ConversationCreatedEvent) => {
         if (!this.sandboxId || ctx !== `/kiloclaw/${this.sandboxId}`) return
-        // Refetch so we capture title/joinedAt/etc.
+        // Newer servers include the full conversation snapshot — splice it in
+        // immediately so the list updates without a roundtrip. Fall back to a
+        // refetch when the snapshot is absent (older servers / safety net).
+        if (e.conversation) {
+          this.conversations = mergeConversations([e.conversation], this.conversations)
+          this.broadcastConversations({ replace: true })
+          return
+        }
         void this.refreshConversations()
-        void e
       }),
     )
 
@@ -754,6 +760,9 @@ export class KiloClawProvider implements vscode.Disposable {
   private async createConversation(title?: string): Promise<void> {
     if (!this.chat || !this.sandboxId) return
     try {
+      // The server now returns the full conversation snapshot alongside
+      // `conversationId`. We rely on `refreshConversations()` to pick up the
+      // canonical list-item shape rather than mapping the detail payload here.
       const res = await this.chat.createConversation({ sandboxId: this.sandboxId, title })
       await this.refreshConversations()
       await this.selectConversation(res.conversationId)
@@ -904,7 +913,7 @@ export class KiloClawProvider implements vscode.Disposable {
       if (this.activeConversationId !== conversationId) return
       const sorted = sortMessagesAscending(res.messages)
       this.messages = mergeMessages(sorted, this.messages)
-      this.hasMoreMessages = res.messages.length >= MESSAGES_PAGE
+      this.hasMoreMessages = res.hasMore
       this.broadcastMessages({ replace: true })
     } catch (err) {
       this.post({ type: "kiloclaw.error", error: this.formatError(err, "Failed to load messages") })
@@ -1005,8 +1014,13 @@ export class KiloClawProvider implements vscode.Disposable {
 
   private async markRead(conversationId: string): Promise<void> {
     if (!this.chat) return
+    if (conversationId !== this.activeConversationId) return
+    // The mark-read endpoint requires `lastSeenMessageId`. With no messages
+    // loaded there is nothing to mark — silently skip.
+    const last = lastNonPendingMessageId(this.messages)
+    if (!last) return
     try {
-      await this.chat.markConversationRead(conversationId)
+      await this.chat.markConversationRead(conversationId, { lastSeenMessageId: last })
     } catch (err) {
       void err
     }
@@ -1043,7 +1057,7 @@ export class KiloClawProvider implements vscode.Disposable {
       // only apply the messages if the active conversation is still `target`.
       if (this.activeConversationId !== target) return
       this.messages = sortMessagesAscending(res.messages)
-      this.hasMoreMessages = res.messages.length >= MESSAGES_PAGE
+      this.hasMoreMessages = res.hasMore
       this.broadcastMessages({ replace: true })
     } catch (err) {
       const message = err instanceof Error ? err.message : String(err)
@@ -1263,3 +1277,19 @@ function mergeConversations(
 function sortMessagesAscending(messages: Message[]): Message[] {
   return [...messages].sort((a, b) => (a.id < b.id ? -1 : a.id > b.id ? 1 : 0))
 }
+
+/**
+ * Pick the latest server-confirmed message id from an ascending list. Pending
+ * (optimistic) messages use `pending-<clientId>` ids that the server doesn't
+ * recognise, so they're skipped — the new mark-read contract requires a real
+ * message id.
+ */
+function lastNonPendingMessageId(messages: Message[]): string | null {
+  for (let i = messages.length - 1; i >= 0; i--) {
+    const m = messages[i]
+    if (!m) continue
+    if (m.id.startsWith("pending-")) continue
+    return m.id
+  }
+  return null
+}

packages/kilo-vscode/src/kiloclaw/event-service-client.ts

@@ -2,9 +2,11 @@
  * Event Service WebSocket client for the VS Code extension host.
  *
  * Minimal inline port of `@kilocode/event-service` (cloud monorepo). Connects
- * to the kilo events Cloudflare Worker over WebSocket, authenticates with a
- * JWT carried in the `kilo.jwt.<base64url>` subprotocol header, subscribes to
- * per-conversation contexts, and dispatches server events to typed handlers.
+ * to the kilo events Cloudflare Worker using a two-step ticket flow:
+ *   1. POST `/connect-ticket` with `Authorization: Bearer <JWT>` to mint a
+ *      single-use ticket (30 s TTL).
+ *   2. Open WebSocket to `/connect?ticket=<ticket>` with subprotocol
+ *      `kilo.events.v1`.
  *
  * Runs in Node.js (the VS Code extension host). `WebSocket` is available in
  * Node 22+ without any import, matching the environment used elsewhere in
@@ -13,9 +15,10 @@
 
 import type { KiloChatEventMap, KiloChatEventName } from "./types"
 
-const SUBPROTOCOL_PREFIX = "kilo.jwt."
+const WS_SUBPROTOCOL = "kilo.events.v1"
 const HANDSHAKE_TIMEOUT_MS = 10_000
 const PING_INTERVAL_MS = 15_000
+const TICKET_FETCH_TIMEOUT_MS = 10_000
 
 export class WebSocketAuthError extends Error {
   constructor(message = "WebSocket authentication failed") {
@@ -58,10 +61,16 @@ export type EventServiceConfig = {
   onUnauthorized?: () => void
 }
 
-function encodeBase64Url(input: string): string {
-  // In the Node extension host btoa is available and JWTs are ASCII.
-  const b64 = typeof btoa === "function" ? btoa(input) : Buffer.from(input, "utf8").toString("base64")
-  return b64.replace(/\+/g, "-").replace(/\//g, "_").replace(/=+$/, "")
+/**
+ * The event-service base URL is configured as a WebSocket URL (`wss://…` /
+ * `ws://…`) but the connect-ticket endpoint is a plain HTTP request. Strip
+ * the trailing slash and swap the protocol so `fetch()` accepts the URL.
+ */
+function toHttpBase(wsBase: string): string {
+  const trimmed = wsBase.replace(/\/$/, "")
+  if (trimmed.startsWith("wss://")) return "https://" + trimmed.slice(6)
+  if (trimmed.startsWith("ws://")) return "http://" + trimmed.slice(5)
+  return trimmed
 }
 
 export class EventServiceClient {
@@ -182,10 +191,10 @@ export class EventServiceClient {
     }
 
     const token = await this.getToken()
-    const subprotocol = `${SUBPROTOCOL_PREFIX}${encodeBase64Url(token)}`
+    const ticket = await this.fetchTicket(token)
 
     return new Promise<void>((resolve, reject) => {
-      const ws = new WebSocket(`${this.url}/connect`, [subprotocol])
+      const ws = new WebSocket(`${this.url}/connect?ticket=${encodeURIComponent(ticket)}`, [WS_SUBPROTOCOL])
       this.ws = ws
 
       let settled = false
@@ -259,6 +268,46 @@ export class EventServiceClient {
     })
   }
 
+  /**
+   * Mint a single-use connection ticket. The event-service issues a 30 s ticket
+   * scoped to the bearer JWT; the WebSocket upgrade then consumes it. We
+   * surface 401/403 as `WebSocketAuthError` so the caller can drop the cached
+   * token and prompt re-auth.
+   *
+   * `this.url` is the WebSocket base (`wss://…` or `ws://…`); `fetch()` only
+   * accepts `http(s)`, so we rewrite the protocol before the HTTP call.
+   */
+  private async fetchTicket(token: string): Promise<string> {
+    const ctrl = new AbortController()
+    const timer = setTimeout(() => ctrl.abort(), TICKET_FETCH_TIMEOUT_MS)
+    try {
+      const res = await fetch(toHttpBase(this.url) + "/connect-ticket", {
+        method: "POST",
+        headers: { Authorization: `Bearer ${token}` },
+        signal: ctrl.signal,
+      })
+      if (res.status === 401 || res.status === 403) {
+        throw new WebSocketAuthError(`Event-service rejected ticket request: ${res.status}`)
+      }
+      if (!res.ok) {
+        throw new WebSocketConnectError(`Failed to mint event-service ticket: ${res.status}`, res.status)
+      }
+      const body = (await res.json().catch(() => null)) as { ticket?: unknown } | null
+      if (!body || typeof body.ticket !== "string" || !body.ticket) {
+        throw new WebSocketConnectError("Malformed event-service ticket response", 0)
+      }
+      return body.ticket
+    } catch (err) {
+      if (err instanceof WebSocketAuthError || err instanceof WebSocketConnectError) throw err
+      if ((err as { name?: string })?.name === "AbortError") {
+        throw new HandshakeTimeoutError()
+      }
+      throw new WebSocketConnectError(`Event-service ticket request failed: ${(err as Error)?.message ?? err}`, 0)
+    } finally {
+      clearTimeout(timer)
+    }
+  }
+
   private clearHandshakeTimer(): void {
     if (this.handshakeTimer !== null) {
       clearTimeout(this.handshakeTimer)

packages/kilo-vscode/src/kiloclaw/kilo-chat-client.ts

@@ -93,7 +93,10 @@ export class KiloChatClient {
     return this.request(`/v1/conversations/${conversationId}`)
   }
 
-  createConversation(req: { sandboxId: string; title?: string }): Promise<{ conversationId: string }> {
+  createConversation(req: {
+    sandboxId: string
+    title?: string
+  }): Promise<{ conversationId: string; conversation?: ConversationDetail }> {
     return this.request("/v1/conversations", { method: "POST", body: req })
   }
 
@@ -105,11 +108,23 @@ export class KiloChatClient {
   }
 
   async leaveConversation(conversationId: string): Promise<void> {
+    // Returns 200 JSON with `{ ok }`-style payload; we don't need the body.
     await this.request<unknown>(`/v1/conversations/${conversationId}/leave`, { method: "POST" })
   }
 
-  async markConversationRead(conversationId: string): Promise<void> {
-    await this.request<unknown>(`/v1/conversations/${conversationId}/mark-read`, { method: "POST" })
+  /**
+   * Mark messages up to `lastSeenMessageId` as read for the current user.
+   * The server enforces monotonic `lastReadAt` and returns whether the read
+   * pointer advanced plus whether the badge bucket was cleared.
+   */
+  markConversationRead(
+    conversationId: string,
+    req: { lastSeenMessageId: string },
+  ): Promise<{ ok: boolean; applied: boolean; lastReadAt: number; badgeClear: boolean }> {
+    return this.request(`/v1/conversations/${conversationId}/mark-read`, {
+      method: "POST",
+      body: req,
+    })
   }
 
   // ── Messages ─────────────────────────────────────────────────────
@@ -119,10 +134,10 @@ export class KiloChatClient {
     content: ContentBlock[]
     inReplyToMessageId?: string
     clientId?: string
-  }): Promise<{ messageId: string; clientId?: string }> {
+  }): Promise<{ messageId: string; clientId?: string; message?: Message }> {
     const prev = this.sendQueues.get(req.conversationId) ?? Promise.resolve()
     const send = () =>
-      this.request<{ messageId: string; clientId?: string }>("/v1/messages", {
+      this.request<{ messageId: string; clientId?: string; message?: Message }>("/v1/messages", {
         method: "POST",
         body: req,
       })
@@ -140,18 +155,22 @@ export class KiloChatClient {
   editMessage(
     messageId: string,
     req: { conversationId: string; content: ContentBlock[]; timestamp: number },
-  ): Promise<{ messageId?: string }> {
+  ): Promise<{ messageId?: string; message?: Message }> {
     return this.request(`/v1/messages/${messageId}`, { method: "PATCH", body: req })
   }
 
   async deleteMessage(messageId: string, conversationId: string): Promise<void> {
+    // Returns 200 JSON with `{ ok }`-style payload; we don't need the body.
     await this.request<unknown>(`/v1/messages/${messageId}`, {
       method: "DELETE",
       query: { conversationId },
     })
   }
 
-  listMessages(conversationId: string, opts?: { before?: string; limit?: number }): Promise<{ messages: Message[] }> {
+  listMessages(
+    conversationId: string,
+    opts?: { before?: string; limit?: number },
+  ): Promise<{ messages: Message[]; hasMore: boolean; nextCursor: string | null }> {
     return this.request(`/v1/conversations/${conversationId}/messages`, {
       query: { before: opts?.before, limit: opts?.limit },
     })
@@ -161,7 +180,7 @@ export class KiloChatClient {
     conversationId: string,
     messageId: string,
     req: { groupId: string; value: ExecApprovalDecision },
-  ): Promise<{ ok: true }> {
+  ): Promise<{ ok?: boolean; message?: Message; content?: ContentBlock[] }> {
     return this.request(`/v1/conversations/${conversationId}/messages/${messageId}/execute-action`, {
       method: "POST",
       body: req,
@@ -170,15 +189,24 @@ export class KiloChatClient {
 
   // ── Reactions ────────────────────────────────────────────────────
 
-  addReaction(messageId: string, req: { conversationId: string; emoji: string }): Promise<{ id: string }> {
+  addReaction(
+    messageId: string,
+    req: { conversationId: string; emoji: string },
+  ): Promise<{ id: string; operationId?: string }> {
     return this.request(`/v1/messages/${messageId}/reactions`, { method: "POST", body: req })
   }
 
-  async removeReaction(messageId: string, req: { conversationId: string; emoji: string }): Promise<void> {
-    await this.request<unknown>(`/v1/messages/${messageId}/reactions`, {
-      method: "DELETE",
-      query: req,
-    })
+  async removeReaction(
+    messageId: string,
+    req: { conversationId: string; emoji: string },
+  ): Promise<{ removed: boolean; id: string | null; operationId?: string }> {
+    return this.request<{ removed: boolean; id: string | null; operationId?: string }>(
+      `/v1/messages/${messageId}/reactions`,
+      {
+        method: "DELETE",
+        query: req,
+      },
+    )
   }
 
   // ── Typing ───────────────────────────────────────────────────────

packages/kilo-vscode/src/kiloclaw/types.ts

@@ -13,7 +13,19 @@
 // ── Instance status (KiloClaw worker) ───────────────────────────────
 
 export type ClawStatus = {
-  status: "provisioned" | "starting" | "restarting" | "running" | "stopped" | "destroying" | null
+  // `recovering` and `restoring` are transitional states the worker reports
+  // while bringing an instance back from an unexpected stop or a snapshot
+  // restore (cloud: `services/kiloclaw/src/index.ts`).
+  status:
+    | "provisioned"
+    | "starting"
+    | "restarting"
+    | "recovering"
+    | "running"
+    | "stopped"
+    | "destroying"
+    | "restoring"
+    | null
   sandboxId?: string
   flyRegion?: string
   machineSize?: { cpus: number; memory_mb: number }
@@ -124,12 +136,25 @@ export type ConversationStatusRecord = {
 // ── Typed Kilo Chat events (server → client) ───────────────────────
 // Event names mirror `@kilocode/kilo-chat/events`.
 
+/**
+ * Snapshot of the message that was replied to. Server includes this on
+ * `message.created` so clients can render a reply preview without a follow-up
+ * fetch. `deleted` mirrors the soft-deletion state at the time of replying.
+ */
+export type ReplyToSnapshot = {
+  messageId: string
+  senderId: string
+  content: ContentBlock[]
+  deleted?: boolean
+}
+
 export type MessageCreatedEvent = {
   messageId: string
   senderId: string
   content: ContentBlock[]
   inReplyToMessageId: string | null
   clientId?: string
+  replyTo?: ReplyToSnapshot | null
 }
 
 export type MessageUpdatedEvent = {
@@ -144,10 +169,18 @@ export type MessageDeliveryFailedEvent = { messageId: string }
 export type TypingEvent = { memberId: string }
 export type TypingStopEvent = { memberId: string }
 
-export type ReactionAddedEvent = { messageId: string; memberId: string; emoji: string }
-export type ReactionRemovedEvent = { messageId: string; memberId: string; emoji: string }
+export type ReactionAddedEvent = { messageId: string; memberId: string; emoji: string; operationId?: string }
+export type ReactionRemovedEvent = { messageId: string; memberId: string; emoji: string; operationId?: string }
 
-export type ConversationCreatedEvent = { conversationId: string }
+/**
+ * Server fans out the full conversation snapshot on `conversation.created` so
+ * clients can append to their list without a follow-up fetch. Older servers may
+ * still send only the `conversationId`, so the snapshot is optional.
+ */
+export type ConversationCreatedEvent = {
+  conversationId: string
+  conversation?: ConversationListItem
+}
 export type ConversationRenamedEvent = { conversationId: string; title: string }
 export type ConversationLeftEvent = { conversationId: string }
 export type ConversationReadEvent = { conversationId: string; memberId: string; lastReadAt: number }