Merge pull request #4066 from Kilo-Org/track-sessions-from-the-extension

Author: iscekic

.changeset/lucky-chicken-doubt.md

@@ -0,0 +1,6 @@
+---
+"kilo-code": minor
+"@kilocode/cli": patch
+---
+
+use shared session manager from extension folder

cli/src/cli.ts

@@ -24,9 +24,9 @@ import type { CLIOptions } from "./types/cli.js"
 import type { CLIConfig, ProviderConfig } from "./config/types.js"
 import { getModelIdKey } from "./constants/providers/models.js"
 import type { ProviderName } from "./types/messages.js"
-import { TrpcClient } from "./services/trpcClient.js"
-import { SessionService } from "./services/session.js"
+import { KiloCodePathProvider, ExtensionMessengerAdapter } from "./services/session-adapters.js"
 import { getKiloToken } from "./config/persistence.js"
+import { SessionManager } from "../../src/shared/kilocode/cli-sessions/core/SessionManager.js"
 
 /**
  * Main application class that orchestrates the CLI lifecycle
@@ -37,7 +37,7 @@ export class CLI {
 	private ui: Instance | null = null
 	private options: CLIOptions
 	private isInitialized = false
-	private sessionService: SessionService | null = null
+	private sessionService: SessionManager | null = null
 
 	constructor(options: CLIOptions = {}) {
 		this.options = options
@@ -135,16 +135,31 @@ export class CLI {
 			const kiloToken = getKiloToken(config)
 
 			if (kiloToken) {
-				TrpcClient.init(kiloToken)
-				logs.debug("TrpcClient initialized with kiloToken", "CLI")
+				const pathProvider = new KiloCodePathProvider()
+				const extensionMessenger = new ExtensionMessengerAdapter(this.service)
+
+				this.sessionService = SessionManager.init({
+					pathProvider,
+					logger: logs,
+					extensionMessenger,
+					getToken: () => Promise.resolve(kiloToken),
+					onSessionCreated: (message) => {
+						if (this.options.json) {
+							console.log(JSON.stringify(message))
+						}
+					},
+					onSessionRestored: () => {
+						if (this.store) {
+							this.store.set(taskResumedViaContinueOrSessionAtom, true)
+						}
+					},
+					platform: "cli",
+				})
+				logs.debug("SessionManager initialized with dependencies", "CLI")
 
-				this.sessionService = SessionService.init(this.service, this.store, this.options.json)
-				logs.debug("SessionService initialized with ExtensionService", "CLI")
-
-				// Set workspace directory for git operations (important for parallel mode/worktrees)
 				const workspace = this.options.workspace || process.cwd()
 				this.sessionService.setWorkspaceDirectory(workspace)
-				logs.debug("SessionService workspace directory set", "CLI", { workspace })
+				logs.debug("SessionManager workspace directory set", "CLI", { workspace })
 
 				if (this.options.session) {
 					await this.sessionService.restoreSession(this.options.session)

cli/src/commands/__tests__/new.test.ts

@@ -6,11 +6,11 @@ import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"
 import { newCommand } from "../new.js"
 import type { CommandContext } from "../core/types.js"
 import { createMockContext } from "./helpers/mockContext.js"
-import { SessionService } from "../../services/session.js"
+import { SessionManager } from "../../../../src/shared/kilocode/cli-sessions/core/SessionManager.js"
 
 describe("/new command", () => {
 	let mockContext: CommandContext
-	let mockSessionService: Partial<SessionService> & { destroy: ReturnType<typeof vi.fn> }
+	let mockSessionManager: Partial<SessionManager> & { destroy: ReturnType<typeof vi.fn> }
 
 	beforeEach(() => {
 		// Mock process.stdout.write to capture terminal clearing
@@ -20,14 +20,14 @@ describe("/new command", () => {
 			input: "/new",
 		})
 
-		// Mock SessionService
-		mockSessionService = {
+		// Mock SessionManager
+		mockSessionManager = {
 			destroy: vi.fn().mockResolvedValue(undefined),
 			sessionId: "test-session-id",
 		}
 
-		// Mock SessionService.init to return our mock
-		vi.spyOn(SessionService, "init").mockReturnValue(mockSessionService as unknown as SessionService)
+		// Mock SessionManager.init to return our mock
+		vi.spyOn(SessionManager, "init").mockReturnValue(mockSessionManager as unknown as SessionManager)
 	})
 
 	afterEach(() => {
@@ -74,13 +74,13 @@ describe("/new command", () => {
 		it("should clear the session", async () => {
 			await newCommand.handler(mockContext)
 
-			expect(SessionService.init).toHaveBeenCalled()
-			expect(mockSessionService.destroy).toHaveBeenCalledTimes(1)
+			expect(SessionManager.init).toHaveBeenCalled()
+			expect(mockSessionManager.destroy).toHaveBeenCalledTimes(1)
 		})
 
 		it("should continue execution even if session clearing fails", async () => {
 			const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {})
-			mockSessionService.destroy.mockRejectedValue(new Error("Session error"))
+			mockSessionManager.destroy.mockRejectedValue(new Error("Session error"))
 
 			await newCommand.handler(mockContext)
 
@@ -121,7 +121,7 @@ describe("/new command", () => {
 				callOrder.push("clearTask")
 			})
 
-			mockSessionService.destroy = vi.fn().mockImplementation(async () => {
+			mockSessionManager.destroy = vi.fn().mockImplementation(async () => {
 				callOrder.push("sessionDestroy")
 			})
 
@@ -164,7 +164,7 @@ describe("/new command", () => {
 
 			// Verify all cleanup operations were performed
 			expect(mockContext.clearTask).toHaveBeenCalled()
-			expect(mockSessionService.destroy).toHaveBeenCalled()
+			expect(mockSessionManager.destroy).toHaveBeenCalled()
 			expect(mockContext.replaceMessages).toHaveBeenCalled()
 
 			// Verify welcome message was replaced