feat(FN-1984): enrich chat and heartbeat prompts with agent context

- Add agent memory support to instruction resolution and introduce/export buildAgentChatPrompt for identity-aware prompt assembly with optional project memory
- Wire ChatManager to AgentStore so chat sessions can enrich system prompts per agent and include recent conversation context in prompt payloads
- Update server chat wiring to provide an AgentStore for prompt enrichment and keep graceful fallback behavior when agent context is unavailable
- Enrich HeartbeatMonitor system prompts with resolved agent instructions/ratings while preserving base-prompt fallback paths
- Expand dashboard and engine tests to cover prompt enrichment, fallback behavior, conversation context assembly, and AgentStore route mocks
This commit is contained in:
Fusion
2026-04-17 02:19:54 -07:00
committed by gsxdsm
parent 0873e3f788
commit a5a65e86fd
11 changed files with 431 additions and 57 deletions

View File

@@ -4,7 +4,7 @@
*/
import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
import { ChatManager, __setCreateKbAgent, __resetChatState } from "../chat.js";
import { ChatManager, __setBuildAgentChatPrompt, __setCreateKbAgent, __resetChatState } from "../chat.js";
// ── Mock Setup ──────────────────────────────────────────────────────────────
@@ -27,6 +27,15 @@ const mockChatStore = {
updateSession: vi.fn(),
};
const mockAgentStore = {
init: vi.fn(),
getAgent: vi.fn(),
};
function createChatManager(): ChatManager {
return new ChatManager(mockChatStore as any, "/tmp/test", mockAgentStore as any);
}
// ── Tests ───────────────────────────────────────────────────────────────────
describe("ChatManager.sendMessage", () => {
@@ -46,6 +55,26 @@ describe("ChatManager.sendMessage", () => {
role: "assistant",
content: "",
});
mockChatStore.getMessages.mockReturnValue([]);
mockAgentStore.init.mockResolvedValue(undefined);
mockAgentStore.getAgent.mockResolvedValue({
id: "agent-001",
name: "Avery",
role: "executor",
soul: "Be calm and precise.",
memory: "Remember to keep test coverage high.",
instructionsText: "Keep replies focused.",
});
__setBuildAgentChatPrompt(async ({ agent, basePrompt }: any) => {
return [
basePrompt,
`## Soul\n\n${agent.soul ?? ""}`,
`## Memory\n\n${agent.memory ?? ""}`,
`## Instructions\n\n${agent.instructionsText ?? ""}`,
].join("\n\n");
});
});
afterEach(() => {
@@ -78,10 +107,7 @@ describe("ChatManager.sendMessage", () => {
});
// Arrange
const chatManager = new ChatManager(
mockChatStore as any,
"/tmp/test",
);
const chatManager = createChatManager();
// Act
await chatManager.sendMessage("chat-001", "Hello");
@@ -114,10 +140,7 @@ describe("ChatManager.sendMessage", () => {
};
});
const chatManager = new ChatManager(
mockChatStore as any,
"/tmp/test",
);
const chatManager = createChatManager();
await chatManager.sendMessage("chat-001", "Hello");
@@ -148,10 +171,7 @@ describe("ChatManager.sendMessage", () => {
};
});
const chatManager = new ChatManager(
mockChatStore as any,
"/tmp/test",
);
const chatManager = createChatManager();
await chatManager.sendMessage("chat-001", "Hello");
@@ -179,10 +199,7 @@ describe("ChatManager.sendMessage", () => {
};
});
const chatManager = new ChatManager(
mockChatStore as any,
"/tmp/test",
);
const chatManager = createChatManager();
await chatManager.sendMessage("chat-001", "Hello");
@@ -216,10 +233,7 @@ describe("ChatManager.sendMessage", () => {
};
});
const chatManager = new ChatManager(
mockChatStore as any,
"/tmp/test",
);
const chatManager = createChatManager();
await chatManager.sendMessage("chat-001", "Hello");
@@ -243,10 +257,7 @@ describe("ChatManager.sendMessage", () => {
};
});
const chatManager = new ChatManager(
mockChatStore as any,
"/tmp/test",
);
const chatManager = createChatManager();
await chatManager.sendMessage("chat-001", "User message");
@@ -264,6 +275,118 @@ describe("ChatManager.sendMessage", () => {
expect(calls[1][1].role).toBe("assistant");
});
it("passes enriched system prompt with agent soul when agent context is available", async () => {
let createOptions: any;
__setCreateKbAgent(async (options: any) => {
createOptions = options;
return {
session: {
prompt: vi.fn().mockResolvedValue(undefined),
dispose: vi.fn(),
state: {
messages: [{ role: "assistant", content: "Done" }],
},
},
};
});
const chatManager = createChatManager();
await chatManager.sendMessage("chat-001", "Hello");
expect(mockAgentStore.init).toHaveBeenCalledTimes(1);
expect(mockAgentStore.getAgent).toHaveBeenCalledWith("agent-001");
expect(createOptions.systemPrompt).toContain("Be calm and precise.");
});
it("passes enriched system prompt with agent memory when agent context is available", async () => {
mockAgentStore.getAgent.mockResolvedValue({
id: "agent-001",
name: "Avery",
role: "executor",
soul: "Be concise.",
memory: "Remember repo conventions from prior tasks.",
instructionsText: "Focus on correctness.",
});
let createOptions: any;
__setCreateKbAgent(async (options: any) => {
createOptions = options;
return {
session: {
prompt: vi.fn().mockResolvedValue(undefined),
dispose: vi.fn(),
state: {
messages: [{ role: "assistant", content: "Done" }],
},
},
};
});
const chatManager = createChatManager();
await chatManager.sendMessage("chat-001", "Hello");
expect(createOptions.systemPrompt).toContain("Remember repo conventions from prior tasks.");
});
it("falls back to generic chat system prompt when agent lookup returns null", async () => {
mockAgentStore.getAgent.mockResolvedValue(null);
let createOptions: any;
__setCreateKbAgent(async (options: any) => {
createOptions = options;
return {
session: {
prompt: vi.fn().mockResolvedValue(undefined),
dispose: vi.fn(),
state: {
messages: [{ role: "assistant", content: "Done" }],
},
},
};
});
const chatManager = createChatManager();
await chatManager.sendMessage("chat-001", "Hello");
expect(createOptions.systemPrompt).toContain("You are a helpful AI assistant integrated into the fn task board system.");
expect(createOptions.systemPrompt).not.toContain("## Soul");
});
it("includes previous user and assistant messages in the prompt context", async () => {
const promptSpy = vi.fn().mockResolvedValue(undefined);
mockChatStore.getMessages.mockReturnValue([
{ role: "user", content: "Earlier user question" },
{ role: "assistant", content: "Earlier assistant answer" },
{ role: "system", content: "System note should be filtered" },
{ role: "user", content: "Current question" },
]);
__setCreateKbAgent(async () => {
return {
session: {
prompt: promptSpy,
dispose: vi.fn(),
state: {
messages: [{ role: "assistant", content: "Done" }],
},
},
};
});
const chatManager = createChatManager();
await chatManager.sendMessage("chat-001", "Current question");
expect(promptSpy).toHaveBeenCalledTimes(1);
const promptArgument = promptSpy.mock.calls[0]?.[0];
expect(promptArgument).toContain("## Previous Conversation");
expect(promptArgument).toContain("[User]: Earlier user question");
expect(promptArgument).toContain("[Assistant]: Earlier assistant answer");
expect(promptArgument).not.toContain("System note should be filtered");
expect(promptArgument).toContain("## Current Message");
expect(promptArgument).toContain("Current question");
});
it("generates title when session has no title", async () => {
mockSummarizeTitle.mockResolvedValue("Short Title");
@@ -279,10 +402,7 @@ describe("ChatManager.sendMessage", () => {
};
});
const chatManager = new ChatManager(
mockChatStore as any,
"/tmp/test",
);
const chatManager = createChatManager();
await chatManager.sendMessage("chat-001", "This is a long message that needs to be summarized");
@@ -316,10 +436,7 @@ describe("ChatManager.sendMessage", () => {
};
});
const chatManager = new ChatManager(
mockChatStore as any,
"/tmp/test",
);
const chatManager = createChatManager();
const longMessage = "A".repeat(300);
await chatManager.sendMessage("chat-001", longMessage);
@@ -354,10 +471,7 @@ describe("ChatManager.sendMessage", () => {
};
});
const chatManager = new ChatManager(
mockChatStore as any,
"/tmp/test",
);
const chatManager = createChatManager();
await chatManager.sendMessage("chat-001", "This is a long message");

View File

@@ -7,6 +7,8 @@ import { request, get } from "../test-request.js";
const mockInit = vi.fn().mockResolvedValue(undefined);
const mockClose = vi.fn().mockResolvedValue(undefined);
const mockGetNode = vi.fn();
const mockAgentStoreInit = vi.fn().mockResolvedValue(undefined);
const mockAgentStoreGetAgent = vi.fn().mockResolvedValue(null);
vi.mock("@fusion/core", () => {
return {
@@ -18,6 +20,10 @@ vi.mock("@fusion/core", () => {
ChatStore: class MockChatStore {
init = vi.fn().mockResolvedValue(undefined);
},
AgentStore: class MockAgentStore {
init = mockAgentStoreInit;
getAgent = mockAgentStoreGetAgent;
},
};
});

View File

@@ -30,6 +30,8 @@ const mockUpdateSettingsSyncState = vi.fn();
const mockApplyRemoteSettings = vi.fn();
const mockGetSettingsForSync = vi.fn();
const mockChatStoreInit = vi.fn().mockResolvedValue(undefined);
const mockAgentStoreInit = vi.fn().mockResolvedValue(undefined);
const mockAgentStoreGetAgent = vi.fn().mockResolvedValue(null);
vi.mock("@fusion/core", () => {
return {
@@ -47,6 +49,10 @@ vi.mock("@fusion/core", () => {
ChatStore: class MockChatStore {
init = mockChatStoreInit;
},
AgentStore: class MockAgentStore {
init = mockAgentStoreInit;
getAgent = mockAgentStoreGetAgent;
},
};
});

View File

@@ -15,6 +15,8 @@ const mockCheckNodeHealth = vi.fn();
const mockIsDiscoveryActive = vi.fn().mockReturnValue(false);
const mockGetDiscoveryConfig = vi.fn().mockReturnValue(null);
const mockChatStoreInit = vi.fn().mockResolvedValue(undefined);
const mockAgentStoreInit = vi.fn().mockResolvedValue(undefined);
const mockAgentStoreGetAgent = vi.fn().mockResolvedValue(null);
vi.mock("@fusion/core", () => {
return {
@@ -33,6 +35,10 @@ vi.mock("@fusion/core", () => {
ChatStore: class MockChatStore {
init = mockChatStoreInit;
},
AgentStore: class MockAgentStore {
init = mockAgentStoreInit;
getAgent = mockAgentStoreGetAgent;
},
};
});

View File

@@ -13,6 +13,8 @@
*/
import type {
Agent,
AgentStore,
ChatStore,
ChatSession,
ChatSessionCreateInput,
@@ -27,10 +29,12 @@ import { SessionEventBuffer } from "./sse-buffer.js";
type AgentResult = any;
// eslint-disable-next-line @typescript-eslint/no-explicit-any
let createKbAgent: any;
// eslint-disable-next-line @typescript-eslint/no-explicit-any
let buildAgentChatPromptFn: any;
// Initialize the import (this runs in actual server, mocked in tests)
async function initEngine() {
if (!createKbAgent) {
if (!createKbAgent || !buildAgentChatPromptFn) {
try {
// Use dynamic import with variable to prevent static analysis
const engineModule = "@fusion/engine";
@@ -38,11 +42,17 @@ async function initEngine() {
if (!createKbAgent) {
createKbAgent = engine.createKbAgent;
}
if (!buildAgentChatPromptFn) {
buildAgentChatPromptFn = engine.buildAgentChatPrompt;
}
} catch {
// Allow failure in test environments - agent functionality will be stubbed
if (!createKbAgent) {
createKbAgent = undefined;
}
if (!buildAgentChatPromptFn) {
buildAgentChatPromptFn = undefined;
}
}
}
}
@@ -255,9 +265,12 @@ export function getRateLimitResetTime(ip: string): Date | null {
* Creates sessions, sends messages, and streams AI responses via SSE.
*/
export class ChatManager {
private agentStoreReady?: Promise<void>;
constructor(
private chatStore: ChatStore,
private rootDir: string,
private agentStore?: AgentStore,
) {}
/**
@@ -357,10 +370,60 @@ export class ChatManager {
throw new Error("AI agent not available");
}
let systemPrompt = CHAT_SYSTEM_PROMPT;
let agent: Agent | null = null;
if (this.agentStore && session.agentId) {
try {
this.agentStoreReady ??= this.agentStore.init();
await this.agentStoreReady;
agent = await this.agentStore.getAgent(session.agentId);
} catch (agentLoadError) {
const message = agentLoadError instanceof Error ? agentLoadError.message : String(agentLoadError);
console.warn(`[chat] Failed to load agent context for ${session.agentId}: ${message}`);
}
}
if (agent && buildAgentChatPromptFn) {
try {
systemPrompt = await buildAgentChatPromptFn({
agent,
rootDir: this.rootDir,
agentStore: this.agentStore,
basePrompt: CHAT_SYSTEM_PROMPT,
includeProjectMemory: true,
});
} catch (promptBuildError) {
const message = promptBuildError instanceof Error ? promptBuildError.message : String(promptBuildError);
console.warn(`[chat] Failed to build enriched system prompt for ${agent.id}: ${message}`);
}
}
const allMessages = this.chatStore.getMessages(sessionId, { limit: 10000 }) ?? [];
const previousMessages = allMessages.slice(-51, -1);
const conversationMessages = previousMessages.filter(
(message) => message.role === "user" || message.role === "assistant",
);
const promptContent = conversationMessages.length > 0
? [
"## Previous Conversation",
"",
...conversationMessages.map((message) => {
const speaker = message.role === "user" ? "User" : "Assistant";
return `[${speaker}]: ${message.content}`;
}),
"",
"## Current Message",
"",
content,
].join("\n")
: content;
// Create AI agent session
agentResult = await createKbAgent({
cwd: this.rootDir,
systemPrompt: CHAT_SYSTEM_PROMPT,
systemPrompt,
tools: "readonly",
...(effectiveModelProvider && effectiveModelId
? {
@@ -385,7 +448,7 @@ export class ChatManager {
});
// Send user message and get response
await agentResult.session.prompt(content);
await agentResult.session.prompt(promptContent);
// Extract response text from agent state
let responseText = "";
@@ -452,10 +515,19 @@ export function __setCreateKbAgent(mock: typeof createKbAgent): void {
createKbAgent = mock;
}
/**
* Inject a mock buildAgentChatPrompt function. Used for testing only.
*/
export function __setBuildAgentChatPrompt(mock: typeof buildAgentChatPromptFn): void {
buildAgentChatPromptFn = mock;
}
/**
* Reset all chat state. Used for testing only.
*/
export function __resetChatState(): void {
chatStreamManager.reset();
rateLimits.clear();
engineReady = undefined;
buildAgentChatPromptFn = undefined;
}

View File

@@ -4,7 +4,7 @@ import { join, dirname } from "node:path";
import { existsSync } from "node:fs";
import { fileURLToPath } from "node:url";
import type { Task, TaskStore, MergeResult, AutomationStore, RoutineStore, CentralCore, MessageStore } from "@fusion/core";
import { ChatStore } from "@fusion/core";
import { AgentStore, ChatStore } from "@fusion/core";
import type { AuthStorageLike, ModelRegistryLike } from "./routes.js";
import { createApiRoutes } from "./routes.js";
import { createSSE } from "./sse.js";
@@ -603,8 +603,11 @@ export function createServer(store: TaskStore, options?: ServerOptions): ReturnT
// Create ChatStore for chat session management
const chatStore = options?.chatStore ?? new ChatStore(store.getFusionDir(), store.getDatabase());
// Create AgentStore for chat prompt enrichment (lazy-initialized inside ChatManager)
const chatAgentStore = new AgentStore({ rootDir: store.getFusionDir() });
// Create ChatManager for AI chat message handling
const chatManager = options?.chatManager ?? new ChatManager(chatStore, store.getRootDir());
const chatManager = options?.chatManager ?? new ChatManager(chatStore, store.getRootDir(), chatAgentStore);
const runAiSessionCleanup = (maxAgeMs: number, source: "initial" | "scheduled") => {
const result = aiSessionStore.cleanupStaleSessions(maxAgeMs);