feat(FN-2017): merge fusion/fn-2017
This commit is contained in:
@@ -779,4 +779,125 @@ describe("ChatManager.sendMessage", () => {
|
||||
// Assert - updateSession was NOT called
|
||||
expect(mockChatStore.updateSession).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("cancelGeneration returns false when no active generation exists", () => {
|
||||
const chatManager = createChatManager();
|
||||
|
||||
expect(chatManager.cancelGeneration("chat-001")).toBe(false);
|
||||
});
|
||||
|
||||
it("cancelGeneration returns true and aborts an active generation", () => {
|
||||
const chatManager = createChatManager();
|
||||
const abortController = new AbortController();
|
||||
const dispose = vi.fn();
|
||||
|
||||
(chatManager as any).activeGenerations.set("chat-001", {
|
||||
abortController,
|
||||
agentResult: { session: { dispose } },
|
||||
});
|
||||
|
||||
const events: Array<{ type: string; data: unknown }> = [];
|
||||
const unsubscribe = chatStreamManager.subscribe("chat-001", (event) => {
|
||||
events.push(event);
|
||||
});
|
||||
|
||||
const result = chatManager.cancelGeneration("chat-001");
|
||||
unsubscribe();
|
||||
|
||||
expect(result).toBe(true);
|
||||
expect(abortController.signal.aborted).toBe(true);
|
||||
expect(dispose).toHaveBeenCalledTimes(1);
|
||||
expect(events).toContainEqual({ type: "error", data: "Generation cancelled" });
|
||||
});
|
||||
|
||||
it("cancelled generation does not persist assistant message", async () => {
|
||||
let rejectPrompt: ((reason?: unknown) => void) | undefined;
|
||||
|
||||
__setCreateKbAgent(async () => {
|
||||
return {
|
||||
session: {
|
||||
prompt: vi.fn().mockImplementation(() => {
|
||||
return new Promise<void>((_resolve, reject) => {
|
||||
rejectPrompt = reject;
|
||||
});
|
||||
}),
|
||||
dispose: vi.fn().mockImplementation(() => {
|
||||
rejectPrompt?.(new Error("Disposed"));
|
||||
}),
|
||||
state: {
|
||||
messages: [{ role: "assistant", content: "Should not persist" }],
|
||||
},
|
||||
},
|
||||
};
|
||||
});
|
||||
|
||||
const chatManager = createChatManager();
|
||||
const sendPromise = chatManager.sendMessage("chat-001", "Hello");
|
||||
|
||||
await new Promise((resolve) => setTimeout(resolve, 0));
|
||||
|
||||
expect(chatManager.cancelGeneration("chat-001")).toBe(true);
|
||||
await sendPromise;
|
||||
|
||||
const assistantCalls = mockChatStore.addMessage.mock.calls.filter((call) => call[1].role === "assistant");
|
||||
expect(assistantCalls).toHaveLength(0);
|
||||
});
|
||||
|
||||
it("cancelled generation broadcasts error event with cancellation message", async () => {
|
||||
let rejectPrompt: ((reason?: unknown) => void) | undefined;
|
||||
|
||||
__setCreateKbAgent(async () => {
|
||||
return {
|
||||
session: {
|
||||
prompt: vi.fn().mockImplementation(() => {
|
||||
return new Promise<void>((_resolve, reject) => {
|
||||
rejectPrompt = reject;
|
||||
});
|
||||
}),
|
||||
dispose: vi.fn().mockImplementation(() => {
|
||||
rejectPrompt?.(new Error("Disposed"));
|
||||
}),
|
||||
state: { messages: [] },
|
||||
},
|
||||
};
|
||||
});
|
||||
|
||||
const events: Array<{ type: string; data: unknown }> = [];
|
||||
const unsubscribe = chatStreamManager.subscribe("chat-001", (event) => {
|
||||
events.push(event);
|
||||
});
|
||||
|
||||
const chatManager = createChatManager();
|
||||
const sendPromise = chatManager.sendMessage("chat-001", "Hello");
|
||||
|
||||
await new Promise((resolve) => setTimeout(resolve, 0));
|
||||
chatManager.cancelGeneration("chat-001");
|
||||
await sendPromise;
|
||||
unsubscribe();
|
||||
|
||||
expect(events.some((event) => event.type === "error" && event.data === "Generation cancelled")).toBe(true);
|
||||
});
|
||||
|
||||
it("cleans active generation state even when dispose fails", async () => {
|
||||
const disposeSpy = vi.fn().mockImplementation(() => {
|
||||
throw new Error("dispose failed");
|
||||
});
|
||||
|
||||
__setCreateKbAgent(async () => {
|
||||
return {
|
||||
session: {
|
||||
prompt: vi.fn().mockResolvedValue(undefined),
|
||||
dispose: disposeSpy,
|
||||
state: {
|
||||
messages: [{ role: "assistant", content: "Done" }],
|
||||
},
|
||||
},
|
||||
};
|
||||
});
|
||||
|
||||
const chatManager = createChatManager();
|
||||
await chatManager.sendMessage("chat-001", "Hello");
|
||||
|
||||
expect((chatManager as any).activeGenerations.has("chat-001")).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -115,7 +115,7 @@ function extractSSEPayload(sseChunk: string): unknown {
|
||||
const mockInit = vi.fn().mockResolvedValue(undefined);
|
||||
|
||||
// Create mock functions before vi.mock
|
||||
const { mockCreateKbAgent, mockChatStreamManager, mockSendMessage } = vi.hoisted(() => {
|
||||
const { mockCreateKbAgent, mockChatStreamManager, mockSendMessage, mockCancelGeneration } = vi.hoisted(() => {
|
||||
// Store subscribers per session for broadcast simulation
|
||||
const subscribers = new Map<string, Set<(event: any, eventId?: number) => void>>();
|
||||
|
||||
@@ -173,6 +173,7 @@ const { mockCreateKbAgent, mockChatStreamManager, mockSendMessage } = vi.hoisted
|
||||
return {
|
||||
mockCreateKbAgent: vi.fn(),
|
||||
mockSendMessage: vi.fn(),
|
||||
mockCancelGeneration: vi.fn(),
|
||||
mockChatStreamManager: chatStreamManager,
|
||||
};
|
||||
});
|
||||
@@ -226,6 +227,7 @@ vi.mock("../chat.js", () => {
|
||||
return {
|
||||
ChatManager: class MockChatManager {
|
||||
sendMessage = mockSendMessage;
|
||||
cancelGeneration = mockCancelGeneration;
|
||||
},
|
||||
chatStreamManager: mockChatStreamManager,
|
||||
checkRateLimit: vi.fn().mockReturnValue(true),
|
||||
@@ -324,7 +326,10 @@ const mockChatStoreInstance = {
|
||||
};
|
||||
|
||||
function createMockChatManager() {
|
||||
return { sendMessage: mockSendMessage };
|
||||
return {
|
||||
sendMessage: mockSendMessage,
|
||||
cancelGeneration: mockCancelGeneration,
|
||||
};
|
||||
}
|
||||
|
||||
// ── Tests ───────────────────────────────────────────────────────────────────
|
||||
@@ -371,6 +376,7 @@ describe("Chat API Routes", () => {
|
||||
mockGetLastMessageForSessions.mockReset();
|
||||
mockDeleteMessage.mockReset();
|
||||
mockSendMessage.mockReset();
|
||||
mockCancelGeneration.mockReset();
|
||||
mockAgentStoreInit.mockResolvedValue(undefined);
|
||||
mockAgentStoreGetAgent.mockReset();
|
||||
mockGetOrCreateProjectStore.mockReset();
|
||||
@@ -379,6 +385,7 @@ describe("Chat API Routes", () => {
|
||||
mockListSessions.mockReturnValue([]);
|
||||
mockGetMessages.mockReturnValue([]);
|
||||
mockGetLastMessageForSessions.mockReturnValue(new Map());
|
||||
mockCancelGeneration.mockReturnValue(false);
|
||||
|
||||
// Default agent mock - agent with model config
|
||||
mockAgentStoreGetAgent.mockResolvedValue({
|
||||
@@ -887,6 +894,45 @@ describe("Chat API Routes", () => {
|
||||
});
|
||||
});
|
||||
|
||||
describe("POST /api/chat/sessions/:id/cancel", () => {
|
||||
it("returns success true when generation is cancelled", async () => {
|
||||
mockCancelGeneration.mockReturnValue(true);
|
||||
|
||||
const response = await request(app, "POST", "/api/chat/sessions/chat-abc123/cancel");
|
||||
|
||||
expect(response.status).toBe(200);
|
||||
expect(response.body).toEqual({ success: true });
|
||||
expect(mockCancelGeneration).toHaveBeenCalledWith("chat-abc123");
|
||||
});
|
||||
|
||||
it("returns success false when no active generation exists", async () => {
|
||||
mockCancelGeneration.mockReturnValue(false);
|
||||
|
||||
const response = await request(app, "POST", "/api/chat/sessions/chat-abc123/cancel");
|
||||
|
||||
expect(response.status).toBe(200);
|
||||
expect(response.body).toEqual({ success: false });
|
||||
expect(mockCancelGeneration).toHaveBeenCalledWith("chat-abc123");
|
||||
});
|
||||
|
||||
it("returns 503 when chat manager is unavailable", async () => {
|
||||
const express = await import("express");
|
||||
const { createApiRoutes } = await import("../routes.js");
|
||||
|
||||
const appWithoutManager = express.default();
|
||||
appWithoutManager.use(express.json());
|
||||
appWithoutManager.use("/api", createApiRoutes(store as any, {
|
||||
chatStore: mockChatStore as any,
|
||||
chatManager: undefined,
|
||||
}));
|
||||
|
||||
const response = await request(appWithoutManager as any, "POST", "/api/chat/sessions/chat-abc123/cancel");
|
||||
|
||||
expect(response.status).toBe(503);
|
||||
expect((response.body as any).error).toContain("Chat manager not available");
|
||||
});
|
||||
});
|
||||
|
||||
describe("DELETE /api/chat/sessions/:id/messages/:messageId", () => {
|
||||
it("deletes message when session exists", async () => {
|
||||
mockGetSession.mockReturnValue(sampleSession);
|
||||
|
||||
@@ -371,6 +371,10 @@ export function getRateLimitResetTime(ip: string): Date | null {
|
||||
*/
|
||||
export class ChatManager {
|
||||
private agentStoreReady?: Promise<void>;
|
||||
private activeGenerations = new Map<string, {
|
||||
abortController: AbortController;
|
||||
agentResult?: AgentResult;
|
||||
}>();
|
||||
|
||||
constructor(
|
||||
private chatStore: ChatStore,
|
||||
@@ -507,72 +511,73 @@ export class ChatManager {
|
||||
modelProvider?: string,
|
||||
modelId?: string,
|
||||
): Promise<void> {
|
||||
// Validate session exists
|
||||
const abortController = new AbortController();
|
||||
this.activeGenerations.set(sessionId, { abortController });
|
||||
|
||||
const session = this.chatStore.getSession(sessionId);
|
||||
if (!session) {
|
||||
chatStreamManager.broadcast(sessionId, {
|
||||
type: "error",
|
||||
data: `Chat session ${sessionId} not found`,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
const hasMentionCandidates = /@[\w-]+/.test(content);
|
||||
const mentionAgents = hasMentionCandidates ? await this.listAgentsForMentions() : [];
|
||||
const mentions = hasMentionCandidates ? await this.parseMentions(content, mentionAgents) : [];
|
||||
|
||||
// Persist user message
|
||||
let _userMessageId: string;
|
||||
try {
|
||||
const userMessage = this.chatStore.addMessage(sessionId, {
|
||||
role: "user",
|
||||
content,
|
||||
metadata: mentions.length > 0 ? { mentions } : undefined,
|
||||
});
|
||||
_userMessageId = userMessage.id;
|
||||
} catch (err) {
|
||||
chatStreamManager.broadcast(sessionId, {
|
||||
type: "error",
|
||||
data: `Failed to save message: ${err instanceof Error ? err.message : "Unknown error"}`,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
// Use model from session if not overridden (needed for both AI response and title generation)
|
||||
const effectiveModelProvider = modelProvider ?? session.modelProvider ?? undefined;
|
||||
const effectiveModelId = modelId ?? session.modelId ?? undefined;
|
||||
|
||||
// Auto-generate chat title on first message if session has no title
|
||||
const needsTitle = session.title === null || session.title === undefined || session.title.trim() === "";
|
||||
if (needsTitle) {
|
||||
// Fire-and-forget title generation (non-blocking)
|
||||
(async () => {
|
||||
try {
|
||||
const generated = await summarizeTitle(
|
||||
content.trim(),
|
||||
this.rootDir,
|
||||
effectiveModelProvider,
|
||||
effectiveModelId,
|
||||
);
|
||||
const title = generated ?? content.trim().slice(0, 60).trim();
|
||||
if (title) {
|
||||
this.chatStore.updateSession(sessionId, { title });
|
||||
}
|
||||
} catch {
|
||||
// Fallback on any error
|
||||
const fallback = content.trim().slice(0, 60).trim();
|
||||
if (fallback) {
|
||||
this.chatStore.updateSession(sessionId, { title: fallback });
|
||||
}
|
||||
}
|
||||
})();
|
||||
}
|
||||
|
||||
let agentResult: AgentResult | undefined;
|
||||
let accumulatedThinking = "";
|
||||
let accumulatedText = "";
|
||||
|
||||
try {
|
||||
// Validate session exists
|
||||
if (!session) {
|
||||
chatStreamManager.broadcast(sessionId, {
|
||||
type: "error",
|
||||
data: `Chat session ${sessionId} not found`,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
const hasMentionCandidates = /@[\w-]+/.test(content);
|
||||
const mentionAgents = hasMentionCandidates ? await this.listAgentsForMentions() : [];
|
||||
const mentions = hasMentionCandidates ? await this.parseMentions(content, mentionAgents) : [];
|
||||
|
||||
// Persist user message
|
||||
try {
|
||||
this.chatStore.addMessage(sessionId, {
|
||||
role: "user",
|
||||
content,
|
||||
metadata: mentions.length > 0 ? { mentions } : undefined,
|
||||
});
|
||||
} catch (err) {
|
||||
chatStreamManager.broadcast(sessionId, {
|
||||
type: "error",
|
||||
data: `Failed to save message: ${err instanceof Error ? err.message : "Unknown error"}`,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
// Use model from session if not overridden (needed for both AI response and title generation)
|
||||
const effectiveModelProvider = modelProvider ?? session.modelProvider ?? undefined;
|
||||
const effectiveModelId = modelId ?? session.modelId ?? undefined;
|
||||
|
||||
// Auto-generate chat title on first message if session has no title
|
||||
const needsTitle = session.title === null || session.title === undefined || session.title.trim() === "";
|
||||
if (needsTitle) {
|
||||
// Fire-and-forget title generation (non-blocking)
|
||||
(async () => {
|
||||
try {
|
||||
const generated = await summarizeTitle(
|
||||
content.trim(),
|
||||
this.rootDir,
|
||||
effectiveModelProvider,
|
||||
effectiveModelId,
|
||||
);
|
||||
const title = generated ?? content.trim().slice(0, 60).trim();
|
||||
if (title) {
|
||||
this.chatStore.updateSession(sessionId, { title });
|
||||
}
|
||||
} catch {
|
||||
// Fallback on any error
|
||||
const fallback = content.trim().slice(0, 60).trim();
|
||||
if (fallback) {
|
||||
this.chatStore.updateSession(sessionId, { title: fallback });
|
||||
}
|
||||
}
|
||||
})();
|
||||
}
|
||||
|
||||
// Ensure engine is loaded
|
||||
await ensureEngineReady();
|
||||
|
||||
@@ -666,10 +671,20 @@ export class ChatManager {
|
||||
});
|
||||
},
|
||||
});
|
||||
this.activeGenerations.set(sessionId, { abortController, agentResult });
|
||||
|
||||
if (abortController.signal.aborted) {
|
||||
agentResult.session.dispose?.();
|
||||
return;
|
||||
}
|
||||
|
||||
// Send user message and get response
|
||||
await agentResult.session.prompt(promptContent);
|
||||
|
||||
if (abortController.signal.aborted) {
|
||||
return;
|
||||
}
|
||||
|
||||
// Extract response text from agent state
|
||||
let responseText = "";
|
||||
interface AgentMessage {
|
||||
@@ -707,6 +722,14 @@ export class ChatManager {
|
||||
data: { messageId: assistantMessage.id },
|
||||
});
|
||||
} catch (err) {
|
||||
if (abortController.signal.aborted) {
|
||||
chatStreamManager.broadcast(sessionId, {
|
||||
type: "error",
|
||||
data: "Generation cancelled",
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
const errorMessage = err instanceof Error ? err.message : "AI processing failed";
|
||||
console.error(`[chat] Error in sendMessage for session ${sessionId}:`, err);
|
||||
|
||||
@@ -728,6 +751,8 @@ export class ChatManager {
|
||||
data: errorMessage,
|
||||
});
|
||||
} finally {
|
||||
this.activeGenerations.delete(sessionId);
|
||||
|
||||
// Always dispose agent session
|
||||
if (agentResult) {
|
||||
try {
|
||||
@@ -738,6 +763,30 @@ export class ChatManager {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
cancelGeneration(sessionId: string): boolean {
|
||||
const entry = this.activeGenerations.get(sessionId);
|
||||
if (!entry) {
|
||||
return false;
|
||||
}
|
||||
|
||||
entry.abortController.abort();
|
||||
|
||||
if (entry.agentResult) {
|
||||
try {
|
||||
entry.agentResult.session.dispose?.();
|
||||
} catch (err) {
|
||||
console.error(`[chat] Error disposing agent session during cancellation:`, err);
|
||||
}
|
||||
}
|
||||
|
||||
chatStreamManager.broadcast(sessionId, {
|
||||
type: "error",
|
||||
data: "Generation cancelled",
|
||||
});
|
||||
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
// ── Test Helpers ────────────────────────────────────────────────────────────
|
||||
|
||||
@@ -9604,6 +9604,28 @@ export function createApiRoutes(store: TaskStore, options?: ServerOptions): Rout
|
||||
}
|
||||
});
|
||||
|
||||
/**
|
||||
* POST /api/chat/sessions/:id/cancel
|
||||
* Cancel an in-flight chat generation.
|
||||
*/
|
||||
router.post("/chat/sessions/:id/cancel", rateLimit(RATE_LIMITS.mutation), async (req, res) => {
|
||||
try {
|
||||
const chatManager = options?.chatManager;
|
||||
if (!chatManager) {
|
||||
throw new ApiError(503, "Chat manager not available");
|
||||
}
|
||||
|
||||
const sessionId = String(req.params.id);
|
||||
const success = chatManager.cancelGeneration(sessionId);
|
||||
res.json({ success });
|
||||
} catch (err: unknown) {
|
||||
if (err instanceof ApiError) {
|
||||
throw err;
|
||||
}
|
||||
rethrowAsApiError(err, "Failed to cancel chat generation");
|
||||
}
|
||||
});
|
||||
|
||||
/**
|
||||
* DELETE /api/chat/sessions/:id/messages/:messageId
|
||||
* Delete a specific message from a chat session.
|
||||
@@ -9653,6 +9675,7 @@ export function createApiRoutes(store: TaskStore, options?: ServerOptions): Rout
|
||||
"DELETE /chat/sessions/:id",
|
||||
"GET /chat/sessions/:id/messages",
|
||||
"POST /chat/sessions/:id/messages",
|
||||
"POST /chat/sessions/:id/cancel",
|
||||
"DELETE /chat/sessions/:id/messages/:messageId",
|
||||
];
|
||||
console.debug("[chat:routes:registered]", chatRoutes);
|
||||
|
||||
Reference in New Issue
Block a user