feat(FN-2017): merge fusion/fn-2017

This commit is contained in:
gsxdsm
2026-04-18 23:22:20 -07:00
parent 83458f1eca
commit 456c8ed6e6
14 changed files with 1095 additions and 90 deletions

View File

@@ -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);
});
});

View File

@@ -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);

View File

@@ -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 ────────────────────────────────────────────────────────────

View File

@@ -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);