mirror of
https://github.com/samiyev/puaros.git
synced 2025-12-28 07:16:53 +05:00
feat(ipuaro): add session management (v0.10.0)
- Add ISessionStorage interface and RedisSessionStorage implementation - Add ContextManager for token budget and compression - Add StartSession, HandleMessage, UndoChange use cases - Update CHANGELOG and TODO documentation - 88 new tests (1174 total), 97.73% coverage
This commit is contained in:
imfozilbek
@@ -0,0 +1,248 @@
|
||||
import { describe, it, expect, vi, beforeEach } from "vitest"
|
||||
import { ContextManager } from "../../../../src/application/use-cases/ContextManager.js"
|
||||
import { Session } from "../../../../src/domain/entities/Session.js"
|
||||
import type { ILLMClient, LLMResponse } from "../../../../src/domain/services/ILLMClient.js"
|
||||
import {
|
||||
createUserMessage,
|
||||
createAssistantMessage,
|
||||
} from "../../../../src/domain/value-objects/ChatMessage.js"
|
||||
|
||||
describe("ContextManager", () => {
|
||||
let manager: ContextManager
|
||||
const CONTEXT_SIZE = 128_000
|
||||
|
||||
beforeEach(() => {
|
||||
manager = new ContextManager(CONTEXT_SIZE)
|
||||
})
|
||||
|
||||
describe("addToContext", () => {
|
||||
it("should add file to context", () => {
|
||||
manager.addToContext("test.ts", 100)
|
||||
|
||||
expect(manager.getFilesInContext()).toContain("test.ts")
|
||||
expect(manager.getTokenCount()).toBe(100)
|
||||
})
|
||||
|
||||
it("should update token count when same file added", () => {
|
||||
manager.addToContext("test.ts", 100)
|
||||
manager.addToContext("test.ts", 200)
|
||||
|
||||
expect(manager.getFilesInContext()).toHaveLength(1)
|
||||
expect(manager.getTokenCount()).toBe(200)
|
||||
})
|
||||
|
||||
it("should accumulate tokens for different files", () => {
|
||||
manager.addToContext("a.ts", 100)
|
||||
manager.addToContext("b.ts", 200)
|
||||
|
||||
expect(manager.getFilesInContext()).toHaveLength(2)
|
||||
expect(manager.getTokenCount()).toBe(300)
|
||||
})
|
||||
})
|
||||
|
||||
describe("removeFromContext", () => {
|
||||
it("should remove file from context", () => {
|
||||
manager.addToContext("test.ts", 100)
|
||||
manager.removeFromContext("test.ts")
|
||||
|
||||
expect(manager.getFilesInContext()).not.toContain("test.ts")
|
||||
expect(manager.getTokenCount()).toBe(0)
|
||||
})
|
||||
|
||||
it("should handle removing non-existent file", () => {
|
||||
manager.removeFromContext("non-existent.ts")
|
||||
|
||||
expect(manager.getTokenCount()).toBe(0)
|
||||
})
|
||||
})
|
||||
|
||||
describe("getUsage", () => {
|
||||
it("should return 0 for empty context", () => {
|
||||
expect(manager.getUsage()).toBe(0)
|
||||
})
|
||||
|
||||
it("should calculate usage ratio correctly", () => {
|
||||
manager.addToContext("test.ts", CONTEXT_SIZE / 2)
|
||||
|
||||
expect(manager.getUsage()).toBe(0.5)
|
||||
})
|
||||
})
|
||||
|
||||
describe("getAvailableTokens", () => {
|
||||
it("should return full context when empty", () => {
|
||||
expect(manager.getAvailableTokens()).toBe(CONTEXT_SIZE)
|
||||
})
|
||||
|
||||
it("should calculate available tokens correctly", () => {
|
||||
manager.addToContext("test.ts", 1000)
|
||||
|
||||
expect(manager.getAvailableTokens()).toBe(CONTEXT_SIZE - 1000)
|
||||
})
|
||||
})
|
||||
|
||||
describe("needsCompression", () => {
|
||||
it("should return false when under threshold", () => {
|
||||
manager.addToContext("test.ts", CONTEXT_SIZE * 0.5)
|
||||
|
||||
expect(manager.needsCompression()).toBe(false)
|
||||
})
|
||||
|
||||
it("should return true when over threshold", () => {
|
||||
manager.addToContext("test.ts", CONTEXT_SIZE * 0.85)
|
||||
|
||||
expect(manager.needsCompression()).toBe(true)
|
||||
})
|
||||
|
||||
it("should return false at exactly threshold", () => {
|
||||
manager.addToContext("test.ts", CONTEXT_SIZE * 0.8)
|
||||
|
||||
expect(manager.needsCompression()).toBe(false)
|
||||
})
|
||||
})
|
||||
|
||||
describe("addTokens", () => {
|
||||
it("should add tokens to current count", () => {
|
||||
manager.addTokens(500)
|
||||
|
||||
expect(manager.getTokenCount()).toBe(500)
|
||||
})
|
||||
|
||||
it("should accumulate tokens", () => {
|
||||
manager.addTokens(100)
|
||||
manager.addTokens(200)
|
||||
|
||||
expect(manager.getTokenCount()).toBe(300)
|
||||
})
|
||||
})
|
||||
|
||||
describe("syncFromSession", () => {
|
||||
it("should sync files from session context", () => {
|
||||
const session = new Session("test", "project")
|
||||
session.context.filesInContext = ["a.ts", "b.ts"]
|
||||
session.context.tokenUsage = 0.5
|
||||
|
||||
manager.syncFromSession(session)
|
||||
|
||||
expect(manager.getFilesInContext()).toContain("a.ts")
|
||||
expect(manager.getFilesInContext()).toContain("b.ts")
|
||||
expect(manager.getTokenCount()).toBe(Math.floor(0.5 * CONTEXT_SIZE))
|
||||
})
|
||||
|
||||
it("should clear previous state on sync", () => {
|
||||
manager.addToContext("old.ts", 1000)
|
||||
|
||||
const session = new Session("test", "project")
|
||||
session.context.filesInContext = ["new.ts"]
|
||||
session.context.tokenUsage = 0.1
|
||||
|
||||
manager.syncFromSession(session)
|
||||
|
||||
expect(manager.getFilesInContext()).not.toContain("old.ts")
|
||||
expect(manager.getFilesInContext()).toContain("new.ts")
|
||||
})
|
||||
})
|
||||
|
||||
describe("updateSession", () => {
|
||||
it("should update session with current context state", () => {
|
||||
const session = new Session("test", "project")
|
||||
|
||||
manager.addToContext("test.ts", 1000)
|
||||
manager.updateSession(session)
|
||||
|
||||
expect(session.context.filesInContext).toContain("test.ts")
|
||||
expect(session.context.tokenUsage).toBeCloseTo(1000 / CONTEXT_SIZE)
|
||||
})
|
||||
|
||||
it("should set needsCompression flag", () => {
|
||||
const session = new Session("test", "project")
|
||||
|
||||
manager.addToContext("large.ts", CONTEXT_SIZE * 0.9)
|
||||
manager.updateSession(session)
|
||||
|
||||
expect(session.context.needsCompression).toBe(true)
|
||||
})
|
||||
})
|
||||
|
||||
describe("compress", () => {
|
||||
let mockLLM: ILLMClient
|
||||
let session: Session
|
||||
|
||||
beforeEach(() => {
|
||||
mockLLM = {
|
||||
chat: vi.fn().mockResolvedValue({
|
||||
content: "Summary of previous conversation",
|
||||
toolCalls: [],
|
||||
tokens: 50,
|
||||
timeMs: 100,
|
||||
truncated: false,
|
||||
stopReason: "end",
|
||||
} as LLMResponse),
|
||||
countTokens: vi.fn().mockResolvedValue(10),
|
||||
isAvailable: vi.fn().mockResolvedValue(true),
|
||||
getModelName: vi.fn().mockReturnValue("test-model"),
|
||||
getContextWindowSize: vi.fn().mockReturnValue(CONTEXT_SIZE),
|
||||
pullModel: vi.fn().mockResolvedValue(undefined),
|
||||
abort: vi.fn(),
|
||||
}
|
||||
|
||||
session = new Session("test", "project")
|
||||
})
|
||||
|
||||
it("should not compress when history is short", async () => {
|
||||
for (let i = 0; i < 5; i++) {
|
||||
session.addMessage(createUserMessage(`Message ${String(i)}`))
|
||||
}
|
||||
|
||||
const result = await manager.compress(session, mockLLM)
|
||||
|
||||
expect(result.compressed).toBe(false)
|
||||
expect(result.removedMessages).toBe(0)
|
||||
})
|
||||
|
||||
it("should compress when history is long enough", async () => {
|
||||
for (let i = 0; i < 15; i++) {
|
||||
session.addMessage(createUserMessage(`Message ${String(i)}`))
|
||||
session.addMessage(createAssistantMessage(`Response ${String(i)}`))
|
||||
}
|
||||
manager.addToContext("test.ts", 10000)
|
||||
|
||||
const result = await manager.compress(session, mockLLM)
|
||||
|
||||
expect(result.compressed).toBe(true)
|
||||
expect(result.removedMessages).toBeGreaterThan(0)
|
||||
expect(result.summary).toBeDefined()
|
||||
})
|
||||
|
||||
it("should keep recent messages after compression", async () => {
|
||||
for (let i = 0; i < 15; i++) {
|
||||
session.addMessage(createUserMessage(`Message ${String(i)}`))
|
||||
}
|
||||
|
||||
await manager.compress(session, mockLLM)
|
||||
|
||||
expect(session.history.length).toBeLessThan(15)
|
||||
expect(session.history[session.history.length - 1].content).toContain("Message 14")
|
||||
})
|
||||
|
||||
it("should add summary as system message", async () => {
|
||||
for (let i = 0; i < 15; i++) {
|
||||
session.addMessage(createUserMessage(`Message ${String(i)}`))
|
||||
}
|
||||
|
||||
await manager.compress(session, mockLLM)
|
||||
|
||||
expect(session.history[0].role).toBe("system")
|
||||
expect(session.history[0].content).toContain("Summary")
|
||||
})
|
||||
})
|
||||
|
||||
describe("createInitialState", () => {
|
||||
it("should create empty initial state", () => {
|
||||
const state = ContextManager.createInitialState()
|
||||
|
||||
expect(state.filesInContext).toEqual([])
|
||||
expect(state.tokenUsage).toBe(0)
|
||||
expect(state.needsCompression).toBe(false)
|
||||
})
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,421 @@
|
||||
import { describe, it, expect, vi, beforeEach } from "vitest"
|
||||
import { HandleMessage } from "../../../../src/application/use-cases/HandleMessage.js"
|
||||
import type { IStorage } from "../../../../src/domain/services/IStorage.js"
|
||||
import type { ISessionStorage } from "../../../../src/domain/services/ISessionStorage.js"
|
||||
import type { ILLMClient, LLMResponse } from "../../../../src/domain/services/ILLMClient.js"
|
||||
import type { IToolRegistry } from "../../../../src/application/interfaces/IToolRegistry.js"
|
||||
import type { ITool, ToolContext } from "../../../../src/domain/services/ITool.js"
|
||||
import { Session } from "../../../../src/domain/entities/Session.js"
|
||||
import { createSuccessResult } from "../../../../src/domain/value-objects/ToolResult.js"
|
||||
|
||||
describe("HandleMessage", () => {
|
||||
let useCase: HandleMessage
|
||||
let mockStorage: IStorage
|
||||
let mockSessionStorage: ISessionStorage
|
||||
let mockLLM: ILLMClient
|
||||
let mockTools: IToolRegistry
|
||||
let session: Session
|
||||
|
||||
const createMockLLMResponse = (content: string, toolCalls = false): LLMResponse => ({
|
||||
content,
|
||||
toolCalls: [],
|
||||
tokens: 100,
|
||||
timeMs: 50,
|
||||
truncated: false,
|
||||
stopReason: toolCalls ? "tool_use" : "end",
|
||||
})
|
||||
|
||||
beforeEach(() => {
|
||||
mockStorage = {
|
||||
getFile: vi.fn().mockResolvedValue(null),
|
||||
setFile: vi.fn().mockResolvedValue(undefined),
|
||||
deleteFile: vi.fn().mockResolvedValue(undefined),
|
||||
getAllFiles: vi.fn().mockResolvedValue(new Map()),
|
||||
getFileCount: vi.fn().mockResolvedValue(0),
|
||||
getAST: vi.fn().mockResolvedValue(null),
|
||||
setAST: vi.fn().mockResolvedValue(undefined),
|
||||
deleteAST: vi.fn().mockResolvedValue(undefined),
|
||||
getAllASTs: vi.fn().mockResolvedValue(new Map()),
|
||||
getMeta: vi.fn().mockResolvedValue(null),
|
||||
setMeta: vi.fn().mockResolvedValue(undefined),
|
||||
deleteMeta: vi.fn().mockResolvedValue(undefined),
|
||||
getAllMetas: vi.fn().mockResolvedValue(new Map()),
|
||||
getSymbolIndex: vi.fn().mockResolvedValue(new Map()),
|
||||
setSymbolIndex: vi.fn().mockResolvedValue(undefined),
|
||||
getDepsGraph: vi.fn().mockResolvedValue({ imports: new Map(), importedBy: new Map() }),
|
||||
setDepsGraph: vi.fn().mockResolvedValue(undefined),
|
||||
getProjectConfig: vi.fn().mockResolvedValue(null),
|
||||
setProjectConfig: vi.fn().mockResolvedValue(undefined),
|
||||
connect: vi.fn().mockResolvedValue(undefined),
|
||||
disconnect: vi.fn().mockResolvedValue(undefined),
|
||||
isConnected: vi.fn().mockReturnValue(true),
|
||||
clear: vi.fn().mockResolvedValue(undefined),
|
||||
}
|
||||
|
||||
mockSessionStorage = {
|
||||
saveSession: vi.fn().mockResolvedValue(undefined),
|
||||
loadSession: vi.fn().mockResolvedValue(null),
|
||||
deleteSession: vi.fn().mockResolvedValue(undefined),
|
||||
listSessions: vi.fn().mockResolvedValue([]),
|
||||
getLatestSession: vi.fn().mockResolvedValue(null),
|
||||
sessionExists: vi.fn().mockResolvedValue(false),
|
||||
pushUndoEntry: vi.fn().mockResolvedValue(undefined),
|
||||
popUndoEntry: vi.fn().mockResolvedValue(null),
|
||||
getUndoStack: vi.fn().mockResolvedValue([]),
|
||||
touchSession: vi.fn().mockResolvedValue(undefined),
|
||||
clearAllSessions: vi.fn().mockResolvedValue(undefined),
|
||||
}
|
||||
|
||||
mockLLM = {
|
||||
chat: vi.fn().mockResolvedValue(createMockLLMResponse("Hello!")),
|
||||
countTokens: vi.fn().mockResolvedValue(10),
|
||||
isAvailable: vi.fn().mockResolvedValue(true),
|
||||
getModelName: vi.fn().mockReturnValue("test-model"),
|
||||
getContextWindowSize: vi.fn().mockReturnValue(128_000),
|
||||
pullModel: vi.fn().mockResolvedValue(undefined),
|
||||
abort: vi.fn(),
|
||||
}
|
||||
|
||||
mockTools = {
|
||||
register: vi.fn(),
|
||||
get: vi.fn().mockReturnValue(undefined),
|
||||
getAll: vi.fn().mockReturnValue([]),
|
||||
getByCategory: vi.fn().mockReturnValue([]),
|
||||
has: vi.fn().mockReturnValue(false),
|
||||
execute: vi.fn(),
|
||||
getToolDefinitions: vi.fn().mockReturnValue([]),
|
||||
}
|
||||
|
||||
session = new Session("test-session", "test-project")
|
||||
useCase = new HandleMessage(mockStorage, mockSessionStorage, mockLLM, mockTools, "/project")
|
||||
})
|
||||
|
||||
describe("execute", () => {
|
||||
it("should add user message to session history", async () => {
|
||||
await useCase.execute(session, "Hello, assistant!")
|
||||
|
||||
expect(session.history.length).toBeGreaterThan(0)
|
||||
expect(session.history[0].role).toBe("user")
|
||||
expect(session.history[0].content).toBe("Hello, assistant!")
|
||||
})
|
||||
|
||||
it("should add user input to input history", async () => {
|
||||
await useCase.execute(session, "Test command")
|
||||
|
||||
expect(session.inputHistory).toContain("Test command")
|
||||
})
|
||||
|
||||
it("should save session after user message", async () => {
|
||||
await useCase.execute(session, "Hello")
|
||||
|
||||
expect(mockSessionStorage.saveSession).toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it("should send messages to LLM", async () => {
|
||||
await useCase.execute(session, "What is 2+2?")
|
||||
|
||||
expect(mockLLM.chat).toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it("should add assistant response to history", async () => {
|
||||
vi.mocked(mockLLM.chat).mockResolvedValue(createMockLLMResponse("The answer is 4!"))
|
||||
|
||||
await useCase.execute(session, "What is 2+2?")
|
||||
|
||||
const assistantMessages = session.history.filter((m) => m.role === "assistant")
|
||||
expect(assistantMessages.length).toBeGreaterThan(0)
|
||||
expect(assistantMessages[0].content).toBe("The answer is 4!")
|
||||
})
|
||||
|
||||
it("should not add empty user messages", async () => {
|
||||
await useCase.execute(session, " ")
|
||||
|
||||
const userMessages = session.history.filter((m) => m.role === "user")
|
||||
expect(userMessages.length).toBe(0)
|
||||
})
|
||||
|
||||
it("should track token usage in message stats", async () => {
|
||||
vi.mocked(mockLLM.chat).mockResolvedValue({
|
||||
content: "Response",
|
||||
toolCalls: [],
|
||||
tokens: 150,
|
||||
timeMs: 200,
|
||||
truncated: false,
|
||||
stopReason: "end",
|
||||
})
|
||||
|
||||
await useCase.execute(session, "Hello")
|
||||
|
||||
const assistantMessage = session.history.find((m) => m.role === "assistant")
|
||||
expect(assistantMessage?.stats?.tokens).toBe(150)
|
||||
expect(assistantMessage?.stats?.timeMs).toBeGreaterThanOrEqual(0)
|
||||
})
|
||||
})
|
||||
|
||||
describe("tool execution", () => {
|
||||
const mockTool: ITool = {
|
||||
name: "get_lines",
|
||||
description: "Get lines from file",
|
||||
parameters: [],
|
||||
requiresConfirmation: false,
|
||||
category: "read",
|
||||
validateParams: vi.fn().mockReturnValue(null),
|
||||
execute: vi.fn().mockResolvedValue(createSuccessResult("test", { lines: [] }, 10)),
|
||||
}
|
||||
|
||||
beforeEach(() => {
|
||||
vi.mocked(mockTools.get).mockReturnValue(mockTool)
|
||||
})
|
||||
|
||||
it("should execute tools when LLM returns tool calls", async () => {
|
||||
vi.mocked(mockLLM.chat)
|
||||
.mockResolvedValueOnce(
|
||||
createMockLLMResponse(
|
||||
'<tool_call name="get_lines"><path>test.ts</path></tool_call>',
|
||||
true,
|
||||
),
|
||||
)
|
||||
.mockResolvedValueOnce(createMockLLMResponse("Done!"))
|
||||
|
||||
await useCase.execute(session, "Show me test.ts")
|
||||
|
||||
expect(mockTool.execute).toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it("should add tool results to session", async () => {
|
||||
vi.mocked(mockLLM.chat)
|
||||
.mockResolvedValueOnce(
|
||||
createMockLLMResponse(
|
||||
'<tool_call name="get_lines"><path>test.ts</path></tool_call>',
|
||||
true,
|
||||
),
|
||||
)
|
||||
.mockResolvedValueOnce(createMockLLMResponse("Done!"))
|
||||
|
||||
await useCase.execute(session, "Show me test.ts")
|
||||
|
||||
const toolMessages = session.history.filter((m) => m.role === "tool")
|
||||
expect(toolMessages.length).toBeGreaterThan(0)
|
||||
})
|
||||
|
||||
it("should return error for unknown tools", async () => {
|
||||
vi.mocked(mockTools.get).mockReturnValue(undefined)
|
||||
vi.mocked(mockLLM.chat)
|
||||
.mockResolvedValueOnce(
|
||||
createMockLLMResponse(
|
||||
'<tool_call name="unknown_tool"><param>value</param></tool_call>',
|
||||
true,
|
||||
),
|
||||
)
|
||||
.mockResolvedValueOnce(createMockLLMResponse("Sorry, that didn't work"))
|
||||
|
||||
await useCase.execute(session, "Do something")
|
||||
|
||||
const toolMessages = session.history.filter((m) => m.role === "tool")
|
||||
expect(toolMessages[0].content).toContain("Unknown tool")
|
||||
})
|
||||
|
||||
it("should stop after max tool calls exceeded", async () => {
|
||||
useCase.setOptions({ maxToolCalls: 2 })
|
||||
|
||||
vi.mocked(mockLLM.chat).mockResolvedValue(
|
||||
createMockLLMResponse(
|
||||
'<tool_call name="get_lines"><path>a.ts</path></tool_call>' +
|
||||
'<tool_call name="get_lines"><path>b.ts</path></tool_call>' +
|
||||
'<tool_call name="get_lines"><path>c.ts</path></tool_call>',
|
||||
true,
|
||||
),
|
||||
)
|
||||
|
||||
await useCase.execute(session, "Show many files")
|
||||
|
||||
const systemMessages = session.history.filter((m) => m.role === "system")
|
||||
const maxExceeded = systemMessages.some((m) => m.content.includes("Maximum tool calls"))
|
||||
expect(maxExceeded).toBe(true)
|
||||
})
|
||||
})
|
||||
|
||||
describe("events", () => {
|
||||
it("should emit message events", async () => {
|
||||
const onMessage = vi.fn()
|
||||
useCase.setEvents({ onMessage })
|
||||
|
||||
await useCase.execute(session, "Hello")
|
||||
|
||||
expect(onMessage).toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it("should emit status changes", async () => {
|
||||
const onStatusChange = vi.fn()
|
||||
useCase.setEvents({ onStatusChange })
|
||||
|
||||
await useCase.execute(session, "Hello")
|
||||
|
||||
expect(onStatusChange).toHaveBeenCalledWith("thinking")
|
||||
expect(onStatusChange).toHaveBeenCalledWith("ready")
|
||||
})
|
||||
|
||||
it("should emit tool call events", async () => {
|
||||
const onToolCall = vi.fn()
|
||||
useCase.setEvents({ onToolCall })
|
||||
|
||||
const mockTool: ITool = {
|
||||
name: "get_lines",
|
||||
description: "Test",
|
||||
parameters: [],
|
||||
requiresConfirmation: false,
|
||||
category: "read",
|
||||
validateParams: vi.fn().mockReturnValue(null),
|
||||
execute: vi.fn().mockResolvedValue(createSuccessResult("test", {}, 10)),
|
||||
}
|
||||
vi.mocked(mockTools.get).mockReturnValue(mockTool)
|
||||
|
||||
vi.mocked(mockLLM.chat)
|
||||
.mockResolvedValueOnce(
|
||||
createMockLLMResponse(
|
||||
'<tool_call name="get_lines"><path>test.ts</path></tool_call>',
|
||||
),
|
||||
)
|
||||
.mockResolvedValueOnce(createMockLLMResponse("Done"))
|
||||
|
||||
await useCase.execute(session, "Show file")
|
||||
|
||||
expect(onToolCall).toHaveBeenCalled()
|
||||
})
|
||||
})
|
||||
|
||||
describe("confirmation handling", () => {
|
||||
const mockEditTool: ITool = {
|
||||
name: "edit_lines",
|
||||
description: "Edit lines",
|
||||
parameters: [],
|
||||
requiresConfirmation: true,
|
||||
category: "edit",
|
||||
validateParams: vi.fn().mockReturnValue(null),
|
||||
execute: vi
|
||||
.fn()
|
||||
.mockImplementation(async (_params: Record<string, unknown>, ctx: ToolContext) => {
|
||||
const confirmed = await ctx.requestConfirmation("Apply edit?", {
|
||||
filePath: "test.ts",
|
||||
oldLines: ["old"],
|
||||
newLines: ["new"],
|
||||
startLine: 1,
|
||||
})
|
||||
if (!confirmed) {
|
||||
return createSuccessResult("test", { cancelled: true }, 10)
|
||||
}
|
||||
return createSuccessResult("test", { applied: true }, 10)
|
||||
}),
|
||||
}
|
||||
|
||||
beforeEach(() => {
|
||||
vi.mocked(mockTools.get).mockReturnValue(mockEditTool)
|
||||
})
|
||||
|
||||
it("should auto-apply when autoApply option is true", async () => {
|
||||
useCase.setOptions({ autoApply: true })
|
||||
|
||||
vi.mocked(mockLLM.chat)
|
||||
.mockResolvedValueOnce(
|
||||
createMockLLMResponse(
|
||||
'<tool_call name="edit_lines"><path>test.ts</path></tool_call>',
|
||||
),
|
||||
)
|
||||
.mockResolvedValueOnce(createMockLLMResponse("Done"))
|
||||
|
||||
await useCase.execute(session, "Edit file")
|
||||
|
||||
expect(mockEditTool.execute).toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it("should ask for confirmation via callback", async () => {
|
||||
const onConfirmation = vi.fn().mockResolvedValue(true)
|
||||
useCase.setEvents({ onConfirmation })
|
||||
|
||||
vi.mocked(mockLLM.chat)
|
||||
.mockResolvedValueOnce(
|
||||
createMockLLMResponse(
|
||||
'<tool_call name="edit_lines"><path>test.ts</path></tool_call>',
|
||||
),
|
||||
)
|
||||
.mockResolvedValueOnce(createMockLLMResponse("Done"))
|
||||
|
||||
await useCase.execute(session, "Edit file")
|
||||
|
||||
expect(onConfirmation).toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it("should create undo entry on confirmation", async () => {
|
||||
const onUndoEntry = vi.fn()
|
||||
useCase.setEvents({
|
||||
onConfirmation: vi.fn().mockResolvedValue(true),
|
||||
onUndoEntry,
|
||||
})
|
||||
|
||||
vi.mocked(mockLLM.chat)
|
||||
.mockResolvedValueOnce(
|
||||
createMockLLMResponse(
|
||||
'<tool_call name="edit_lines"><path>test.ts</path></tool_call>',
|
||||
),
|
||||
)
|
||||
.mockResolvedValueOnce(createMockLLMResponse("Done"))
|
||||
|
||||
await useCase.execute(session, "Edit file")
|
||||
|
||||
expect(onUndoEntry).toHaveBeenCalled()
|
||||
expect(mockSessionStorage.pushUndoEntry).toHaveBeenCalled()
|
||||
})
|
||||
})
|
||||
|
||||
describe("abort", () => {
|
||||
it("should stop processing when aborted", async () => {
|
||||
vi.mocked(mockLLM.chat).mockImplementation(async () => {
|
||||
await new Promise((resolve) => setTimeout(resolve, 100))
|
||||
return createMockLLMResponse("Response")
|
||||
})
|
||||
|
||||
const promise = useCase.execute(session, "Hello")
|
||||
|
||||
setTimeout(() => useCase.abort(), 10)
|
||||
|
||||
await promise
|
||||
|
||||
expect(mockLLM.abort).toHaveBeenCalled()
|
||||
})
|
||||
})
|
||||
|
||||
describe("error handling", () => {
|
||||
it("should handle LLM errors gracefully", async () => {
|
||||
vi.mocked(mockLLM.chat).mockRejectedValue(new Error("LLM unavailable"))
|
||||
|
||||
await useCase.execute(session, "Hello")
|
||||
|
||||
const systemMessages = session.history.filter((m) => m.role === "system")
|
||||
expect(systemMessages.some((m) => m.content.includes("Error"))).toBe(true)
|
||||
})
|
||||
|
||||
it("should emit error status on LLM failure", async () => {
|
||||
const onStatusChange = vi.fn()
|
||||
useCase.setEvents({ onStatusChange })
|
||||
|
||||
vi.mocked(mockLLM.chat).mockRejectedValue(new Error("LLM error"))
|
||||
|
||||
await useCase.execute(session, "Hello")
|
||||
|
||||
expect(onStatusChange).toHaveBeenCalledWith("error")
|
||||
})
|
||||
|
||||
it("should allow retry on error", async () => {
|
||||
const onError = vi.fn().mockResolvedValue("retry")
|
||||
useCase.setEvents({ onError })
|
||||
|
||||
vi.mocked(mockLLM.chat)
|
||||
.mockRejectedValueOnce(new Error("Temporary error"))
|
||||
.mockResolvedValueOnce(createMockLLMResponse("Success!"))
|
||||
|
||||
await useCase.execute(session, "Hello")
|
||||
|
||||
expect(onError).toHaveBeenCalled()
|
||||
})
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,112 @@
|
||||
import { describe, it, expect, vi, beforeEach } from "vitest"
|
||||
import { StartSession } from "../../../../src/application/use-cases/StartSession.js"
|
||||
import type { ISessionStorage } from "../../../../src/domain/services/ISessionStorage.js"
|
||||
import { Session } from "../../../../src/domain/entities/Session.js"
|
||||
|
||||
describe("StartSession", () => {
|
||||
let useCase: StartSession
|
||||
let mockSessionStorage: ISessionStorage
|
||||
|
||||
beforeEach(() => {
|
||||
mockSessionStorage = {
|
||||
saveSession: vi.fn().mockResolvedValue(undefined),
|
||||
loadSession: vi.fn().mockResolvedValue(null),
|
||||
deleteSession: vi.fn().mockResolvedValue(undefined),
|
||||
listSessions: vi.fn().mockResolvedValue([]),
|
||||
getLatestSession: vi.fn().mockResolvedValue(null),
|
||||
sessionExists: vi.fn().mockResolvedValue(false),
|
||||
pushUndoEntry: vi.fn().mockResolvedValue(undefined),
|
||||
popUndoEntry: vi.fn().mockResolvedValue(null),
|
||||
getUndoStack: vi.fn().mockResolvedValue([]),
|
||||
touchSession: vi.fn().mockResolvedValue(undefined),
|
||||
clearAllSessions: vi.fn().mockResolvedValue(undefined),
|
||||
}
|
||||
|
||||
useCase = new StartSession(mockSessionStorage)
|
||||
})
|
||||
|
||||
describe("execute", () => {
|
||||
it("should create new session when no existing session", async () => {
|
||||
const result = await useCase.execute("test-project")
|
||||
|
||||
expect(result.isNew).toBe(true)
|
||||
expect(result.session.projectName).toBe("test-project")
|
||||
expect(mockSessionStorage.saveSession).toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it("should return latest session when one exists", async () => {
|
||||
const existingSession = new Session("existing-id", "test-project")
|
||||
vi.mocked(mockSessionStorage.getLatestSession).mockResolvedValue(existingSession)
|
||||
|
||||
const result = await useCase.execute("test-project")
|
||||
|
||||
expect(result.isNew).toBe(false)
|
||||
expect(result.session.id).toBe("existing-id")
|
||||
expect(mockSessionStorage.touchSession).toHaveBeenCalledWith("existing-id")
|
||||
})
|
||||
|
||||
it("should load specific session by ID", async () => {
|
||||
const specificSession = new Session("specific-id", "test-project")
|
||||
vi.mocked(mockSessionStorage.loadSession).mockResolvedValue(specificSession)
|
||||
|
||||
const result = await useCase.execute("test-project", { sessionId: "specific-id" })
|
||||
|
||||
expect(result.isNew).toBe(false)
|
||||
expect(result.session.id).toBe("specific-id")
|
||||
expect(mockSessionStorage.loadSession).toHaveBeenCalledWith("specific-id")
|
||||
})
|
||||
|
||||
it("should create new session when specified session not found", async () => {
|
||||
vi.mocked(mockSessionStorage.loadSession).mockResolvedValue(null)
|
||||
|
||||
const result = await useCase.execute("test-project", { sessionId: "non-existent" })
|
||||
|
||||
expect(result.isNew).toBe(true)
|
||||
expect(mockSessionStorage.saveSession).toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it("should force new session when forceNew is true", async () => {
|
||||
const existingSession = new Session("existing-id", "test-project")
|
||||
vi.mocked(mockSessionStorage.getLatestSession).mockResolvedValue(existingSession)
|
||||
|
||||
const result = await useCase.execute("test-project", { forceNew: true })
|
||||
|
||||
expect(result.isNew).toBe(true)
|
||||
expect(result.session.id).not.toBe("existing-id")
|
||||
expect(mockSessionStorage.saveSession).toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it("should generate unique session IDs", async () => {
|
||||
const result1 = await useCase.execute("test-project", { forceNew: true })
|
||||
const result2 = await useCase.execute("test-project", { forceNew: true })
|
||||
|
||||
expect(result1.session.id).not.toBe(result2.session.id)
|
||||
})
|
||||
|
||||
it("should set correct project name on new session", async () => {
|
||||
const result = await useCase.execute("my-special-project")
|
||||
|
||||
expect(result.session.projectName).toBe("my-special-project")
|
||||
})
|
||||
|
||||
it("should initialize new session with empty history", async () => {
|
||||
const result = await useCase.execute("test-project")
|
||||
|
||||
expect(result.session.history).toEqual([])
|
||||
})
|
||||
|
||||
it("should initialize new session with empty undo stack", async () => {
|
||||
const result = await useCase.execute("test-project")
|
||||
|
||||
expect(result.session.undoStack).toEqual([])
|
||||
})
|
||||
|
||||
it("should initialize new session with zero stats", async () => {
|
||||
const result = await useCase.execute("test-project")
|
||||
|
||||
expect(result.session.stats.totalTokens).toBe(0)
|
||||
expect(result.session.stats.toolCalls).toBe(0)
|
||||
expect(result.session.stats.editsApplied).toBe(0)
|
||||
})
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,234 @@
|
||||
import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"
|
||||
import { promises as fs } from "node:fs"
|
||||
import { UndoChange } from "../../../../src/application/use-cases/UndoChange.js"
|
||||
import type { ISessionStorage } from "../../../../src/domain/services/ISessionStorage.js"
|
||||
import type { IStorage } from "../../../../src/domain/services/IStorage.js"
|
||||
import { Session } from "../../../../src/domain/entities/Session.js"
|
||||
import type { UndoEntry } from "../../../../src/domain/value-objects/UndoEntry.js"
|
||||
|
||||
vi.mock("node:fs", () => ({
|
||||
promises: {
|
||||
readFile: vi.fn(),
|
||||
writeFile: vi.fn(),
|
||||
stat: vi.fn(),
|
||||
},
|
||||
}))
|
||||
|
||||
describe("UndoChange", () => {
|
||||
let useCase: UndoChange
|
||||
let mockSessionStorage: ISessionStorage
|
||||
let mockStorage: IStorage
|
||||
let session: Session
|
||||
|
||||
const createUndoEntry = (overrides: Partial<UndoEntry> = {}): UndoEntry => ({
|
||||
id: "undo-1",
|
||||
timestamp: Date.now(),
|
||||
filePath: "/project/test.ts",
|
||||
previousContent: ["const a = 1"],
|
||||
newContent: ["const a = 2"],
|
||||
description: "Edit test.ts",
|
||||
...overrides,
|
||||
})
|
||||
|
||||
beforeEach(() => {
|
||||
mockSessionStorage = {
|
||||
saveSession: vi.fn().mockResolvedValue(undefined),
|
||||
loadSession: vi.fn().mockResolvedValue(null),
|
||||
deleteSession: vi.fn().mockResolvedValue(undefined),
|
||||
listSessions: vi.fn().mockResolvedValue([]),
|
||||
getLatestSession: vi.fn().mockResolvedValue(null),
|
||||
sessionExists: vi.fn().mockResolvedValue(false),
|
||||
pushUndoEntry: vi.fn().mockResolvedValue(undefined),
|
||||
popUndoEntry: vi.fn().mockResolvedValue(null),
|
||||
getUndoStack: vi.fn().mockResolvedValue([]),
|
||||
touchSession: vi.fn().mockResolvedValue(undefined),
|
||||
clearAllSessions: vi.fn().mockResolvedValue(undefined),
|
||||
}
|
||||
|
||||
mockStorage = {
|
||||
getFile: vi.fn().mockResolvedValue(null),
|
||||
setFile: vi.fn().mockResolvedValue(undefined),
|
||||
deleteFile: vi.fn().mockResolvedValue(undefined),
|
||||
getAllFiles: vi.fn().mockResolvedValue(new Map()),
|
||||
getFileCount: vi.fn().mockResolvedValue(0),
|
||||
getAST: vi.fn().mockResolvedValue(null),
|
||||
setAST: vi.fn().mockResolvedValue(undefined),
|
||||
deleteAST: vi.fn().mockResolvedValue(undefined),
|
||||
getAllASTs: vi.fn().mockResolvedValue(new Map()),
|
||||
getMeta: vi.fn().mockResolvedValue(null),
|
||||
setMeta: vi.fn().mockResolvedValue(undefined),
|
||||
deleteMeta: vi.fn().mockResolvedValue(undefined),
|
||||
getAllMetas: vi.fn().mockResolvedValue(new Map()),
|
||||
getSymbolIndex: vi.fn().mockResolvedValue(new Map()),
|
||||
setSymbolIndex: vi.fn().mockResolvedValue(undefined),
|
||||
getDepsGraph: vi.fn().mockResolvedValue({ imports: new Map(), importedBy: new Map() }),
|
||||
setDepsGraph: vi.fn().mockResolvedValue(undefined),
|
||||
getProjectConfig: vi.fn().mockResolvedValue(null),
|
||||
setProjectConfig: vi.fn().mockResolvedValue(undefined),
|
||||
connect: vi.fn().mockResolvedValue(undefined),
|
||||
disconnect: vi.fn().mockResolvedValue(undefined),
|
||||
isConnected: vi.fn().mockReturnValue(true),
|
||||
clear: vi.fn().mockResolvedValue(undefined),
|
||||
}
|
||||
|
||||
session = new Session("test-session", "test-project")
|
||||
session.stats.editsApplied = 1
|
||||
|
||||
useCase = new UndoChange(mockSessionStorage, mockStorage)
|
||||
|
||||
vi.mocked(fs.stat).mockResolvedValue({
|
||||
size: 100,
|
||||
mtimeMs: Date.now(),
|
||||
} as unknown as Awaited<ReturnType<typeof fs.stat>>)
|
||||
})
|
||||
|
||||
afterEach(() => {
|
||||
vi.clearAllMocks()
|
||||
})
|
||||
|
||||
describe("execute", () => {
|
||||
it("should return error when no undo entries", async () => {
|
||||
vi.mocked(mockSessionStorage.popUndoEntry).mockResolvedValue(null)
|
||||
|
||||
const result = await useCase.execute(session)
|
||||
|
||||
expect(result.success).toBe(false)
|
||||
expect(result.error).toBe("No changes to undo")
|
||||
})
|
||||
|
||||
it("should restore previous content when file matches", async () => {
|
||||
const entry = createUndoEntry()
|
||||
vi.mocked(mockSessionStorage.popUndoEntry).mockResolvedValue(entry)
|
||||
vi.mocked(fs.readFile).mockResolvedValue("const a = 2")
|
||||
vi.mocked(fs.writeFile).mockResolvedValue(undefined)
|
||||
|
||||
session.addUndoEntry(entry)
|
||||
|
||||
const result = await useCase.execute(session)
|
||||
|
||||
expect(result.success).toBe(true)
|
||||
expect(result.entry).toBe(entry)
|
||||
expect(fs.writeFile).toHaveBeenCalledWith(entry.filePath, "const a = 1", "utf-8")
|
||||
})
|
||||
|
||||
it("should update storage after undo", async () => {
|
||||
const entry = createUndoEntry()
|
||||
vi.mocked(mockSessionStorage.popUndoEntry).mockResolvedValue(entry)
|
||||
vi.mocked(fs.readFile).mockResolvedValue("const a = 2")
|
||||
vi.mocked(fs.writeFile).mockResolvedValue(undefined)
|
||||
|
||||
session.addUndoEntry(entry)
|
||||
|
||||
await useCase.execute(session)
|
||||
|
||||
expect(mockStorage.setFile).toHaveBeenCalledWith(
|
||||
entry.filePath,
|
||||
expect.objectContaining({
|
||||
lines: entry.previousContent,
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("should decrement editsApplied counter", async () => {
|
||||
const entry = createUndoEntry()
|
||||
vi.mocked(mockSessionStorage.popUndoEntry).mockResolvedValue(entry)
|
||||
vi.mocked(fs.readFile).mockResolvedValue("const a = 2")
|
||||
vi.mocked(fs.writeFile).mockResolvedValue(undefined)
|
||||
|
||||
session.addUndoEntry(entry)
|
||||
const initialEdits = session.stats.editsApplied
|
||||
|
||||
await useCase.execute(session)
|
||||
|
||||
expect(session.stats.editsApplied).toBe(initialEdits - 1)
|
||||
})
|
||||
|
||||
it("should fail when file has been modified externally", async () => {
|
||||
const entry = createUndoEntry()
|
||||
vi.mocked(mockSessionStorage.popUndoEntry).mockResolvedValue(entry)
|
||||
vi.mocked(fs.readFile).mockResolvedValue("const a = 999")
|
||||
|
||||
const result = await useCase.execute(session)
|
||||
|
||||
expect(result.success).toBe(false)
|
||||
expect(result.error).toContain("modified since the change")
|
||||
})
|
||||
|
||||
it("should re-push undo entry on conflict", async () => {
|
||||
const entry = createUndoEntry()
|
||||
vi.mocked(mockSessionStorage.popUndoEntry).mockResolvedValue(entry)
|
||||
vi.mocked(fs.readFile).mockResolvedValue("const a = 999")
|
||||
|
||||
await useCase.execute(session)
|
||||
|
||||
expect(mockSessionStorage.pushUndoEntry).toHaveBeenCalledWith(session.id, entry)
|
||||
})
|
||||
|
||||
it("should handle empty file for undo", async () => {
|
||||
const entry = createUndoEntry({
|
||||
previousContent: [],
|
||||
newContent: ["new content"],
|
||||
})
|
||||
vi.mocked(mockSessionStorage.popUndoEntry).mockResolvedValue(entry)
|
||||
vi.mocked(fs.readFile).mockResolvedValue("new content")
|
||||
vi.mocked(fs.writeFile).mockResolvedValue(undefined)
|
||||
|
||||
session.addUndoEntry(entry)
|
||||
|
||||
const result = await useCase.execute(session)
|
||||
|
||||
expect(result.success).toBe(true)
|
||||
expect(fs.writeFile).toHaveBeenCalledWith(entry.filePath, "", "utf-8")
|
||||
})
|
||||
|
||||
it("should handle file not found during undo", async () => {
|
||||
const entry = createUndoEntry()
|
||||
vi.mocked(mockSessionStorage.popUndoEntry).mockResolvedValue(entry)
|
||||
const error = new Error("ENOENT") as NodeJS.ErrnoException
|
||||
error.code = "ENOENT"
|
||||
vi.mocked(fs.readFile).mockRejectedValue(error)
|
||||
|
||||
const result = await useCase.execute(session)
|
||||
|
||||
expect(result.success).toBe(false)
|
||||
})
|
||||
})
|
||||
|
||||
describe("canUndo", () => {
|
||||
it("should return false when stack is empty", async () => {
|
||||
vi.mocked(mockSessionStorage.getUndoStack).mockResolvedValue([])
|
||||
|
||||
const result = await useCase.canUndo(session)
|
||||
|
||||
expect(result).toBe(false)
|
||||
})
|
||||
|
||||
it("should return true when stack has entries", async () => {
|
||||
vi.mocked(mockSessionStorage.getUndoStack).mockResolvedValue([createUndoEntry()])
|
||||
|
||||
const result = await useCase.canUndo(session)
|
||||
|
||||
expect(result).toBe(true)
|
||||
})
|
||||
})
|
||||
|
||||
describe("peekUndoEntry", () => {
|
||||
it("should return null when stack is empty", async () => {
|
||||
vi.mocked(mockSessionStorage.getUndoStack).mockResolvedValue([])
|
||||
|
||||
const result = await useCase.peekUndoEntry(session)
|
||||
|
||||
expect(result).toBeNull()
|
||||
})
|
||||
|
||||
it("should return last entry without removing", async () => {
|
||||
const entry = createUndoEntry()
|
||||
vi.mocked(mockSessionStorage.getUndoStack).mockResolvedValue([entry])
|
||||
|
||||
const result = await useCase.peekUndoEntry(session)
|
||||
|
||||
expect(result).toBe(entry)
|
||||
expect(mockSessionStorage.popUndoEntry).not.toHaveBeenCalled()
|
||||
})
|
||||
})
|
||||
})
|
||||
Reference in New Issue
Block a user