From d26624ab9342e619c840ac8589a8346305eb382c Mon Sep 17 00:00:00 2001 From: Rob Colbert Date: Sat, 9 May 2026 09:58:47 -0400 Subject: [PATCH] chat session auto-naming with IDE update --- docs/socket-protocol.md | 23 +- gadget-code/frontend/src/lib/socket.ts | 6 + .../frontend/src/pages/ChatSessionView.tsx | 6 + gadget-code/src/lib/code-session.ts | 22 +- gadget-code/src/services/chat-session.ts | 16 +- gadget-code/tests/code-session.test.ts | 284 ++++++++++++------ gadget-drone/src/gadget-drone.ts | 16 +- packages/ai/src/ollama.ts | 21 +- packages/api/src/index.ts | 1 + packages/api/src/messages/socket.ts | 2 + packages/api/src/messages/web.ts | 16 + 11 files changed, 292 insertions(+), 121 deletions(-) create mode 100644 packages/api/src/messages/web.ts diff --git a/docs/socket-protocol.md b/docs/socket-protocol.md index a8459df..e230c47 100644 --- a/docs/socket-protocol.md +++ b/docs/socket-protocol.md @@ -34,6 +34,9 @@ Defined in `packages/api/src/messages/socket.ts`. * `crashRecoveryResponse`: Command to `discard` or `retry` a stalled work order. * `requestTermination`: Command to immediately terminate the drone process. +### Web -> IDE (Server to Client) +* `sessionUpdated`: Notify the IDE that a chat session property has changed (e.g. auto-generated name). + --- ## 3. Core Sequences & Routing @@ -42,8 +45,12 @@ Defined in `packages/api/src/messages/socket.ts`. 1. **IDE** emits `submitPrompt(content)`. 2. **Web (`CodeSession.ts`)**: * Creates a `ChatTurn` document (status: `processing`). + * Increments the chat session's `stats.turnCount`. * Finds the target `DroneSession`. + * Caches the updated session and signals the **IDE** to enter Processing state. * Emits `processWorkOrder` to the **Drone**. + * On first prompt (name is still the default), calls AI API to auto-generate session name. + * Emits `sessionUpdated({ name })` to **IDE** if the name changed. 3. **Drone (`gadget-drone.ts`)**: * Writes a local `.gadget/work-order.json` cache (for crash recovery). * Calls `AgentService.process()`. @@ -117,6 +124,13 @@ type RequestTerminationMessage = ( ) => void; ``` +### Web -> IDE +```typescript +type SessionUpdatedMessage = ( + updates: Partial +) => void; +``` + ### Drone -> Web (Streaming) ```typescript type ThinkingMessage = (content: string) => void; @@ -183,8 +197,9 @@ All indexes are kept in sync during connection and disconnection. ## 7. Extending the Protocol To add a new message: -1. Add the message type to `packages/api/src/messages/ide.ts` or `drone.ts`. +1. Add the message type to `packages/api/src/messages/ide.ts`, `drone.ts`, or `web.ts`. 2. Register it in `ClientToServerEvents` or `ServerToClientEvents` in `packages/api/src/messages/socket.ts`. -3. Implement the sender (emit) in the Client (`ide` or `drone`). -4. Implement the handler in the corresponding `CodeSession` or `DroneSession` on the Web server. -5. Implement the forward-path routing if needed. +3. Re-export from `packages/api/src/index.ts`. +4. Implement the sender (emit) in the Client (`ide` or `drone`) or Server (`CodeSession`/`DroneSession`). +5. Implement the handler in the corresponding class or frontend component. +6. Implement the forward-path routing if needed. diff --git a/gadget-code/frontend/src/lib/socket.ts b/gadget-code/frontend/src/lib/socket.ts index b33cbfe..505d8ac 100644 --- a/gadget-code/frontend/src/lib/socket.ts +++ b/gadget-code/frontend/src/lib/socket.ts @@ -1,5 +1,6 @@ import { createContext } from "react"; import { io, Socket } from "socket.io-client"; +import type { ChatSession } from "./api"; const SOCKET_URL = ""; @@ -90,6 +91,7 @@ export interface SocketEvents { role: "user" | "assistant" | "system"; }) => void; workspaceModeChanged: (mode: string) => void; + sessionUpdated: (updates: Partial) => void; connect: () => void; disconnect: (reason: string) => void; error: (error: Error) => void; @@ -161,6 +163,10 @@ class SocketClient { this.emit("workspaceModeChanged", mode); }); + this.socket.on("sessionUpdated", (updates: unknown) => { + this.emit("sessionUpdated", updates as Partial); + }); + this.socket.on( "log", ( diff --git a/gadget-code/frontend/src/pages/ChatSessionView.tsx b/gadget-code/frontend/src/pages/ChatSessionView.tsx index 4159d5a..5e7182b 100644 --- a/gadget-code/frontend/src/pages/ChatSessionView.tsx +++ b/gadget-code/frontend/src/pages/ChatSessionView.tsx @@ -189,12 +189,17 @@ export default function ChatSessionView() { } }; + const handleSessionUpdated = useCallback((updates: Partial) => { + setSession(prev => prev ? { ...prev, ...updates } : null); + }, []); + const setupSocketListeners = () => { socketClient.on('thinking', handleThinking); socketClient.on('response', handleResponse); socketClient.on('toolCall', handleToolCall); socketClient.on('workOrderComplete', handleWorkOrderComplete); socketClient.on('workspaceModeChanged', handleWorkspaceModeChanged); + socketClient.on('sessionUpdated', handleSessionUpdated); socketClient.on('log:entry', handleLogEntry); socketClient.on('status', handleStatus); }; @@ -205,6 +210,7 @@ export default function ChatSessionView() { socketClient.off('toolCall', handleToolCall); socketClient.off('workOrderComplete', handleWorkOrderComplete); socketClient.off('workspaceModeChanged', handleWorkspaceModeChanged); + socketClient.off('sessionUpdated', handleSessionUpdated); socketClient.off('log:entry', handleLogEntry); socketClient.off('status', handleStatus); }; diff --git a/gadget-code/src/lib/code-session.ts b/gadget-code/src/lib/code-session.ts index a649a8c..c942a1a 100644 --- a/gadget-code/src/lib/code-session.ts +++ b/gadget-code/src/lib/code-session.ts @@ -253,13 +253,23 @@ export class CodeSession extends SocketSession { ); /* - * Call out to have the session's name auto-generated from this prompt - * if this is the first prompt. + * Auto-generate a session name from the first prompt. Only do this when + * the name is still the default (user hasn't set a custom name) and + * we're on the first turn (turnCount === 1 after the increment above). */ - this.chatSession = await ChatSessionService.generateSessionNameFromPrompt( - this.chatSession, - content, - ); + if ( + this.chatSession.name === "New Chat Session" && + this.chatSession.stats.turnCount === 1 + ) { + this.chatSession = + await ChatSessionService.generateSessionNameFromPrompt( + this.chatSession, + content, + ); + const update: Partial = { name: this.chatSession.name }; + this.log.debug("emitting sessionUpdated message", { update }); + this.socket.emit("sessionUpdated", update); + } } catch (error) { this.log.error("prompt rejected", { error }); cb(false, {}); diff --git a/gadget-code/src/services/chat-session.ts b/gadget-code/src/services/chat-session.ts index 0e67274..d46f3f1 100644 --- a/gadget-code/src/services/chat-session.ts +++ b/gadget-code/src/services/chat-session.ts @@ -21,12 +21,12 @@ import { import { IAiProvider } from "@gadget/api"; -import Project from "../models/project.js"; -import ChatTurn from "../models/chat-turn.js"; -import ChatSession from "../models/chat-session.js"; -import AiProvider from "../models/ai-provider.js"; +import Project from "../models/project.ts"; +import ChatTurn from "../models/chat-turn.ts"; +import ChatSession from "../models/chat-session.ts"; +import AiProvider from "../models/ai-provider.ts"; -import { DtpService } from "../lib/service.js"; +import { DtpService } from "../lib/service.ts"; import { PopulateOptions } from "mongoose"; import { AiApi, @@ -425,7 +425,7 @@ class ChatSessionService extends DtpService { const newSession = await ChatSession.findOneAndUpdate( { _id: session._id }, { $set: { name: response.response || "New Session" } }, - { new: true, populate: this.populateChatSession, lean: true }, + { new: true, populate: this.populateChatSession }, ); if (!newSession) { const error = new Error("chat session has been removed"); @@ -433,10 +433,6 @@ class ChatSessionService extends DtpService { throw error; } - //TODO: emit the `sessionUpdated` message to the CodeSession in the IDE - // for this chat session, letting it know the name of the chat session has - // changed. The IDE will then update it's displays and state for the User. - return newSession; } diff --git a/gadget-code/tests/code-session.test.ts b/gadget-code/tests/code-session.test.ts index abeadde..ca90ff9 100644 --- a/gadget-code/tests/code-session.test.ts +++ b/gadget-code/tests/code-session.test.ts @@ -9,12 +9,16 @@ import { ChatTurnStatus, } from "@gadget/api"; import SocketService from "../src/services/socket"; +import ChatSessionService from "../src/services/chat-session"; import { ChatTurn } from "../src/models/chat-turn"; +import ChatSession from "../src/models/chat-session"; import { nanoid } from "nanoid"; // Mock dependencies vi.mock("../src/services/socket"); +vi.mock("../src/services/chat-session"); vi.mock("../src/models/chat-turn"); +vi.mock("../src/models/chat-session"); describe("CodeSession", () => { let mockSocket: any; @@ -54,12 +58,18 @@ describe("CodeSession", () => { mockChatSession = { _id: nanoid(), - name: "Test Session", + name: "New Chat Session", mode: "build", provider: nanoid(), selectedModel: "llama3.1", user: mockUser, project: mockProject, + stats: { + turnCount: 0, + toolCallCount: 0, + inputTokens: 0, + outputTokens: 0, + }, } as IChatSession; codeSession = new CodeSession(mockSocket, mockUser); @@ -68,7 +78,6 @@ describe("CodeSession", () => { describe("setSelectedDrone", () => { it("should set the selected drone", () => { codeSession.setSelectedDrone(mockDrone); - // Can't directly access private property, but we can verify through behavior expect(() => codeSession.setSelectedDrone(mockDrone)).not.toThrow(); }); }); @@ -83,91 +92,12 @@ describe("CodeSession", () => { }); describe("onSubmitPrompt", () => { - it("should return error if no drone is selected", async () => { - codeSession.setChatSession(mockChatSession, mockProject); - const cb = vi.fn(); + let mockTurn: any; + let mockDroneSession: any; + let cb: any; - await codeSession.onSubmitPrompt("test prompt", cb); - - expect(cb).toHaveBeenCalledWith(false, { message: "No drone selected" }); - }); - - it("should return error if no chat session is active", async () => { - codeSession.setSelectedDrone(mockDrone); - const cb = vi.fn(); - - await codeSession.onSubmitPrompt("test prompt", cb); - - expect(cb).toHaveBeenCalledWith(false, { message: "No chat session active" }); - }); - - it("should return error if no project is selected", async () => { - codeSession.setSelectedDrone(mockDrone); - codeSession.setChatSession(mockChatSession, undefined as any); - const cb = vi.fn(); - - await codeSession.onSubmitPrompt("test prompt", cb); - - expect(cb).toHaveBeenCalledWith(false, { message: "No project selected" }); - }); - - it("should create a ChatTurn and emit processWorkOrder to drone", async () => { - codeSession.setSelectedDrone(mockDrone); - codeSession.setChatSession(mockChatSession, mockProject); - - const mockTurn = { - _id: nanoid(), - save: vi.fn().mockResolvedValue(undefined), - }; - vi.mocked(ChatTurn).mockImplementation(function () { - return mockTurn as any; - }); - (vi.mocked(ChatTurn) as any).populate = vi - .fn() - .mockResolvedValue(mockTurn); - - const mockDroneSession = { - socket: { - emit: vi.fn(), - }, - setChatSessionId: vi.fn(), - setCurrentTurnId: vi.fn(), - }; - vi.mocked(SocketService.getDroneSession).mockReturnValue( - mockDroneSession as any, - ); - - const cb = vi.fn(); - await codeSession.onSubmitPrompt("test prompt", cb); - - expect(ChatTurn).toHaveBeenCalledWith( - expect.objectContaining({ - user: mockUser._id, - project: mockProject._id, - session: mockChatSession._id, - provider: mockChatSession.provider, - llm: mockChatSession.selectedModel, - status: ChatTurnStatus.Processing, - prompts: expect.objectContaining({ - user: "test prompt", - }), - }), - ); - - expect(mockTurn.save).toHaveBeenCalled(); - expect(mockDroneSession.socket.emit).toHaveBeenCalledWith( - "processWorkOrder", - mockDrone, - mockTurn, - expect.any(Function), - ); - }); - - it("should update ChatTurn to Error status if drone rejects work order", async () => { - codeSession.setSelectedDrone(mockDrone); - codeSession.setChatSession(mockChatSession, mockProject); - - const mockTurn = { + beforeEach(() => { + mockTurn = { _id: nanoid(), status: ChatTurnStatus.Processing, errorMessage: "", @@ -180,12 +110,9 @@ describe("CodeSession", () => { .fn() .mockResolvedValue(mockTurn); - const mockDroneSession = { + mockDroneSession = { socket: { - emit: vi.fn((event, ...args) => { - const callback = args[args.length - 1]; - callback(false, "Drone is busy"); - }), + emit: vi.fn(), }, setChatSessionId: vi.fn(), setCurrentTurnId: vi.fn(), @@ -194,20 +121,189 @@ describe("CodeSession", () => { mockDroneSession as any, ); - const cb = vi.fn(); + vi.mocked(ChatSessionService.createTurn).mockResolvedValue( + mockTurn as any, + ); + + cb = vi.fn(); + + codeSession.setSelectedDrone(mockDrone); + codeSession.setChatSession(mockChatSession, mockProject); + }); + + it("should return error if no drone is selected", async () => { + codeSession = new CodeSession(mockSocket, mockUser); + codeSession.setChatSession(mockChatSession, mockProject); + + await codeSession.onSubmitPrompt("test prompt", cb); + + expect(cb).toHaveBeenCalledWith(false, { message: "No drone selected" }); + }); + + it("should return error if no chat session is active", async () => { + codeSession = new CodeSession(mockSocket, mockUser); + codeSession.setSelectedDrone(mockDrone); + + await codeSession.onSubmitPrompt("test prompt", cb); + + expect(cb).toHaveBeenCalledWith(false, { + message: "No chat session active", + }); + }); + + it("should return error if no project is selected", async () => { + codeSession.setChatSession(mockChatSession, undefined as any); + await codeSession.onSubmitPrompt("test prompt", cb); + + expect(cb).toHaveBeenCalledWith(false, { message: "No project selected" }); + }); + + it("should create a ChatTurn, increment turnCount, and emit processWorkOrder", async () => { + const updatedSession = { + ...mockChatSession, + stats: { ...mockChatSession.stats, turnCount: 1 }, + }; + vi.mocked(ChatSession.findOneAndUpdate).mockResolvedValue( + updatedSession as any, + ); + vi.mocked( + ChatSessionService.generateSessionNameFromPrompt, + ).mockResolvedValue(updatedSession as any); + + await codeSession.onSubmitPrompt("test prompt", cb); + + expect(ChatSessionService.createTurn).toHaveBeenCalledWith( + mockChatSession, + "test prompt", + ); + expect(ChatSession.findOneAndUpdate).toHaveBeenCalledWith( + { _id: mockChatSession._id }, + { $inc: { "stats.turnCount": 1 } }, + expect.objectContaining({ new: true }), + ); + expect(mockDroneSession.socket.emit).toHaveBeenCalledWith( + "processWorkOrder", + mockDrone, + mockTurn, + expect.any(Function), + ); + }); + + it("should update ChatTurn to Error status if drone rejects work order", async () => { + const updatedSession = { + ...mockChatSession, + stats: { ...mockChatSession.stats, turnCount: 1 }, + }; + vi.mocked(ChatSession.findOneAndUpdate).mockResolvedValue( + updatedSession as any, + ); + mockDroneSession.socket.emit = vi.fn((event: string, ...args: any[]) => { + const callback = args[args.length - 1]; + callback(false, "Drone is busy"); + }); + await codeSession.onSubmitPrompt("test prompt", cb); expect(mockTurn.status).toBe(ChatTurnStatus.Error); expect(mockTurn.errorMessage).toBe("Drone is busy"); expect(mockTurn.save).toHaveBeenCalled(); }); + + it("should reject prompt and delete turn if session was removed during increment", async () => { + vi.mocked(ChatSession.findOneAndUpdate).mockResolvedValue(null); + vi.mocked(ChatSessionService.delete).mockResolvedValue(undefined); + + await codeSession.onSubmitPrompt("test prompt", cb); + + expect(ChatSessionService.delete).toHaveBeenCalledWith(mockChatSession._id); + expect(cb).toHaveBeenCalledWith(false, {}); + }); + + it("should auto-generate session name and emit sessionUpdated on first prompt with default name", async () => { + const updatedSession = { + ...mockChatSession, + stats: { ...mockChatSession.stats, turnCount: 1 }, + }; + vi.mocked(ChatSession.findOneAndUpdate).mockResolvedValue( + updatedSession as any, + ); + + const namedSession = { + ...updatedSession, + name: "Build a REST API", + }; + vi.mocked( + ChatSessionService.generateSessionNameFromPrompt, + ).mockResolvedValue(namedSession as any); + + await codeSession.onSubmitPrompt("build a rest api", cb); + + expect( + ChatSessionService.generateSessionNameFromPrompt, + ).toHaveBeenCalledWith(updatedSession, "build a rest api"); + expect(mockSocket.emit).toHaveBeenCalledWith("sessionUpdated", { + name: "Build a REST API", + }); + }); + + it("should NOT auto-generate name if session already has a custom name", async () => { + const customSession = { + ...mockChatSession, + name: "My Custom Session", + }; + codeSession.setChatSession(customSession, mockProject); + + const updatedSession = { + ...customSession, + stats: { ...customSession.stats, turnCount: 1 }, + }; + vi.mocked(ChatSession.findOneAndUpdate).mockResolvedValue( + updatedSession as any, + ); + + await codeSession.onSubmitPrompt("test prompt", cb); + + expect( + ChatSessionService.generateSessionNameFromPrompt, + ).not.toHaveBeenCalled(); + expect(mockSocket.emit).not.toHaveBeenCalledWith( + "sessionUpdated", + expect.anything(), + ); + }); + + it("should NOT auto-generate name on subsequent prompts (turnCount > 1)", async () => { + const multiTurnSession = { + ...mockChatSession, + stats: { ...mockChatSession.stats, turnCount: 5 }, + }; + codeSession.setChatSession(multiTurnSession, mockProject); + + const updatedSession = { + ...multiTurnSession, + stats: { ...multiTurnSession.stats, turnCount: 6 }, + }; + vi.mocked(ChatSession.findOneAndUpdate).mockResolvedValue( + updatedSession as any, + ); + + await codeSession.onSubmitPrompt("test prompt", cb); + + expect( + ChatSessionService.generateSessionNameFromPrompt, + ).not.toHaveBeenCalled(); + expect(mockSocket.emit).not.toHaveBeenCalledWith( + "sessionUpdated", + expect.anything(), + ); + }); }); describe("onRequestSessionLock", () => { it("should set selected drone, chat session, and project on success", () => { const mockDroneSession = { socket: { - emit: vi.fn((event, ...args) => { + emit: vi.fn((event: string, ...args: any[]) => { const callback = args[args.length - 1]; callback(true, mockChatSession._id); }), @@ -233,7 +329,7 @@ describe("CodeSession", () => { it("should not set session data on failure", () => { const mockDroneSession = { socket: { - emit: vi.fn((event, ...args) => { + emit: vi.fn((event: string, ...args: any[]) => { const callback = args[args.length - 1]; callback(false, ""); }), diff --git a/gadget-drone/src/gadget-drone.ts b/gadget-drone/src/gadget-drone.ts index d203d21..265f92e 100644 --- a/gadget-drone/src/gadget-drone.ts +++ b/gadget-drone/src/gadget-drone.ts @@ -12,7 +12,12 @@ import AgentService, { IAgentWorkOrder } from "./services/agent.ts"; import AiService from "./services/ai.ts"; import PlatformService from "./services/platform.ts"; import WorkspaceService from "./services/workspace.ts"; -import { DroneStatus, GadgetLog, GadgetLogTransportSocket, IUser } from "@gadget/api"; +import { + DroneStatus, + GadgetLog, + GadgetLogTransportSocket, + IUser, +} from "@gadget/api"; import { GadgetProcess } from "./lib/process.ts"; import { @@ -210,7 +215,14 @@ class GadgetDrone extends GadgetProcess { const socketTransport = new GadgetLogTransportSocket( (event, timestamp, component, level, message, metadata): void => { - (this.socket as any)?.emit(event, timestamp, component, level, message, metadata); + (this.socket as any)?.emit( + event, + timestamp, + { name: component.name, slug: component.slug }, + level, + message, + metadata, + ); }, ); GadgetLog.addDefaultTransport(socketTransport); diff --git a/packages/ai/src/ollama.ts b/packages/ai/src/ollama.ts index a3d30bc..3201117 100644 --- a/packages/ai/src/ollama.ts +++ b/packages/ai/src/ollama.ts @@ -147,18 +147,27 @@ export class OllamaAiApi extends AiApi { stream: true, }); + const content = { + response: "", + thinking: "", + }; let lastChunk; for await (const chunk of response) { lastChunk = chunk; - if (streamCallback) { - if (chunk.thinking) { + if (chunk.thinking) { + content.thinking += chunk.thinking; + if (streamCallback) { await streamCallback({ type: "thinking", data: chunk.thinking, }); } - if (chunk.response) { + } + + if (chunk.response) { + content.response += chunk.response; + if (streamCallback) { await streamCallback({ type: "response", data: chunk.response, @@ -166,13 +175,15 @@ export class OllamaAiApi extends AiApi { } } } + + this.log.debug("generate call is done", content); assert(lastChunk, "no stream response chunks received"); return { done: lastChunk.done, doneReason: lastChunk.done_reason, - response: lastChunk.response, - thinking: lastChunk.thinking, + response: content.response, + thinking: content.thinking, stats: { duration: { seconds: lastChunk.total_duration, diff --git a/packages/api/src/index.ts b/packages/api/src/index.ts index 4d77750..25ebf07 100644 --- a/packages/api/src/index.ts +++ b/packages/api/src/index.ts @@ -30,6 +30,7 @@ export * from "./interfaces/user.ts"; export * from "./messages/ide.ts"; export * from "./messages/drone.ts"; +export * from "./messages/web.ts"; export * from "./messages/socket.ts"; /* diff --git a/packages/api/src/messages/socket.ts b/packages/api/src/messages/socket.ts index 25de98d..3765d22 100644 --- a/packages/api/src/messages/socket.ts +++ b/packages/api/src/messages/socket.ts @@ -15,6 +15,7 @@ import { WorkspaceModeChangedMessage, LogMessage, } from "./drone.ts"; +import { SessionUpdatedMessage } from "./web.ts"; import { ReleaseSessionLockMessage, RequestSessionLockMessage, @@ -91,6 +92,7 @@ export interface ServerToClientEvents { toolCall: ToolCallMessage; workOrderComplete: WorkOrderCompleteMessage; workspaceModeChanged: WorkspaceModeChangedMessage; + sessionUpdated: SessionUpdatedMessage; } export interface SocketData { diff --git a/packages/api/src/messages/web.ts b/packages/api/src/messages/web.ts new file mode 100644 index 0000000..7f97891 --- /dev/null +++ b/packages/api/src/messages/web.ts @@ -0,0 +1,16 @@ +// src/messages/web.ts +// Copyright (C) 2026 Robert Colbert +// Licensed under the Apache License, Version 2.0 + +import { IChatSession } from "../interfaces/chat-session.ts"; + +/* + * sessionUpdated + * Sent from gadget-code:web to gadget-code:ide when a chat session property + * has changed (e.g. auto-generated name). The IDE should merge these updates + * into its local state. + */ + +export type SessionUpdatedMessage = ( + updates: Partial, +) => void;