From 89183c8dd8f91c2b4eebc7716933111f0310faed Mon Sep 17 00:00:00 2001 From: CC1227871 <2812624878@qq.com> Date: Fri, 11 Sep 2026 14:52:48 +0800 Subject: [PATCH 1/7] feat: add conversation branch lineage foundation --- packages/core/src/index.ts | 70 ++++++++++ packages/shared/src/kernel/agent-kernel.ts | 2 + .../kernel/conversation-projection.test.ts | 50 +++++++ .../src/kernel/conversation-projection.ts | 45 +++++++ .../shared/src/kernel/run-manager.test.ts | 2 + packages/shared/src/kernel/session-manager.ts | 125 +++++++++++++++++- .../shared/src/storage/branch-repository.ts | 49 +++++++ packages/shared/src/storage/index.ts | 1 + .../shared/src/storage/repositories.test.ts | 35 ++++- .../shared/src/storage/session-repository.ts | 1 + 10 files changed, 376 insertions(+), 4 deletions(-) create mode 100644 packages/shared/src/kernel/conversation-projection.test.ts create mode 100644 packages/shared/src/kernel/conversation-projection.ts create mode 100644 packages/shared/src/storage/branch-repository.ts diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index 0cc4058..40706a6 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -114,6 +114,8 @@ export interface Session { status: SessionStatus; createdAt: number; updatedAt: number; + /** Persisted branch selected when the session was last opened. */ + activeBranchId?: string; /** Runtime-context identity (e.g. Pi session id) so the runtime conversation can be recovered after restart. */ runtimeSessionId?: string; /** Runtime session file path (e.g. Pi JSONL session file). */ @@ -128,6 +130,41 @@ export interface SessionMeta extends Session { lastMessageAt?: number; } +/** + * A user-visible conversation action. Every action creates new artifacts; + * none of the values below mean "update the previous answer in place". + */ +export type ConversationOperation = 'send' | 'retry' | 'regenerate' | 'edit' | 'fork'; + +/** + * Snapshot of the workspace/security context used to start a run. The + * current workspace may change while a historical branch is being viewed, so + * replaying a run must use this persisted value instead of live UI state. + */ +export interface RunContextSnapshot { + capturedAt: number; + workspaceContext?: WorkspaceContext; + recentSymbols?: string[]; +} + +/** + * A logical conversation branch. `forkMessageId` is the last message from + * the parent branch visible in this branch; messages after it are local to the + * new branch. This makes branch materialization deterministic and preserves + * the original branch unchanged. + */ +export interface ConversationBranch { + id: string; + sessionId: string; + name: string; + createdAt: number; + updatedAt: number; + parentBranchId?: string; + forkMessageId?: string; + /** Pi's tree leaf for this branch, when the Pi runtime is in use. */ + runtimeLeafId?: string; +} + export interface Message { id: string; role: 'user' | 'assistant' | 'tool'; @@ -136,6 +173,17 @@ export interface Message { toolName?: string; toolCalls?: ToolCallRecord[]; trace?: AgentTraceEvent[]; + /** Branch containing this immutable message artifact. */ + branchId?: string; + /** Previous message in the visible branch at creation time. */ + parentMessageId?: string; + /** Run/generation that produced this message. */ + runId?: string; + generationId?: string; + operation?: ConversationOperation; + /** Message/run this operation was derived from (edit, retry, regenerate). */ + sourceMessageId?: string; + contextSnapshot?: RunContextSnapshot; } export type RunStatus = 'running' | 'completed' | 'failed' | 'cancelled'; @@ -150,6 +198,28 @@ export interface Run { completedAt?: number; answer?: string; error?: ApiError; + /** Conversation branch in which this generation ran. */ + branchId?: string; + operation?: ConversationOperation; + parentRunId?: string; + sourceMessageId?: string; + userMessageId?: string; + /** Explicit generation identity; normally equal to `id`, kept for APIs. */ + generationId?: string; + contextSnapshot?: RunContextSnapshot; + /** Immutable run manifest for audit/evaluation consumers. */ + manifest?: RunManifest; +} + +export interface RunManifest { + runId: string; + branchId: string; + operation: ConversationOperation; + inputMessageId?: string; + parentRunId?: string; + sourceMessageId?: string; + contextSnapshot?: RunContextSnapshot; + toolCallIds: string[]; } /** Live tool call state, streamed through agent events. */ diff --git a/packages/shared/src/kernel/agent-kernel.ts b/packages/shared/src/kernel/agent-kernel.ts index 8e6d4c2..5e0b63d 100644 --- a/packages/shared/src/kernel/agent-kernel.ts +++ b/packages/shared/src/kernel/agent-kernel.ts @@ -1,6 +1,7 @@ import type { AgentRuntime, ApiResult, ToolDefinition } from '@finagent/core'; import type { SkillHub } from '@finagent/skill-hub'; import { JsonFileStore } from '../storage/json-file-store.ts'; +import { BranchRepository } from '../storage/branch-repository.ts'; import { MessageRepository } from '../storage/message-repository.ts'; import { RunRepository } from '../storage/run-repository.ts'; import { SessionRepository } from '../storage/session-repository.ts'; @@ -57,6 +58,7 @@ export class AgentKernel { sessions: new SessionRepository(store), messages: new MessageRepository(store), runs: new RunRepository(store), + branches: new BranchRepository(store), piSessionDir: options.piSessionDir, now, }); diff --git a/packages/shared/src/kernel/conversation-projection.test.ts b/packages/shared/src/kernel/conversation-projection.test.ts new file mode 100644 index 0000000..f44fead --- /dev/null +++ b/packages/shared/src/kernel/conversation-projection.test.ts @@ -0,0 +1,50 @@ +import { describe, expect, it } from 'bun:test'; +import type { ConversationBranch, Message } from '@finagent/core'; +import { materializeBranchMessages } from './conversation-projection.ts'; + +const main: ConversationBranch = { + id: 'main', + sessionId: 's1', + name: 'Main', + createdAt: 1, + updatedAt: 1, +}; + +const alternative: ConversationBranch = { + id: 'alternative', + sessionId: 's1', + name: 'Alternative', + createdAt: 2, + updatedAt: 2, + parentBranchId: 'main', + forkMessageId: 'm2', +}; + +const messages: Message[] = [ + { id: 'm1', role: 'user', content: 'first', timestamp: 1, branchId: 'main' }, + { id: 'm2', role: 'assistant', content: 'first answer', timestamp: 2, branchId: 'main' }, + { id: 'm3', role: 'user', content: 'second', timestamp: 3, branchId: 'main' }, + { id: 'm4', role: 'assistant', content: 'second answer', timestamp: 4, branchId: 'main' }, + { id: 'm5', role: 'assistant', content: 'alternative answer', timestamp: 5, branchId: 'alternative' }, +]; + +describe('materializeBranchMessages', () => { + it('inherits only the parent prefix through the fork cursor', () => { + expect(materializeBranchMessages(messages, [main, alternative], 'alternative').map((message) => message.id)) + .toEqual(['m1', 'm2', 'm5']); + }); + + it('maps pre-branch legacy messages to Main', () => { + const legacy = messages.map(({ branchId: _branchId, ...message }) => message); + expect(materializeBranchMessages(legacy, [main], 'main').map((message) => message.id)) + .toEqual(['m1', 'm2', 'm3', 'm4', 'm5']); + }); + + it('rejects cyclic branch metadata', () => { + const cyclic: ConversationBranch[] = [ + { ...main, parentBranchId: 'alternative' }, + alternative, + ]; + expect(() => materializeBranchMessages(messages, cyclic, 'main')).toThrow('cycle'); + }); +}); diff --git a/packages/shared/src/kernel/conversation-projection.ts b/packages/shared/src/kernel/conversation-projection.ts new file mode 100644 index 0000000..3cfc201 --- /dev/null +++ b/packages/shared/src/kernel/conversation-projection.ts @@ -0,0 +1,45 @@ +import type { ConversationBranch, Message } from '@finagent/core'; + +/** + * Materialize one branch from append-only artifacts. A child branch inherits + * its parent transcript only through `forkMessageId`; its own messages are + * then appended. Missing branch ids are treated as legacy Main-branch data. + */ +export function materializeBranchMessages( + messages: Message[], + branches: ConversationBranch[], + branchId: string +): Message[] { + const byId = new Map(branches.map((branch) => [branch.id, branch])); + const rootId = branches.find((branch) => !branch.parentBranchId)?.id; + const normalized = messages.map((message) => ({ + ...message, + branchId: message.branchId ?? rootId, + })); + + const visit = (currentId: string, stack: Set): Message[] => { + if (stack.has(currentId)) { + throw new Error(`Conversation branch cycle detected at ${currentId}.`); + } + const branch = byId.get(currentId); + if (!branch) return normalized.filter((message) => message.branchId === currentId); + + const nextStack = new Set(stack); + nextStack.add(currentId); + let inherited: Message[] = []; + if (branch.parentBranchId) { + inherited = visit(branch.parentBranchId, nextStack); + if (branch.forkMessageId) { + const forkIndex = inherited.findIndex((message) => message.id === branch.forkMessageId); + inherited = forkIndex >= 0 ? inherited.slice(0, forkIndex + 1) : []; + } + } + + return [ + ...inherited, + ...normalized.filter((message) => message.branchId === branch.id), + ]; + }; + + return visit(branchId, new Set()); +} diff --git a/packages/shared/src/kernel/run-manager.test.ts b/packages/shared/src/kernel/run-manager.test.ts index 6d13af6..b20e991 100644 --- a/packages/shared/src/kernel/run-manager.test.ts +++ b/packages/shared/src/kernel/run-manager.test.ts @@ -11,6 +11,7 @@ import type { RuntimeSession, ToolDefinition, } from '@finagent/core'; +import { BranchRepository } from '../storage/branch-repository.ts'; import { JsonFileStore } from '../storage/json-file-store.ts'; import { MessageRepository } from '../storage/message-repository.ts'; import { RunRepository } from '../storage/run-repository.ts'; @@ -37,6 +38,7 @@ function makeKernel(script: (input: AgentRunInput) => AsyncIterable) sessions: new SessionRepository(store), messages: new MessageRepository(store), runs: new RunRepository(store), + branches: new BranchRepository(store), piSessionDir: join(dir, 'pi-sessions'), now: () => clock, }); diff --git a/packages/shared/src/kernel/session-manager.ts b/packages/shared/src/kernel/session-manager.ts index 2230be5..9baff9c 100644 --- a/packages/shared/src/kernel/session-manager.ts +++ b/packages/shared/src/kernel/session-manager.ts @@ -1,12 +1,19 @@ import { randomUUID } from 'node:crypto'; import { join } from 'node:path'; -import type { Message, Run, Session, SessionMeta } from '@finagent/core'; -import type { MessageRepository, RunRepository, SessionRepository } from '../storage/index.ts'; +import type { ConversationBranch, Message, Run, Session, SessionMeta } from '@finagent/core'; +import type { + BranchRepository, + MessageRepository, + RunRepository, + SessionRepository, +} from '../storage/index.ts'; +import { materializeBranchMessages } from './conversation-projection.ts'; export interface SessionManagerOptions { sessions: SessionRepository; messages: MessageRepository; runs: RunRepository; + branches: BranchRepository; /** Directory holding one runtime session file per Folio session. */ piSessionDir: string; now?: () => number; @@ -21,6 +28,7 @@ export class SessionManager { private readonly sessions: SessionRepository; private readonly messages: MessageRepository; private readonly runs: RunRepository; + private readonly branches: BranchRepository; private readonly piSessionDir: string; private readonly now: () => number; @@ -28,6 +36,7 @@ export class SessionManager { this.sessions = options.sessions; this.messages = options.messages; this.runs = options.runs; + this.branches = options.branches; this.piSessionDir = options.piSessionDir; this.now = options.now ?? Date.now; } @@ -42,6 +51,7 @@ export class SessionManager { async createSession(title?: string): Promise { const id = randomUUID(); + const branchId = randomUUID(); const now = this.now(); const session: Session = { id, @@ -49,10 +59,18 @@ export class SessionManager { status: 'idle', createdAt: now, updatedAt: now, + activeBranchId: branchId, runtimeSessionPath: join(this.piSessionDir, `${id}.jsonl`), }; const meta: SessionMeta = { ...session, messageCount: 0 }; await this.sessions.upsert(meta); + await this.branches.create({ + id: branchId, + sessionId: id, + name: 'Main', + createdAt: now, + updatedAt: now, + }); return meta; } @@ -73,7 +91,23 @@ export class SessionManager { return updated; } - async listMessages(sessionId: string): Promise { + /** + * Return the visible transcript for a branch. With no branch argument the + * session's persisted active branch is used, which is what the renderer + * should display after a reload. + */ + async listMessages(sessionId: string, branchId?: string): Promise { + const allMessages = await this.messages.list(sessionId); + const branch = branchId + ? await this.branches.get(sessionId, branchId) + : await this.getActiveBranch(sessionId); + if (!branch) return allMessages; + const branches = await this.branches.list(sessionId); + return materializeBranchMessages(allMessages, branches, branch.id); + } + + /** Raw append-only message artifacts, including messages hidden by branches. */ + async listAllMessages(sessionId: string): Promise { return this.messages.list(sessionId); } @@ -99,4 +133,89 @@ export class SessionManager { async getRun(sessionId: string, runId: string): Promise { return this.runs.get(sessionId, runId); } + + /** Return all branches, creating a Main branch for legacy sessions on first read. */ + async listBranches(sessionId: string): Promise { + const session = await this.sessions.get(sessionId); + if (!session) return []; + await this.ensureDefaultBranch(session); + return this.branches.list(sessionId); + } + + async getBranch(sessionId: string, branchId: string): Promise { + await this.listBranches(sessionId); + return this.branches.get(sessionId, branchId); + } + + async getActiveBranch(sessionId: string): Promise { + const session = await this.sessions.get(sessionId); + if (!session) return null; + const branch = await this.ensureDefaultBranch(session); + return branch; + } + + /** Persist the branch selected by the user and return it after validation. */ + async setActiveBranch(sessionId: string, branchId: string): Promise { + const session = await this.sessions.get(sessionId); + if (!session) return null; + const branch = await this.branches.get(sessionId, branchId); + if (!branch) return null; + await this.sessions.upsert({ ...session, activeBranchId: branchId, updatedAt: this.now() }); + return branch; + } + + /** Create an immutable child branch from the selected parent cursor. */ + async createBranch(input: { + sessionId: string; + name?: string; + parentBranchId?: string; + forkMessageId?: string; + runtimeLeafId?: string; + }): Promise { + const session = await this.sessions.get(input.sessionId); + if (!session) throw new Error(`Session ${input.sessionId} was not found.`); + + const now = this.now(); + const branch: ConversationBranch = { + id: randomUUID(), + sessionId: input.sessionId, + name: input.name?.trim() || `Branch ${now}`, + createdAt: now, + updatedAt: now, + parentBranchId: input.parentBranchId, + forkMessageId: input.forkMessageId, + runtimeLeafId: input.runtimeLeafId, + }; + if (branch.parentBranchId && !(await this.branches.get(input.sessionId, branch.parentBranchId))) { + throw new Error(`Parent branch ${branch.parentBranchId} was not found.`); + } + await this.branches.create(branch); + return branch; + } + + private async ensureDefaultBranch(session: SessionMeta): Promise { + const existing = await this.branches.list(session.id); + const selected = session.activeBranchId + ? existing.find((branch) => branch.id === session.activeBranchId) + : undefined; + if (selected) return selected; + + const root = existing.find((branch) => !branch.parentBranchId); + if (root) { + await this.sessions.upsert({ ...session, activeBranchId: root.id, updatedAt: this.now() }); + return root; + } + + const now = this.now(); + const branch: ConversationBranch = { + id: randomUUID(), + sessionId: session.id, + name: 'Main', + createdAt: session.createdAt || now, + updatedAt: now, + }; + await this.branches.create(branch); + await this.sessions.upsert({ ...session, activeBranchId: branch.id, updatedAt: now }); + return branch; + } } diff --git a/packages/shared/src/storage/branch-repository.ts b/packages/shared/src/storage/branch-repository.ts new file mode 100644 index 0000000..6840e7c --- /dev/null +++ b/packages/shared/src/storage/branch-repository.ts @@ -0,0 +1,49 @@ +import type { ConversationBranch } from '@finagent/core'; +import type { JsonFileStore } from './json-file-store.ts'; + +interface BranchesFile { + branches: ConversationBranch[]; +} + +/** Persists immutable conversation-branch metadata for one session. */ +export class BranchRepository { + private readonly store: JsonFileStore; + + constructor(store: JsonFileStore) { + this.store = store; + } + + private fileFor(sessionId: string): string { + return `sessions/${sessionId}/branches.json`; + } + + async list(sessionId: string): Promise { + const file = await this.store.read(this.fileFor(sessionId), { branches: [] }); + return [...file.branches].sort((a, b) => a.createdAt - b.createdAt); + } + + async get(sessionId: string, branchId: string): Promise { + const branches = await this.list(sessionId); + return branches.find((branch) => branch.id === branchId) ?? null; + } + + async create(branch: ConversationBranch): Promise { + const file = await this.store.read(this.fileFor(branch.sessionId), { branches: [] }); + if (file.branches.some((existing) => existing.id === branch.id)) { + throw new Error(`Conversation branch ${branch.id} already exists.`); + } + file.branches.push(branch); + await this.store.write(this.fileFor(branch.sessionId), file); + } + + async update(branch: ConversationBranch): Promise { + const file = await this.store.read(this.fileFor(branch.sessionId), { branches: [] }); + const index = file.branches.findIndex((existing) => existing.id === branch.id); + if (index >= 0) { + file.branches[index] = branch; + } else { + file.branches.push(branch); + } + await this.store.write(this.fileFor(branch.sessionId), file); + } +} diff --git a/packages/shared/src/storage/index.ts b/packages/shared/src/storage/index.ts index fce31eb..734ff56 100644 --- a/packages/shared/src/storage/index.ts +++ b/packages/shared/src/storage/index.ts @@ -2,3 +2,4 @@ export { JsonFileStore } from './json-file-store.ts'; export { SessionRepository } from './session-repository.ts'; export { MessageRepository } from './message-repository.ts'; export { RunRepository } from './run-repository.ts'; +export { BranchRepository } from './branch-repository.ts'; diff --git a/packages/shared/src/storage/repositories.test.ts b/packages/shared/src/storage/repositories.test.ts index e84068e..e914a59 100644 --- a/packages/shared/src/storage/repositories.test.ts +++ b/packages/shared/src/storage/repositories.test.ts @@ -2,7 +2,8 @@ import { mkdtemp, readdir, rm } from 'node:fs/promises'; import { tmpdir } from 'node:os'; import { join } from 'node:path'; import { afterEach, beforeEach, describe, expect, it } from 'bun:test'; -import type { Message, Run, SessionMeta } from '@finagent/core'; +import type { ConversationBranch, Message, Run, SessionMeta } from '@finagent/core'; +import { BranchRepository } from './branch-repository.ts'; import { JsonFileStore } from './json-file-store.ts'; import { MessageRepository } from './message-repository.ts'; import { RunRepository } from './run-repository.ts'; @@ -142,3 +143,35 @@ describe('RunRepository', () => { expect(runs.map((run) => run.id)).toEqual(['r2', 'r1']); }); }); + +describe('BranchRepository', () => { + it('persists branches and keeps them ordered by creation time', async () => { + const repo = new BranchRepository(store); + const main: ConversationBranch = { + id: 'b-main', + sessionId: 's1', + name: 'Main', + createdAt: 1000, + updatedAt: 1000, + }; + const alternative: ConversationBranch = { + id: 'b-alt', + sessionId: 's1', + name: 'Alternative', + createdAt: 2000, + updatedAt: 2000, + parentBranchId: main.id, + forkMessageId: 'm1', + }; + + await repo.create(alternative); + await repo.create(main); + const reloaded = new BranchRepository(new JsonFileStore(dir)); + + expect((await reloaded.list('s1')).map((branch) => branch.id)).toEqual(['b-main', 'b-alt']); + expect(await reloaded.get('s1', 'b-alt')).toMatchObject({ + parentBranchId: 'b-main', + forkMessageId: 'm1', + }); + }); +}); diff --git a/packages/shared/src/storage/session-repository.ts b/packages/shared/src/storage/session-repository.ts index babd570..85038cc 100644 --- a/packages/shared/src/storage/session-repository.ts +++ b/packages/shared/src/storage/session-repository.ts @@ -44,5 +44,6 @@ export class SessionRepository { await this.store.write(SessionRepository.FILE, file); await this.store.remove(`sessions/${id}/messages.json`); await this.store.remove(`sessions/${id}/runs.json`); + await this.store.remove(`sessions/${id}/branches.json`); } } From 6862779a6c6f125da0ce73bb93f80bcf5865922e Mon Sep 17 00:00:00 2001 From: CC1227871 <2812624878@qq.com> Date: Fri, 11 Sep 2026 15:06:24 +0800 Subject: [PATCH 2/7] feat: add conversation operations and branch-aware IPC --- apps/electron/src/main/index.ts | 24 ++ apps/electron/src/main/kernelHost.ts | 111 ++++-- apps/electron/src/preload/index.cjs | 6 + apps/electron/src/preload/index.ts | 26 ++ apps/electron/src/renderer/finagentClient.ts | 15 + packages/core/src/index.ts | 10 +- .../src/kernel/conversation-projection.ts | 4 +- .../shared/src/kernel/run-manager.test.ts | 104 ++++++ packages/shared/src/kernel/run-manager.ts | 346 ++++++++++++++++-- packages/shared/src/kernel/session-manager.ts | 4 +- packages/ui/src/atoms/runAtoms.ts | 36 +- packages/ui/src/atoms/sessionAtoms.test.ts | 6 + packages/ui/src/client.tsx | 30 ++ 13 files changed, 655 insertions(+), 67 deletions(-) diff --git a/apps/electron/src/main/index.ts b/apps/electron/src/main/index.ts index fca92c9..4c67ee5 100644 --- a/apps/electron/src/main/index.ts +++ b/apps/electron/src/main/index.ts @@ -116,6 +116,14 @@ ipcMain.handle('sessions:getMessages', async (_event, sessionId: unknown) => toIpcResult(() => agentKernelHost.getMessages(sessionId)) ); +ipcMain.handle('sessions:listBranches', async (_event, sessionId: unknown) => + toIpcResult(() => agentKernelHost.listBranches(sessionId)) +); + +ipcMain.handle('sessions:setActiveBranch', async (_event, input: unknown) => + toIpcResult(() => agentKernelHost.setActiveBranch(input)) +); + ipcMain.handle('sessions:listRuns', async (_event, sessionId: unknown) => toIpcResult(() => agentKernelHost.listRuns(sessionId)) ); @@ -124,6 +132,22 @@ ipcMain.handle('runs:start', async (_event, input: unknown) => toIpcResult(() => agentKernelHost.startRun(input)) ); +ipcMain.handle('runs:retry', async (_event, input: unknown) => + toIpcResult(() => agentKernelHost.retryRun(input)) +); + +ipcMain.handle('messages:edit', async (_event, input: unknown) => + toIpcResult(() => agentKernelHost.editMessage(input)) +); + +ipcMain.handle('messages:regenerate', async (_event, input: unknown) => + toIpcResult(() => agentKernelHost.regenerateMessage(input)) +); + +ipcMain.handle('branches:fork', async (_event, input: unknown) => + toIpcResult(() => agentKernelHost.forkBranch(input)) +); + ipcMain.handle('runs:cancel', async (_event, input: unknown) => toIpcResult(() => agentKernelHost.cancelRun(input)) ); diff --git a/apps/electron/src/main/kernelHost.ts b/apps/electron/src/main/kernelHost.ts index e9930b5..86d2eff 100644 --- a/apps/electron/src/main/kernelHost.ts +++ b/apps/electron/src/main/kernelHost.ts @@ -8,6 +8,7 @@ import type { ApiResult, CapabilityRegistry, Comparison, + ConversationBranch, CredentialInfo, CustomProviderConfig, FinancialProviderStatus, @@ -449,40 +450,25 @@ export class AgentKernelHost { return this.kernel.sessions.listMessages(requireString(sessionId, 'sessionId')); } + async listBranches(sessionId: unknown): Promise { + return this.kernel.sessions.listBranches(requireString(sessionId, 'sessionId')); + } + + async setActiveBranch(input: unknown): Promise { + const request = requireObject(input); + return this.kernel.runs.setActiveBranch( + requireString(request.sessionId, 'sessionId'), + requireString(request.branchId, 'branchId') + ); + } + async listRuns(sessionId: unknown): Promise { return this.kernel.sessions.listRuns(requireString(sessionId, 'sessionId')); } async startRun(input: unknown): Promise { const request = requireObject(input) as Partial; - let workspaceContext: WorkspaceContext | undefined; - if (request.workspaceContext && typeof request.workspaceContext === 'object') { - const context = request.workspaceContext as Record; - workspaceContext = {}; - if (typeof context.activeSymbol === 'string') { - workspaceContext.activeSymbol = context.activeSymbol.toUpperCase(); - } - if ( - context.activeView === 'overview' || - context.activeView === 'chart' || - context.activeView === 'financials' || - context.activeView === 'news' || - context.activeView === 'portfolio' - ) { - workspaceContext.activeView = context.activeView; - } - if (typeof context.selectedPosition === 'string') { - workspaceContext.selectedPosition = context.selectedPosition; - } - if ( - Array.isArray(context.comparisonSymbols) && - context.comparisonSymbols.every((entry) => typeof entry === 'string') - ) { - workspaceContext.comparisonSymbols = context.comparisonSymbols.map((entry) => - entry.toUpperCase() - ); - } - } + const workspaceContext = readWorkspaceContext(request.workspaceContext); // V8: new agent responses follow the *effective* app locale unless the // user explicitly requests another language in the prompt (spec §41–42). // Resolved after validation and with a safe fallback so a prefs failure @@ -495,6 +481,47 @@ export class AgentKernelHost { ); } + async retryRun(input: unknown): Promise { + const request = requireObject(input); + return this.kernel.runs.retryRun( + requireString(request.sessionId, 'sessionId'), + requireString(request.runId, 'runId'), + readWorkspaceContext(request.workspaceContext), + await this.effectiveRunLocale() + ); + } + + async editMessage(input: unknown): Promise { + const request = requireObject(input); + return this.kernel.runs.editMessage( + requireString(request.sessionId, 'sessionId'), + requireString(request.messageId, 'messageId'), + requireString(request.content, 'content'), + readWorkspaceContext(request.workspaceContext), + await this.effectiveRunLocale() + ); + } + + async regenerateMessage(input: unknown): Promise { + const request = requireObject(input); + return this.kernel.runs.regenerateMessage( + requireString(request.sessionId, 'sessionId'), + requireString(request.messageId, 'messageId'), + readWorkspaceContext(request.workspaceContext), + await this.effectiveRunLocale() + ); + } + + async forkBranch(input: unknown): Promise { + const request = requireObject(input); + const name = request.name === undefined ? undefined : requireString(request.name, 'name'); + return this.kernel.runs.forkBranch( + requireString(request.sessionId, 'sessionId'), + requireString(request.messageId, 'messageId'), + name + ); + } + private async effectiveRunLocale(): Promise { try { return (await this.getAppPreferences()).effectiveLocale; @@ -2241,6 +2268,34 @@ function requireObject(value: unknown): Record { return value as Record; } +function readWorkspaceContext(value: unknown): WorkspaceContext | undefined { + if (typeof value !== 'object' || value === null || Array.isArray(value)) return undefined; + const context = value as Record; + const result: WorkspaceContext = {}; + if (typeof context.activeSymbol === 'string') { + result.activeSymbol = context.activeSymbol.toUpperCase(); + } + if ( + context.activeView === 'overview' || + context.activeView === 'chart' || + context.activeView === 'financials' || + context.activeView === 'news' || + context.activeView === 'portfolio' + ) { + result.activeView = context.activeView; + } + if (typeof context.selectedPosition === 'string') { + result.selectedPosition = context.selectedPosition; + } + if ( + Array.isArray(context.comparisonSymbols) && + context.comparisonSymbols.every((entry) => typeof entry === 'string') + ) { + result.comparisonSymbols = context.comparisonSymbols.map((entry) => String(entry).toUpperCase()); + } + return result; +} + function createCodeError(code: string, message: string, action?: string) { const error = new Error(message) as Error & { code: string; action?: string }; error.code = code; diff --git a/apps/electron/src/preload/index.cjs b/apps/electron/src/preload/index.cjs index 5cc792a..2472ac1 100644 --- a/apps/electron/src/preload/index.cjs +++ b/apps/electron/src/preload/index.cjs @@ -39,8 +39,14 @@ var electronAPI = { createSession: (title) => import_electron.ipcRenderer.invoke("sessions:create", title), deleteSession: (sessionId) => import_electron.ipcRenderer.invoke("sessions:delete", sessionId), getMessages: (sessionId) => import_electron.ipcRenderer.invoke("sessions:getMessages", sessionId), + listBranches: (sessionId) => import_electron.ipcRenderer.invoke("sessions:listBranches", sessionId), + setActiveBranch: (input) => import_electron.ipcRenderer.invoke("sessions:setActiveBranch", input), listRuns: (sessionId) => import_electron.ipcRenderer.invoke("sessions:listRuns", sessionId), startRun: (input) => import_electron.ipcRenderer.invoke("runs:start", input), + retryRun: (input) => import_electron.ipcRenderer.invoke("runs:retry", input), + editMessage: (input) => import_electron.ipcRenderer.invoke("messages:edit", input), + regenerateMessage: (input) => import_electron.ipcRenderer.invoke("messages:regenerate", input), + forkBranch: (input) => import_electron.ipcRenderer.invoke("branches:fork", input), cancelRun: (input) => import_electron.ipcRenderer.invoke("runs:cancel", input), onAgentEvent: (callback) => { const listener = (_event, agentEvent) => callback(agentEvent); diff --git a/apps/electron/src/preload/index.ts b/apps/electron/src/preload/index.ts index 7acffc2..908391b 100644 --- a/apps/electron/src/preload/index.ts +++ b/apps/electron/src/preload/index.ts @@ -12,12 +12,23 @@ export interface ElectronAPI { createSession: (title?: string) => Promise; deleteSession: (sessionId: string) => Promise; getMessages: (sessionId: string) => Promise; + listBranches: (sessionId: string) => Promise; + setActiveBranch: (input: { sessionId: string; branchId: string }) => Promise; listRuns: (sessionId: string) => Promise; startRun: (input: { sessionId: string; content: string; workspaceContext?: unknown; }) => Promise; + retryRun: (input: { sessionId: string; runId: string; workspaceContext?: unknown }) => Promise; + editMessage: (input: { + sessionId: string; + messageId: string; + content: string; + workspaceContext?: unknown; + }) => Promise; + regenerateMessage: (input: { sessionId: string; messageId: string; workspaceContext?: unknown }) => Promise; + forkBranch: (input: { sessionId: string; messageId: string; name?: string }) => Promise; cancelRun: (input: { sessionId: string; runId: string }) => Promise; onAgentEvent: (callback: (event: unknown) => void) => () => void; }; @@ -186,12 +197,27 @@ const electronAPI: ElectronAPI = { createSession: (title?: string) => ipcRenderer.invoke('sessions:create', title), deleteSession: (sessionId: string) => ipcRenderer.invoke('sessions:delete', sessionId), getMessages: (sessionId: string) => ipcRenderer.invoke('sessions:getMessages', sessionId), + listBranches: (sessionId: string) => ipcRenderer.invoke('sessions:listBranches', sessionId), + setActiveBranch: (input: { sessionId: string; branchId: string }) => + ipcRenderer.invoke('sessions:setActiveBranch', input), listRuns: (sessionId: string) => ipcRenderer.invoke('sessions:listRuns', sessionId), startRun: (input: { sessionId: string; content: string; workspaceContext?: unknown; }) => ipcRenderer.invoke('runs:start', input), + retryRun: (input: { sessionId: string; runId: string; workspaceContext?: unknown }) => + ipcRenderer.invoke('runs:retry', input), + editMessage: (input: { + sessionId: string; + messageId: string; + content: string; + workspaceContext?: unknown; + }) => ipcRenderer.invoke('messages:edit', input), + regenerateMessage: (input: { sessionId: string; messageId: string; workspaceContext?: unknown }) => + ipcRenderer.invoke('messages:regenerate', input), + forkBranch: (input: { sessionId: string; messageId: string; name?: string }) => + ipcRenderer.invoke('branches:fork', input), cancelRun: (input: { sessionId: string; runId: string }) => ipcRenderer.invoke('runs:cancel', input), onAgentEvent: (callback: (event: unknown) => void) => { const listener = (_event: Electron.IpcRendererEvent, agentEvent: unknown) => callback(agentEvent); diff --git a/apps/electron/src/renderer/finagentClient.ts b/apps/electron/src/renderer/finagentClient.ts index 1fe075e..1ef9d58 100644 --- a/apps/electron/src/renderer/finagentClient.ts +++ b/apps/electron/src/renderer/finagentClient.ts @@ -29,9 +29,24 @@ function createElectronClient(): FinagentClient { createSession: (title?: string) => ipcResult(window.electronAPI.kernel.createSession(title)), deleteSession: (sessionId: string) => ipcResult(window.electronAPI.kernel.deleteSession(sessionId)), getMessages: (sessionId: string) => ipcResult(window.electronAPI.kernel.getMessages(sessionId)), + listBranches: (sessionId: string) => ipcResult(window.electronAPI.kernel.listBranches(sessionId)), + setActiveBranch: (sessionId: string, branchId: string) => + ipcResult(window.electronAPI.kernel.setActiveBranch({ sessionId, branchId })), listRuns: (sessionId: string) => ipcResult(window.electronAPI.kernel.listRuns(sessionId)), startRun: (sessionId: string, content: string, workspaceContext?: WorkspaceContext) => ipcResult(window.electronAPI.kernel.startRun({ sessionId, content, workspaceContext })), + retryRun: (sessionId: string, runId: string, workspaceContext?: WorkspaceContext) => + ipcResult(window.electronAPI.kernel.retryRun({ sessionId, runId, workspaceContext })), + editMessage: ( + sessionId: string, + messageId: string, + content: string, + workspaceContext?: WorkspaceContext + ) => ipcResult(window.electronAPI.kernel.editMessage({ sessionId, messageId, content, workspaceContext })), + regenerateMessage: (sessionId: string, messageId: string, workspaceContext?: WorkspaceContext) => + ipcResult(window.electronAPI.kernel.regenerateMessage({ sessionId, messageId, workspaceContext })), + forkBranch: (sessionId: string, messageId: string, name?: string) => + ipcResult(window.electronAPI.kernel.forkBranch({ sessionId, messageId, name })), cancelRun: (sessionId: string, runId: string) => ipcResult(window.electronAPI.kernel.cancelRun({ sessionId, runId })), onAgentEvent: (callback: (event: AgentEvent) => void) => diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index 40706a6..928e677 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -150,8 +150,10 @@ export interface RunContextSnapshot { /** * A logical conversation branch. `forkMessageId` is the last message from * the parent branch visible in this branch; messages after it are local to the - * new branch. This makes branch materialization deterministic and preserves - * the original branch unchanged. + * new branch. `null` explicitly means "start before the first parent + * message", while `undefined` is reserved for legacy metadata. This makes + * branch materialization deterministic and preserves the original branch + * unchanged. */ export interface ConversationBranch { id: string; @@ -160,7 +162,7 @@ export interface ConversationBranch { createdAt: number; updatedAt: number; parentBranchId?: string; - forkMessageId?: string; + forkMessageId?: string | null; /** Pi's tree leaf for this branch, when the Pi runtime is in use. */ runtimeLeafId?: string; } @@ -266,6 +268,8 @@ export type AgentEvent = export interface RunStartedPayload { run: Run; userMessage: Message; + /** False when a new generation reuses an inherited user message. */ + userMessageIsNew?: boolean; } /** diff --git a/packages/shared/src/kernel/conversation-projection.ts b/packages/shared/src/kernel/conversation-projection.ts index 3cfc201..15de9e5 100644 --- a/packages/shared/src/kernel/conversation-projection.ts +++ b/packages/shared/src/kernel/conversation-projection.ts @@ -29,7 +29,9 @@ export function materializeBranchMessages( let inherited: Message[] = []; if (branch.parentBranchId) { inherited = visit(branch.parentBranchId, nextStack); - if (branch.forkMessageId) { + if (branch.forkMessageId === null) { + inherited = []; + } else if (branch.forkMessageId !== undefined) { const forkIndex = inherited.findIndex((message) => message.id === branch.forkMessageId); inherited = forkIndex >= 0 ? inherited.slice(0, forkIndex + 1) : []; } diff --git a/packages/shared/src/kernel/run-manager.test.ts b/packages/shared/src/kernel/run-manager.test.ts index b20e991..49ac515 100644 --- a/packages/shared/src/kernel/run-manager.test.ts +++ b/packages/shared/src/kernel/run-manager.test.ts @@ -282,6 +282,110 @@ describe('RunManager', () => { code: 'SESSION_NOT_FOUND', }); }); + + it('regenerates without duplicating the inherited user message or overwriting the old answer', async () => { + const { sessions, runs } = makeKernel((input) => completedScript(`answer-${input.runId}`)(input)); + const session = await sessions.createSession('A'); + + await runs.startRun(session.id, 'compare NVDA'); + await waitFor(async () => !runs.isRunning()); + const originalMessages = await sessions.listMessages(session.id); + const originalAnswer = originalMessages[1]; + const regenerated = await runs.regenerateMessage(session.id, originalAnswer.id); + await waitFor(async () => !runs.isRunning()); + + const branches = await sessions.listBranches(session.id); + expect(branches).toHaveLength(2); + expect(regenerated.operation).toBe('regenerate'); + expect(regenerated.parentRunId).toBeTruthy(); + expect(regenerated.manifest).toMatchObject({ operation: 'regenerate', toolCallIds: ['t1'] }); + expect((await sessions.listMessages(session.id)).map((message) => message.id)).toEqual([ + originalMessages[0].id, + `assistant-${regenerated.id}`, + ]); + expect((await sessions.listAllMessages(session.id)).map((message) => message.id)).toEqual([ + originalMessages[0].id, + originalAnswer.id, + `assistant-${regenerated.id}`, + ]); + }); + + it('edits from the selected historical point and excludes later answers from the new branch', async () => { + const { sessions, runs } = makeKernel(completedScript('answer')); + const session = await sessions.createSession('A'); + + await runs.startRun(session.id, 'first question'); + await waitFor(async () => !runs.isRunning()); + await runs.startRun(session.id, 'second question'); + await waitFor(async () => !runs.isRunning()); + const original = await sessions.listMessages(session.id); + const editedRun = await runs.editMessage(session.id, original[2].id, 'edited second question'); + await waitFor(async () => !runs.isRunning()); + + expect(editedRun.operation).toBe('edit'); + expect(editedRun.sourceMessageId).toBe(original[2].id); + expect((await sessions.listMessages(session.id)).map((message) => message.content)).toEqual([ + 'first question', + 'answer', + 'edited second question', + 'answer', + ]); + expect((await sessions.listMessages(session.id)).some((message) => message.content === 'second question')).toBe(false); + }); + + it('retries only failed runs and creates a new generation on a child branch', async () => { + let attempts = 0; + const { sessions, runs } = makeKernel(async function* (input) { + attempts += 1; + if (attempts === 1) { + yield event(input.sessionId, input.runId, 'run_failed', { + error: { code: 'TOOL_EXECUTION_ERROR', message: 'temporary failure' }, + }); + return; + } + yield* completedScript('retried answer')(input); + }); + const session = await sessions.createSession('A'); + + const failed = await runs.startRun(session.id, 'retry me'); + await waitFor(async () => !runs.isRunning()); + const retried = await runs.retryRun(session.id, failed.id); + await waitFor(async () => !runs.isRunning()); + + expect(retried.operation).toBe('retry'); + expect(retried.parentRunId).toBe(failed.id); + expect(retried.id).not.toBe(failed.id); + expect((await sessions.listMessages(session.id)).map((message) => message.content)).toEqual([ + 'retry me', + 'retried answer', + ]); + await expect(runs.retryRun(session.id, retried.id)).rejects.toMatchObject({ code: 'INVALID_ARGUMENT' }); + }); + + it('forks from an earlier message and keeps the fork cursor after reload', async () => { + const { sessions, runs } = makeKernel(completedScript('answer')); + const session = await sessions.createSession('A'); + await runs.startRun(session.id, 'first'); + await waitFor(async () => !runs.isRunning()); + await runs.startRun(session.id, 'second'); + await waitFor(async () => !runs.isRunning()); + const original = await sessions.listMessages(session.id); + + const branch = await runs.forkBranch(session.id, original[1].id, 'Research alternative'); + expect(branch.parentBranchId).toBeTruthy(); + expect((await sessions.listMessages(session.id)).map((message) => message.content)).toEqual(['first', 'answer']); + + await runs.startRun(session.id, 'alternative question'); + await waitFor(async () => !runs.isRunning()); + const reloaded = new SessionRepository(new JsonFileStore(dir)); + expect((await reloaded.get(session.id))?.activeBranchId).toBe(branch.id); + expect((await sessions.listMessages(session.id)).map((message) => message.content)).toEqual([ + 'first', + 'answer', + 'alternative question', + 'answer', + ]); + }); }); async function waitFor(predicate: () => Promise, timeoutMs = 2000) { diff --git a/packages/shared/src/kernel/run-manager.ts b/packages/shared/src/kernel/run-manager.ts index 91621eb..5b4c128 100644 --- a/packages/shared/src/kernel/run-manager.ts +++ b/packages/shared/src/kernel/run-manager.ts @@ -4,8 +4,11 @@ import type { AgentEventPayload, AgentRuntime, ApiError, + ConversationBranch, + ConversationOperation, Message, Run, + RunContextSnapshot, SessionMeta, ToolCall, ToolCallRecord, @@ -29,6 +32,19 @@ interface ActiveRun { cancelRequested: boolean; } +interface GenerationOptions { + session: SessionMeta; + branch: ConversationBranch; + content: string; + operation: ConversationOperation; + sourceMessageId?: string; + parentRunId?: string; + parentMessageId?: string; + existingUserMessage?: Message; + workspaceContext?: WorkspaceContext; + locale?: SupportedLocale; +} + /** * Starts, observes, persists, and terminates runs. * @@ -74,10 +90,187 @@ export class RunManager { workspaceContext?: WorkspaceContext, locale?: SupportedLocale ): Promise { + const session = await this.sessions.getSession(sessionId); + if (!session) { + throw createCodeError('SESSION_NOT_FOUND', `Session ${sessionId} was not found.`); + } + const branch = await this.sessions.getActiveBranch(sessionId); + if (!branch) { + throw createCodeError('BRANCH_NOT_FOUND', `The active branch for session ${sessionId} was not found.`); + } + const visibleMessages = await this.sessions.listMessages(sessionId, branch.id); + return this.startGeneration({ + session, + branch, + content, + operation: 'send', + parentMessageId: visibleMessages.at(-1)?.id, + workspaceContext, + locale, + }); + } + + /** Retry a failed/cancelled generation on a new child branch. */ + async retryRun( + sessionId: string, + runId: string, + workspaceContext?: WorkspaceContext, + locale?: SupportedLocale + ): Promise { + this.assertNoActiveRun(); + const sourceRun = await this.sessions.getRun(sessionId, runId); + if (!sourceRun) throw createCodeError('RUN_NOT_FOUND', `Run ${runId} was not found.`); + if (sourceRun.status !== 'failed' && sourceRun.status !== 'cancelled') { + throw createCodeError('INVALID_ARGUMENT', 'Only failed or cancelled runs can be retried.'); + } + + const sourceBranch = await this.branchForRun(sessionId, sourceRun); + const visible = await this.sessions.listMessages(sessionId, sourceBranch.id); + const userIndex = findUserMessageIndex(visible, sourceRun); + const sourceUser = userIndex >= 0 ? visible[userIndex] : undefined; + const child = await this.createDerivedBranch( + sessionId, + sourceBranch, + userIndex > 0 ? visible[userIndex - 1].id : null, + `Retry · ${sourceRun.input.slice(0, 32)}` + ); + const session = await this.requireSession(sessionId); + return this.startGeneration({ + session, + branch: child, + content: sourceRun.input, + operation: 'retry', + sourceMessageId: sourceUser?.id, + parentRunId: sourceRun.id, + parentMessageId: userIndex > 0 ? visible[userIndex - 1].id : undefined, + workspaceContext, + locale, + }); + } + + /** Regenerate an assistant artifact while inheriting the same user context. */ + async regenerateMessage( + sessionId: string, + assistantMessageId: string, + workspaceContext?: WorkspaceContext, + locale?: SupportedLocale + ): Promise { + this.assertNoActiveRun(); + const sourceBranch = await this.sessions.getActiveBranch(sessionId); + if (!sourceBranch) throw createCodeError('BRANCH_NOT_FOUND', 'The active branch was not found.'); + const visible = await this.sessions.listMessages(sessionId, sourceBranch.id); + const assistant = visible.find((message) => message.id === assistantMessageId); + if (!assistant || assistant.role !== 'assistant') { + throw createCodeError('MESSAGE_NOT_FOUND', `Assistant message ${assistantMessageId} was not found.`); + } + const sourceRun = assistant.runId ? await this.sessions.getRun(sessionId, assistant.runId) : null; + if (!sourceRun) { + throw createCodeError('RUN_NOT_FOUND', `The generation for message ${assistantMessageId} was not found.`); + } + const userMessage = sourceRun.userMessageId + ? visible.find((message) => message.id === sourceRun.userMessageId) + : findPreviousUserMessage(visible, visible.indexOf(assistant)); + if (!userMessage) { + throw createCodeError('MESSAGE_NOT_FOUND', 'The user context for this answer was not found.'); + } + const child = await this.createDerivedBranch( + sessionId, + sourceBranch, + userMessage.id, + `Regenerate · ${userMessage.content.slice(0, 28)}` + ); + const session = await this.requireSession(sessionId); + return this.startGeneration({ + session, + branch: child, + content: userMessage.content, + operation: 'regenerate', + sourceMessageId: assistant.id, + parentRunId: sourceRun.id, + parentMessageId: userMessage.id, + existingUserMessage: userMessage, + workspaceContext, + locale, + }); + } + + /** Edit a historical user message and rerun from the preceding cursor. */ + async editMessage( + sessionId: string, + userMessageId: string, + content: string, + workspaceContext?: WorkspaceContext, + locale?: SupportedLocale + ): Promise { + this.assertNoActiveRun(); const text = content.trim(); - if (!text) { - throw createCodeError('INVALID_ARGUMENT', 'Message content is required.'); + if (!text) throw createCodeError('INVALID_ARGUMENT', 'Message content is required.'); + const sourceBranch = await this.sessions.getActiveBranch(sessionId); + if (!sourceBranch) throw createCodeError('BRANCH_NOT_FOUND', 'The active branch was not found.'); + const visible = await this.sessions.listMessages(sessionId, sourceBranch.id); + const sourceIndex = visible.findIndex((message) => message.id === userMessageId); + const sourceMessage = sourceIndex >= 0 ? visible[sourceIndex] : undefined; + if (!sourceMessage || sourceMessage.role !== 'user') { + throw createCodeError('MESSAGE_NOT_FOUND', `User message ${userMessageId} was not found.`); } + const child = await this.createDerivedBranch( + sessionId, + sourceBranch, + sourceIndex > 0 ? visible[sourceIndex - 1].id : null, + `Edit · ${text.slice(0, 32)}` + ); + const session = await this.requireSession(sessionId); + return this.startGeneration({ + session, + branch: child, + content: text, + operation: 'edit', + sourceMessageId: sourceMessage.id, + parentMessageId: sourceIndex > 0 ? visible[sourceIndex - 1].id : undefined, + workspaceContext, + locale, + }); + } + + /** Create and activate a branch from an earlier visible message. */ + async forkBranch(sessionId: string, messageId: string, name?: string): Promise { + this.assertNoActiveRun(); + const sourceBranch = await this.sessions.getActiveBranch(sessionId); + if (!sourceBranch) throw createCodeError('BRANCH_NOT_FOUND', 'The active branch was not found.'); + const visible = await this.sessions.listMessages(sessionId, sourceBranch.id); + if (!visible.some((message) => message.id === messageId)) { + throw createCodeError('MESSAGE_NOT_FOUND', `Message ${messageId} was not found.`); + } + const child = await this.createDerivedBranch( + sessionId, + sourceBranch, + messageId, + name?.trim() || `Fork · ${new Date(this.now()).toISOString()}` + ); + await this.sessions.setActiveBranch(sessionId, child.id); + return child; + } + + async setActiveBranch(sessionId: string, branchId: string): Promise { + this.assertNoActiveRun(); + const branch = await this.sessions.setActiveBranch(sessionId, branchId); + if (!branch) throw createCodeError('BRANCH_NOT_FOUND', `Branch ${branchId} was not found.`); + return branch; + } + + /** Abort the given run if it is the one currently executing. */ + async cancelRun(sessionId: string, runId: string): Promise { + const active = this.activeRun; + if (!active || active.sessionId !== sessionId || active.runId !== runId) { + return; + } + active.cancelRequested = true; + await this.runtime.cancel({ sessionId, runId }); + } + + private async startGeneration(options: GenerationOptions): Promise { + const text = options.content.trim(); + if (!text) throw createCodeError('INVALID_ARGUMENT', 'Message content is required.'); if (this.activeRun) { throw createCodeError( 'RUN_IN_PROGRESS', @@ -85,53 +278,113 @@ export class RunManager { ); } - const session = await this.sessions.getSession(sessionId); - if (!session) { - throw createCodeError('SESSION_NOT_FOUND', `Session ${sessionId} was not found.`); - } - const now = this.now(); - const run: Run = { + const runId = randomUUID(); + const contextSnapshot: RunContextSnapshot = { + capturedAt: now, + workspaceContext: options.workspaceContext, + recentSymbols: options.session.recentSymbols ? [...options.session.recentSymbols] : undefined, + }; + const userMessage = options.existingUserMessage ?? { id: randomUUID(), - sessionId, + role: 'user' as const, + content: text, + timestamp: now, + branchId: options.branch.id, + parentMessageId: options.parentMessageId, + runId, + generationId: runId, + operation: options.operation, + sourceMessageId: options.sourceMessageId, + contextSnapshot, + }; + const run: Run = { + id: runId, + sessionId: options.session.id, status: 'running', input: text, startedAt: now, + branchId: options.branch.id, + operation: options.operation, + parentRunId: options.parentRunId, + sourceMessageId: options.sourceMessageId, + userMessageId: userMessage.id, + generationId: runId, + contextSnapshot, + manifest: { + runId, + branchId: options.branch.id, + operation: options.operation, + inputMessageId: userMessage.id, + parentRunId: options.parentRunId, + sourceMessageId: options.sourceMessageId, + contextSnapshot, + toolCallIds: [], + }, }; - await this.runs.create(run); - const userMessage: Message = { - id: randomUUID(), - role: 'user', - content: text, - timestamp: now, - }; - await this.sessions.appendMessage(sessionId, userMessage); - await this.sessions.updateSession(sessionId, { status: 'running' }); + await this.runs.create(run); + if (!options.existingUserMessage) { + await this.sessions.appendMessage(options.session.id, userMessage); + } + await this.sessions.updateSession(options.session.id, { status: 'running' }); - this.activeRun = { sessionId, runId: run.id, cancelRequested: false }; + this.activeRun = { sessionId: options.session.id, runId: run.id, cancelRequested: false }; this.emit({ id: randomUUID(), - sessionId, + sessionId: options.session.id, runId: run.id, type: 'run_started', timestamp: now, sequence: 1, - payload: { run, userMessage }, + payload: { + run, + userMessage, + userMessageIsNew: !options.existingUserMessage, + }, }); - void this.execute(run, session, workspaceContext, locale); + void this.execute(run, options.session, options.workspaceContext, options.locale); return run; } - /** Abort the given run if it is the one currently executing. */ - async cancelRun(sessionId: string, runId: string): Promise { - const active = this.activeRun; - if (!active || active.sessionId !== sessionId || active.runId !== runId) { - return; + private assertNoActiveRun(): void { + if (this.activeRun) { + throw createCodeError( + 'RUN_IN_PROGRESS', + 'Another run is still in progress. Stop it before changing conversation branches.' + ); } - active.cancelRequested = true; - await this.runtime.cancel({ sessionId, runId }); + } + + private async requireSession(sessionId: string): Promise { + const session = await this.sessions.getSession(sessionId); + if (!session) throw createCodeError('SESSION_NOT_FOUND', `Session ${sessionId} was not found.`); + return session; + } + + private async branchForRun(sessionId: string, run: Run): Promise { + const branch = run.branchId + ? await this.sessions.getBranch(sessionId, run.branchId) + : await this.sessions.getActiveBranch(sessionId); + if (!branch) throw createCodeError('BRANCH_NOT_FOUND', 'The source branch was not found.'); + return branch; + } + + private async createDerivedBranch( + sessionId: string, + parent: ConversationBranch, + forkMessageId: string | null | undefined, + name: string + ): Promise { + const child = await this.sessions.createBranch({ + sessionId, + name, + parentBranchId: parent.id, + forkMessageId, + }); + await this.sessions.setActiveBranch(sessionId, child.id); + return child; } private async execute( @@ -195,6 +448,12 @@ export class RunManager { run.answer = answer; } run.completedAt = now; + if (run.manifest) { + run.manifest = { + ...run.manifest, + toolCallIds: toolCalls.map((toolCall) => toolCall.id), + }; + } await this.runs.update(run); @@ -205,11 +464,20 @@ export class RunManager { const isInfraFailure = run.status === 'failed' && isRuntimeInfraCode(run.error?.code); if (!isInfraFailure) { const assistantMessage: Message = { - id: randomUUID(), + // The run id gives the assistant artifact a deterministic stable id + // across the live event projection and a subsequent reload. + id: `assistant-${run.id}`, role: 'assistant', content: answer || (run.status === 'failed' ? run.error?.message ?? 'Run failed.' : ''), timestamp: now, toolCalls: toolCalls.map(toRecord), + branchId: run.branchId, + parentMessageId: run.userMessageId, + runId: run.id, + generationId: run.generationId ?? run.id, + operation: run.operation, + sourceMessageId: run.sourceMessageId, + contextSnapshot: run.contextSnapshot, }; await this.sessions.appendMessage(run.sessionId, assistantMessage); } @@ -281,3 +549,21 @@ function collectSymbols(toolCalls: ToolCall[]): string[] { } return symbols.slice(0, 5); } + +function findUserMessageIndex(messages: Message[], run: Run): number { + if (run.userMessageId) { + const exact = messages.findIndex((message) => message.id === run.userMessageId); + if (exact >= 0) return exact; + } + for (let index = messages.length - 1; index >= 0; index -= 1) { + if (messages[index].role === 'user' && messages[index].content === run.input) return index; + } + return -1; +} + +function findPreviousUserMessage(messages: Message[], beforeIndex: number): Message | undefined { + for (let index = beforeIndex - 1; index >= 0; index -= 1) { + if (messages[index].role === 'user') return messages[index]; + } + return undefined; +} diff --git a/packages/shared/src/kernel/session-manager.ts b/packages/shared/src/kernel/session-manager.ts index 9baff9c..38c1d43 100644 --- a/packages/shared/src/kernel/session-manager.ts +++ b/packages/shared/src/kernel/session-manager.ts @@ -169,7 +169,7 @@ export class SessionManager { sessionId: string; name?: string; parentBranchId?: string; - forkMessageId?: string; + forkMessageId?: string | null; runtimeLeafId?: string; }): Promise { const session = await this.sessions.get(input.sessionId); @@ -183,7 +183,7 @@ export class SessionManager { createdAt: now, updatedAt: now, parentBranchId: input.parentBranchId, - forkMessageId: input.forkMessageId, + forkMessageId: input.parentBranchId ? input.forkMessageId ?? null : undefined, runtimeLeafId: input.runtimeLeafId, }; if (branch.parentBranchId && !(await this.branches.get(input.sessionId, branch.parentBranchId))) { diff --git a/packages/ui/src/atoms/runAtoms.ts b/packages/ui/src/atoms/runAtoms.ts index 9b4ebf2..d0f2a21 100644 --- a/packages/ui/src/atoms/runAtoms.ts +++ b/packages/ui/src/atoms/runAtoms.ts @@ -1,5 +1,5 @@ import { atom } from 'jotai'; -import type { AgentEvent, ApiError, Message, ToolCall, WorkspaceContext } from '@finagent/core'; +import type { AgentEvent, ApiError, Message, Run, ToolCall, WorkspaceContext } from '@finagent/core'; import { isRuntimeInfraCode } from '@finagent/core'; import type { FinagentClient } from '../client'; import { activeSessionIdAtom, messagesAtomFamily, sessionsAtom } from './sessionAtoms'; @@ -8,6 +8,11 @@ import { activeSessionIdAtom, messagesAtomFamily, sessionsAtom } from './session export interface RunView { runId: string; sessionId: string; + branchId?: string; + operation?: Run['operation']; + generationId?: string; + userMessageId?: string; + contextSnapshot?: Run['contextSnapshot']; answer: string; toolCalls: ToolCall[]; error?: ApiError; @@ -56,13 +61,26 @@ export const applyAgentEventAtom = atom( if (event.type === 'run_started') { // Kernel persists the user message; surface it in the UI here. - set(messages, [...get(messages), event.payload.userMessage]); + if (event.payload.userMessageIsNew !== false) { + set(messages, [...get(messages), event.payload.userMessage]); + } set(sessionsAtom, (sessions) => sessions.map((session) => - session.id === sessionId ? { ...session, status: 'running' as const, messageCount: session.messageCount + 1 } : session + session.id === sessionId + ? { + ...session, + status: 'running' as const, + messageCount: session.messageCount + (event.payload.userMessageIsNew === false ? 0 : 1), + } + : session )); set(runViewAtom, { runId: event.runId, sessionId, + branchId: event.payload.run.branchId, + operation: event.payload.run.operation, + generationId: event.payload.run.generationId, + userMessageId: event.payload.run.userMessageId, + contextSnapshot: event.payload.run.contextSnapshot, answer: '', toolCalls: [], infraError: undefined, @@ -121,6 +139,12 @@ export const applyAgentEventAtom = atom( role: 'assistant', content: event.payload.answer, timestamp: event.timestamp, + branchId: run.branchId, + parentMessageId: run.userMessageId, + runId: event.runId, + generationId: run.generationId ?? event.runId, + operation: run.operation, + contextSnapshot: run.contextSnapshot, toolCalls: event.payload.toolCalls.map((toolCall) => ({ id: toolCall.id, toolName: toolCall.toolName, @@ -181,6 +205,12 @@ export const applyAgentEventAtom = atom( ? (run.answer || '(run stopped)') : `Error: ${error.message}`, timestamp: event.timestamp, + branchId: run.branchId, + parentMessageId: run.userMessageId, + runId: event.runId, + generationId: run.generationId ?? event.runId, + operation: run.operation, + contextSnapshot: run.contextSnapshot, toolCalls: run.toolCalls.map((toolCall) => ({ id: toolCall.id, toolName: toolCall.toolName, diff --git a/packages/ui/src/atoms/sessionAtoms.test.ts b/packages/ui/src/atoms/sessionAtoms.test.ts index 82b77c3..a0ac704 100644 --- a/packages/ui/src/atoms/sessionAtoms.test.ts +++ b/packages/ui/src/atoms/sessionAtoms.test.ts @@ -45,8 +45,14 @@ function makeClient(): FinagentClient { ok: true as const, data: (savedMessages[sessionId] ?? []) as never[], }), + listBranches: async () => ({ ok: true as const, data: [] }), + setActiveBranch: async () => ({ ok: false as const, error: { code: 'TEST', message: 'no-op' } }), listRuns: async () => ({ ok: true as const, data: [] as Run[] }), startRun: async () => ({ ok: false as const, error: { code: 'TEST', message: 'no-op' } }), + retryRun: async () => ({ ok: false as const, error: { code: 'TEST', message: 'no-op' } }), + editMessage: async () => ({ ok: false as const, error: { code: 'TEST', message: 'no-op' } }), + regenerateMessage: async () => ({ ok: false as const, error: { code: 'TEST', message: 'no-op' } }), + forkBranch: async () => ({ ok: false as const, error: { code: 'TEST', message: 'no-op' } }), cancelRun: async () => ({ ok: true as const, data: undefined }), onAgentEvent: () => () => undefined, }, diff --git a/packages/ui/src/client.tsx b/packages/ui/src/client.tsx index 52baea9..88f6bcc 100644 --- a/packages/ui/src/client.tsx +++ b/packages/ui/src/client.tsx @@ -5,6 +5,7 @@ import type { AlertTriggerEvent, ApiResult, CalcIndex, + ConversationBranch, Comparison, CredentialInfo, CustomProviderConfig, @@ -167,12 +168,35 @@ export interface FinagentClient { createSession: (title?: string) => Promise>; deleteSession: (sessionId: string) => Promise>; getMessages: (sessionId: string) => Promise>; + listBranches: (sessionId: string) => Promise>; + setActiveBranch: (sessionId: string, branchId: string) => Promise>; listRuns: (sessionId: string) => Promise>; startRun: ( sessionId: string, content: string, workspaceContext?: WorkspaceContext ) => Promise>; + retryRun: ( + sessionId: string, + runId: string, + workspaceContext?: WorkspaceContext + ) => Promise>; + editMessage: ( + sessionId: string, + messageId: string, + content: string, + workspaceContext?: WorkspaceContext + ) => Promise>; + regenerateMessage: ( + sessionId: string, + messageId: string, + workspaceContext?: WorkspaceContext + ) => Promise>; + forkBranch: ( + sessionId: string, + messageId: string, + name?: string + ) => Promise>; cancelRun: (sessionId: string, runId: string) => Promise>; onAgentEvent: (callback: (event: AgentEvent) => void) => () => void; }; @@ -297,8 +321,14 @@ export const fallbackClient: FinagentClient = { createSession: missingClient('kernel.createSession'), deleteSession: missingClient('kernel.deleteSession'), getMessages: missingClient('kernel.getMessages'), + listBranches: missingClient('kernel.listBranches'), + setActiveBranch: missingClient('kernel.setActiveBranch'), listRuns: missingClient('kernel.listRuns'), startRun: missingClient('kernel.startRun'), + retryRun: missingClient('kernel.retryRun'), + editMessage: missingClient('kernel.editMessage'), + regenerateMessage: missingClient('kernel.regenerateMessage'), + forkBranch: missingClient('kernel.forkBranch'), cancelRun: missingClient('kernel.cancelRun'), onAgentEvent: () => () => undefined, }, From f12bfdebe9707664b7b6fb0a36385e8a647f036b Mon Sep 17 00:00:00 2001 From: CC1227871 <2812624878@qq.com> Date: Fri, 11 Sep 2026 15:15:08 +0800 Subject: [PATCH 3/7] feat: expose native Pi conversation cursors --- packages/core/src/index.ts | 50 ++++++++++++ packages/shared/src/agent/pi-rpc-client.ts | 77 +++++++++++++++++++ .../agent/pi-runtime-agent-backend.test.ts | 52 +++++++++++++ 3 files changed, 179 insertions(+) diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index 928e677..9020231 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -165,6 +165,10 @@ export interface ConversationBranch { forkMessageId?: string | null; /** Pi's tree leaf for this branch, when the Pi runtime is in use. */ runtimeLeafId?: string; + /** Runtime session file created for this logical branch. */ + runtimeSessionPath?: string; + /** Runtime entry used to create this branch from its parent. */ + runtimeForkEntryId?: string; } export interface Message { @@ -186,6 +190,8 @@ export interface Message { /** Message/run this operation was derived from (edit, retry, regenerate). */ sourceMessageId?: string; contextSnapshot?: RunContextSnapshot; + /** Stable entry id in the backing runtime's conversation tree. */ + runtimeEntryId?: string; } export type RunStatus = 'running' | 'completed' | 'failed' | 'cancelled'; @@ -209,6 +215,12 @@ export interface Run { /** Explicit generation identity; normally equal to `id`, kept for APIs. */ generationId?: string; contextSnapshot?: RunContextSnapshot; + /** Runtime identities captured for this generation's audit manifest. */ + runtimeSessionId?: string; + runtimeSessionPath?: string; + runtimeLeafId?: string; + runtimeUserEntryId?: string; + runtimeAssistantEntryId?: string; /** Immutable run manifest for audit/evaluation consumers. */ manifest?: RunManifest; } @@ -222,6 +234,11 @@ export interface RunManifest { sourceMessageId?: string; contextSnapshot?: RunContextSnapshot; toolCallIds: string[]; + runtimeSessionId?: string; + runtimeSessionPath?: string; + runtimeLeafId?: string; + runtimeUserEntryId?: string; + runtimeAssistantEntryId?: string; } /** Live tool call state, streamed through agent events. */ @@ -456,11 +473,39 @@ export interface AgentRunInput { sessionId: string; runId: string; content: string; + /** Logical Folio branch that owns this runtime generation. */ + branchId?: string; + /** Runtime session file selected for this branch. */ + sessionPath?: string; workspaceContext?: WorkspaceContext; /** V8: effective UI locale for new agent responses (spec §41–42). */ locale?: SupportedLocale; } +/** Request to materialize a logical branch in a runtime conversation tree. */ +export interface RuntimeBranchPreparationInput { + sessionId: string; + branchId: string; + parentBranchId?: string; + parentSessionPath?: string; + forkMessageId?: string | null; + /** Runtime entry id of the parent message used as the native fork cursor. */ + forkRuntimeEntryId?: string; +} + +/** Runtime identities returned after a branch has been prepared. */ +export interface RuntimeBranchState { + runtimeSessionId?: string; + runtimeSessionPath?: string; + runtimeLeafId?: string; +} + +/** Stable runtime identities captured for one generation. */ +export interface RuntimeRunArtifacts extends RuntimeBranchState { + runtimeUserEntryId?: string; + runtimeAssistantEntryId?: string; +} + /** A model as reported by the Pi model registry. */ export interface LlmModel { provider: string; @@ -552,10 +597,15 @@ export interface AgentRuntime { ensureSession: (session: { id: string; title?: string; + branchId?: string; sessionPath?: string; recentSymbols?: string[]; }) => Promise; run: (input: AgentRunInput) => AsyncIterable; + /** Optional native branch operation (Pi uses its session-tree fork). */ + prepareBranch?: (input: RuntimeBranchPreparationInput) => Promise; + /** Optional runtime identities captured after a generation settles. */ + getRunArtifacts?: (input: { sessionId: string; runId: string }) => Promise; cancel: (input: { sessionId: string; runId: string }) => Promise; disposeSession?: (sessionId: string) => Promise; dispose: () => Promise; diff --git a/packages/shared/src/agent/pi-rpc-client.ts b/packages/shared/src/agent/pi-rpc-client.ts index 7190503..879e2a1 100644 --- a/packages/shared/src/agent/pi-rpc-client.ts +++ b/packages/shared/src/agent/pi-rpc-client.ts @@ -65,6 +65,32 @@ export interface PiState { thinkingLevel?: string; } +export interface PiForkResult { + text?: string; + cancelled: boolean; +} + +export interface PiForkMessage { + entryId: string; + text: string; +} + +/** Append-only entry returned by Pi's session-tree RPC. */ +export interface PiSessionEntry { + type?: string; + id: string; + parentId?: string | null; + message?: { + role?: string; + content?: unknown; + }; +} + +export interface PiEntriesResult { + entries: PiSessionEntry[]; + leafId?: string | null; +} + /** Events yielded by {@link PiRpcClient.promptStreaming}. */ export type PiStreamEvent = | { kind: 'event'; event: Record } @@ -164,6 +190,57 @@ export class PiRpcClient { return this.getState(); } + /** Fork the active Pi conversation before a previous user entry. */ + async fork(entryId: string): Promise { + const data = await this.sendControl({ type: 'fork', entryId }, this.controlTimeoutMs); + const record = readRecord(data); + return { + text: typeof record.text === 'string' ? record.text : undefined, + cancelled: record.cancelled === true, + }; + } + + /** Return user entries that Pi exposes as native fork cursors. */ + async getForkMessages(): Promise { + const data = await this.sendControl({ type: 'get_fork_messages' }, this.controlTimeoutMs); + const record = readRecord(data); + if (!Array.isArray(record.messages)) { + throw createCodeError('PI_PROTOCOL_ERROR', 'Pi get_fork_messages returned no messages list.'); + } + return record.messages.flatMap((message) => { + const item = readRecord(message); + if (typeof item.entryId !== 'string' || typeof item.text !== 'string') return []; + return [{ entryId: item.entryId, text: item.text }]; + }); + } + + /** Read the append-only Pi session tree and its current leaf cursor. */ + async getEntries(since?: string): Promise { + const data = await this.sendControl( + since ? { type: 'get_entries', since } : { type: 'get_entries' }, + this.controlTimeoutMs + ); + const record = readRecord(data); + if (!Array.isArray(record.entries)) { + throw createCodeError('PI_PROTOCOL_ERROR', 'Pi get_entries returned no entries list.'); + } + return { + entries: record.entries.flatMap((entry) => { + const item = readRecord(entry); + if (typeof item.id !== 'string') return []; + return [{ + type: typeof item.type === 'string' ? item.type : undefined, + id: item.id, + parentId: typeof item.parentId === 'string' ? item.parentId : item.parentId === null ? null : undefined, + message: item.message && typeof item.message === 'object' + ? readRecord(item.message) as PiSessionEntry['message'] + : undefined, + }]; + }), + leafId: typeof record.leafId === 'string' ? record.leafId : record.leafId === null ? null : undefined, + }; + } + /** * List models available in the runtime (auth-configured only). * First call can be slow while the runtime warms its model catalog. diff --git a/packages/shared/src/agent/pi-runtime-agent-backend.test.ts b/packages/shared/src/agent/pi-runtime-agent-backend.test.ts index 8649b6e..8ebeea1 100644 --- a/packages/shared/src/agent/pi-runtime-agent-backend.test.ts +++ b/packages/shared/src/agent/pi-runtime-agent-backend.test.ts @@ -142,6 +142,58 @@ describe('PiRpcClient', () => { expect(state).toMatchObject({ sessionId: 'pi-123', sessionFile: '/tmp/s1.jsonl' }); }); + it('exposes Pi native fork cursors and append-only entry identities', async () => { + const client = new PiRpcClient({ + spawnProcess: createSpawn(() => + new FakePiProcess((line, proc) => { + if (line.type === 'fork') { + proc.writeEvent({ + id: line.id, + type: 'response', + command: 'fork', + success: true, + data: { text: 'Original prompt', cancelled: false }, + }); + } + if (line.type === 'get_fork_messages') { + proc.writeEvent({ + id: line.id, + type: 'response', + command: 'get_fork_messages', + success: true, + data: { messages: [{ entryId: 'u-1', text: 'Original prompt' }] }, + }); + } + if (line.type === 'get_entries') { + proc.writeEvent({ + id: line.id, + type: 'response', + command: 'get_entries', + success: true, + data: { + entries: [ + { type: 'message', id: 'u-1', parentId: null, message: { role: 'user', content: 'Original prompt' } }, + { type: 'message', id: 'a-1', parentId: 'u-1', message: { role: 'assistant', content: 'Original answer' } }, + ], + leafId: 'a-1', + }, + }); + } + }) + ), + }); + + await expect(client.fork('u-1')).resolves.toEqual({ text: 'Original prompt', cancelled: false }); + await expect(client.getForkMessages()).resolves.toEqual([{ entryId: 'u-1', text: 'Original prompt' }]); + await expect(client.getEntries()).resolves.toMatchObject({ + leafId: 'a-1', + entries: [ + { id: 'u-1', parentId: null, message: { role: 'user' } }, + { id: 'a-1', parentId: 'u-1', message: { role: 'assistant' } }, + ], + }); + }); + it('streams raw Pi events and settles with the aggregated result', async () => { const client = new PiRpcClient({ spawnProcess: createSpawn(() => From 443145d3272a744f03e8b318c685242f34762c3b Mon Sep 17 00:00:00 2001 From: CC1227871 <2812624878@qq.com> Date: Fri, 11 Sep 2026 15:25:18 +0800 Subject: [PATCH 4/7] feat: isolate native Pi conversation branches --- .../shared/src/agent/pi-runtime-adapter.ts | 202 ++++++++++++++++-- .../agent/pi-runtime-agent-backend.test.ts | 143 ++++++++++++- .../shared/src/kernel/run-manager.test.ts | 46 +++- packages/shared/src/kernel/run-manager.ts | 96 ++++++++- packages/shared/src/kernel/session-manager.ts | 23 ++ 5 files changed, 478 insertions(+), 32 deletions(-) diff --git a/packages/shared/src/agent/pi-runtime-adapter.ts b/packages/shared/src/agent/pi-runtime-adapter.ts index 08c0eff..6b0d5e8 100644 --- a/packages/shared/src/agent/pi-runtime-adapter.ts +++ b/packages/shared/src/agent/pi-runtime-adapter.ts @@ -11,6 +11,9 @@ import type { LlmRuntimeState, LlmTestResult, RuntimeSession, + RuntimeBranchPreparationInput, + RuntimeBranchState, + RuntimeRunArtifacts, SkillReadiness, ToolCall, ToolDefinition, @@ -21,7 +24,13 @@ import type { SkillHub } from '@finagent/skill-hub'; import { FinanceToolRegistry } from './finance-tool-registry.ts'; import { createPhaseOneRegistry } from '../capabilities/index.ts'; import { MarketDataService } from './market-data-service.ts'; -import { PiRpcClient, type PiRpcClientOptions, type PiState } from './pi-rpc-client.ts'; +import { + PiRpcClient, + type PiEntriesResult, + type PiRpcClientOptions, + type PiSessionEntry, + type PiState, +} from './pi-rpc-client.ts'; import { PiEventAdapter } from './pi-event-adapter.ts'; import { createCodeError, toApiError } from './errors.ts'; @@ -63,8 +72,10 @@ export interface LlmRuntimeApi { interface RuntimeSessionState { sessionId: string; + branchId?: string; sessionPath: string; runtimeSessionId?: string; + runtimeLeafId?: string; recentSymbols: string[]; } @@ -86,8 +97,12 @@ export class PiRuntimeAdapter implements AgentRuntime { private readonly readinessProvider?: (skillId: string) => SkillReadiness | undefined; private readonly now: () => number; private readonly sessions = new Map(); + private readonly runStartLeaves = new Map(); + private readonly runArtifacts = new Map(); /** Active session file in the runtime. */ private activePath: string | null = null; + /** Active Folio runtime state key. */ + private activeStateKey: string | null = null; /** Last configured Pi extension list + the Finagent-core-only subset. */ private extensions: string[] = []; private coreExtensions: string[] = []; @@ -137,18 +152,66 @@ export class PiRuntimeAdapter implements AgentRuntime { }; } - async ensureSession(session: { id: string; title?: string; sessionPath?: string }): Promise { + async ensureSession(session: { + id: string; + title?: string; + branchId?: string; + sessionPath?: string; + recentSymbols?: string[]; + }): Promise { const sessionPath = session.sessionPath ?? this.sessionPathFor(session.id); - const state = this.getOrCreateState(session.id, sessionPath); + const state = this.getOrCreateState(session.id, session.branchId, sessionPath); + if (session.recentSymbols && session.recentSymbols.length > 0 && state.recentSymbols.length === 0) { + state.recentSymbols = [...session.recentSymbols]; + } await this.activate(state); return { sessionId: session.id, runtimeSessionId: state.runtimeSessionId, - sessionPath, + sessionPath: state.sessionPath, status: 'active', }; } + /** + * Materialize a Folio child branch in Pi. Pi's `fork` creates a new JSONL + * session file; a branch without a known native cursor gets a fresh empty + * session so historical answers can never leak into the new generation. + */ + async prepareBranch(input: RuntimeBranchPreparationInput): Promise { + if (input.parentSessionPath && this.activePath !== input.parentSessionPath) { + await this.rpcClient.switchSession(input.parentSessionPath); + this.activePath = input.parentSessionPath; + } + + if (input.forkRuntimeEntryId) { + const result = await this.rpcClient.fork(input.forkRuntimeEntryId); + if (result.cancelled) { + throw createCodeError('RUN_CANCELLED', 'Pi cancelled the conversation branch fork.'); + } + } else { + const freshPath = this.branchPathFor(input.sessionId, input.branchId); + await this.rpcClient.switchSession(freshPath); + this.activePath = freshPath; + } + + const state = await this.rpcClient.getState(); + const entries = await this.rpcClient.getEntries().catch(() => undefined); + const runtimeSessionPath = state.sessionFile ?? this.activePath ?? undefined; + this.activePath = runtimeSessionPath ?? null; + return { + runtimeSessionId: state.sessionId, + runtimeSessionPath, + runtimeLeafId: entries?.leafId ?? undefined, + }; + } + + async getRunArtifacts(input: { sessionId: string; runId: string }): Promise { + const artifacts = this.runArtifacts.get(input.runId); + this.runArtifacts.delete(input.runId); + return artifacts; + } + /** * Run one prompt, with one retry at most. If the FIRST attempt dies at * startup with an optional-extension load failure (V8.1 §37), the retry @@ -219,7 +282,11 @@ export class PiRuntimeAdapter implements AgentRuntime { outcome: { error: ApiError | undefined; terminated: boolean }, yieldInfraFailure: boolean ): AsyncIterable { - const state = this.getOrCreateState(input.sessionId, this.sessionPathFor(input.sessionId)); + const state = this.getOrCreateState( + input.sessionId, + input.branchId, + input.sessionPath ?? this.sessionPathFor(input.sessionId) + ); const now = this.now; const fail = async function* (error: unknown, emit: boolean): AsyncIterable { outcome.error = toApiError(error); @@ -237,6 +304,9 @@ export class PiRuntimeAdapter implements AgentRuntime { return; } + const beforeEntries = await this.readEntries(); + this.runStartLeaves.set(input.runId, beforeEntries?.leafId ?? undefined); + const adapter = new PiEventAdapter({ sessionId: input.sessionId, runId: input.runId, now: this.now }); const stream = this.rpcClient.promptStreaming( buildPrompt(input.content, state, input.workspaceContext, this.skillHub, this.readinessProvider, input.locale) @@ -259,6 +329,8 @@ export class PiRuntimeAdapter implements AgentRuntime { runError = error; } + await this.captureRunArtifacts(input, state).catch(() => undefined); + if (runError !== undefined) { const emitTerminal = yieldInfraFailure || !(await this.isOptionalExtensionFailure(runError)); @@ -272,6 +344,49 @@ export class PiRuntimeAdapter implements AgentRuntime { } } + private async captureRunArtifacts( + input: AgentRunInput, + state: RuntimeSessionState + ): Promise { + const since = this.runStartLeaves.get(input.runId); + this.runStartLeaves.delete(input.runId); + + const initialEntries = await this.readEntries(since); + if (!initialEntries) return; + let entries = initialEntries; + + // A Pi restart or a stale cursor can make the incremental query empty; + // the full tree still gives us the durable identities for this run. + if (entries.entries.length === 0) { + entries = (await this.readEntries()) ?? entries; + } + + const messageEntries = entries.entries.filter((entry) => entry.type === 'message' && entry.message); + const userIndex = findLatestMatchingUser(messageEntries, input.content); + const userEntry = userIndex >= 0 ? messageEntries[userIndex] : undefined; + const assistantEntry = userIndex >= 0 + ? [...messageEntries.slice(userIndex + 1)].reverse().find((entry) => entry.message?.role === 'assistant') + : undefined; + + state.runtimeLeafId = entries.leafId ?? state.runtimeLeafId; + const artifacts: RuntimeRunArtifacts = { + runtimeSessionId: state.runtimeSessionId, + runtimeSessionPath: state.sessionPath, + runtimeLeafId: state.runtimeLeafId, + runtimeUserEntryId: userEntry?.id, + runtimeAssistantEntryId: assistantEntry?.id, + }; + this.runArtifacts.set(input.runId, artifacts); + } + + private async readEntries(since?: string): Promise { + const client = this.rpcClient as PiRpcClient & { + getEntries?: (cursor?: string) => Promise; + }; + if (typeof client.getEntries !== 'function') return undefined; + return client.getEntries(since).catch(() => undefined); + } + /** True when the error signature points at a broken optional-extension load. */ private async isOptionalExtensionFailure(error: unknown): Promise { if (this.extensions.length <= this.coreExtensions.length) return false; @@ -298,17 +413,18 @@ export class PiRuntimeAdapter implements AgentRuntime { } async disposeSession(sessionId: string): Promise { - const state = this.sessions.get(sessionId); - this.sessions.delete(sessionId); - // Remove the Pi conversation file together with the Folio session. The - // path is deterministic per session, so this works even when the session - // was created but never ran. - await unlink(join(this.sessionDir, `${sessionId}.jsonl`)).catch(() => undefined); - if (state?.sessionPath && state.sessionPath !== join(this.sessionDir, `${sessionId}.jsonl`)) { - await unlink(state.sessionPath).catch(() => undefined); + const states = Array.from(this.sessions.values()).filter((state) => state.sessionId === sessionId); + for (const [key, state] of this.sessions.entries()) { + if (state.sessionId === sessionId) this.sessions.delete(key); } - if (this.activePath === state?.sessionPath) { + const paths = new Set([ + join(this.sessionDir, `${sessionId}.jsonl`), + ...states.map((state) => state.sessionPath), + ]); + await Promise.all(Array.from(paths, (path) => unlink(path).catch(() => undefined))); + if (states.some((state) => state.sessionPath === this.activePath)) { this.activePath = null; + this.activeStateKey = null; } } @@ -325,22 +441,46 @@ export class PiRuntimeAdapter implements AgentRuntime { return join(this.sessionDir, `${sessionId}.jsonl`); } - private getOrCreateState(sessionId: string, sessionPath: string): RuntimeSessionState { - const existing = this.sessions.get(sessionId); - if (existing) return existing; - const state: RuntimeSessionState = { sessionId, sessionPath, recentSymbols: [] }; - this.sessions.set(sessionId, state); + private branchPathFor(sessionId: string, branchId: string): string { + return join(this.sessionDir, `${sessionId}-${branchId}.jsonl`); + } + + private stateKey(sessionId: string, branchId?: string): string { + return `${sessionId}:${branchId ?? 'main'}`; + } + + private getOrCreateState( + sessionId: string, + branchId: string | undefined, + sessionPath: string + ): RuntimeSessionState { + const key = this.stateKey(sessionId, branchId); + const existing = this.sessions.get(key); + if (existing) { + if (existing.sessionPath !== sessionPath) { + existing.sessionPath = sessionPath; + existing.runtimeSessionId = undefined; + existing.runtimeLeafId = undefined; + } + return existing; + } + const state: RuntimeSessionState = { sessionId, branchId, sessionPath, recentSymbols: [] }; + this.sessions.set(key, state); return state; } /** Ensure the Pi runtime has the session's conversation file loaded. */ private async activate(state: RuntimeSessionState): Promise { if (this.activePath === state.sessionPath && state.runtimeSessionId) { + this.activeStateKey = this.stateKey(state.sessionId, state.branchId); return; } const piState = await this.rpcClient.switchSession(state.sessionPath); state.runtimeSessionId = piState.sessionId; + state.sessionPath = piState.sessionFile ?? state.sessionPath; + state.runtimeLeafId = undefined; this.activePath = state.sessionPath; + this.activeStateKey = this.stateKey(state.sessionId, state.branchId); } private rememberSymbols(event: AgentEvent) { @@ -350,7 +490,10 @@ export class PiRuntimeAdapter implements AgentRuntime { ? toolCall.args.symbol.toUpperCase() : undefined; if (!symbol) return; - const state = this.sessions.get(event.sessionId); + const activeState = this.activeStateKey ? this.sessions.get(this.activeStateKey) : undefined; + const state = activeState?.sessionId === event.sessionId + ? activeState + : Array.from(this.sessions.values()).find((candidate) => candidate.sessionId === event.sessionId); if (!state) return; state.recentSymbols = [ symbol, @@ -468,6 +611,25 @@ function toLlmRuntimeState(piState: PiState): LlmRuntimeState { }; } +function findLatestMatchingUser(entries: PiSessionEntry[], input: string): number { + for (let index = entries.length - 1; index >= 0; index -= 1) { + const entry = entries[index]; + if (entry.message?.role !== 'user') continue; + if (piMessageText(entry.message.content).includes(input)) return index; + } + return -1; +} + +function piMessageText(content: unknown): string { + if (typeof content === 'string') return content; + if (!Array.isArray(content)) return ''; + return content.flatMap((part) => { + if (!part || typeof part !== 'object') return []; + const record = part as Record; + return record.type === 'text' && typeof record.text === 'string' ? [record.text] : []; + }).join(''); +} + function buildPrompt( content: string, state: RuntimeSessionState, diff --git a/packages/shared/src/agent/pi-runtime-agent-backend.test.ts b/packages/shared/src/agent/pi-runtime-agent-backend.test.ts index 8ebeea1..c27d27f 100644 --- a/packages/shared/src/agent/pi-runtime-agent-backend.test.ts +++ b/packages/shared/src/agent/pi-runtime-agent-backend.test.ts @@ -27,6 +27,8 @@ class FakePiProcess extends EventEmitter { killed = false; pid = 1234; private received: string[] = []; + private activeRequestId: string | undefined; + private activeRequestResponded = false; constructor( private readonly handler: (line: Record, proc: FakePiProcess) => void @@ -41,7 +43,20 @@ class FakePiProcess extends EventEmitter { buffer = buffer.slice(newlineIndex + 1); if (line) { this.received.push(line); - this.handler(JSON.parse(line) as Record, this); + const request = JSON.parse(line) as Record; + this.activeRequestId = typeof request.id === 'string' ? request.id : undefined; + this.activeRequestResponded = false; + this.handler(request, this); + if (request.type === 'get_entries' && !this.activeRequestResponded) { + this.writeEvent({ + id: request.id, + type: 'response', + command: 'get_entries', + success: true, + data: { entries: [], leafId: null }, + }); + } + this.activeRequestId = undefined; } newlineIndex = buffer.indexOf('\n'); } @@ -55,6 +70,7 @@ class FakePiProcess extends EventEmitter { } writeEvent(event: Record) { + if (event.id === this.activeRequestId) this.activeRequestResponded = true; this.stdout.write(`${JSON.stringify(event)}\n`); } @@ -581,6 +597,131 @@ describe('PiRuntimeAdapter', () => { expect(switched).toEqual(['/tmp/pi/s1.jsonl', '/tmp/pi/s2.jsonl', '/tmp/pi/s1.jsonl']); }); + it('captures stable runtime entry identities for a generation', async () => { + let entriesCalls = 0; + const client = new PiRpcClient({ + spawnProcess: createSpawn(() => + new FakePiProcess((line, proc) => { + if (line.type === 'switch_session') { + proc.writeEvent({ id: line.id, type: 'response', command: 'switch_session', success: true }); + } + if (line.type === 'get_state') { + proc.writeEvent({ + id: line.id, + type: 'response', + command: 'get_state', + success: true, + data: { sessionId: 'pi-1', sessionFile: '/tmp/s.jsonl' }, + }); + } + if (line.type === 'get_entries') { + entriesCalls += 1; + proc.writeEvent({ + id: line.id, + type: 'response', + command: 'get_entries', + success: true, + data: entriesCalls === 1 + ? { entries: [], leafId: null } + : { + entries: [ + { type: 'message', id: 'u-2', parentId: 'a-1', message: { role: 'user', content: 'wrapper hello' } }, + { type: 'message', id: 'a-2', parentId: 'u-2', message: { role: 'assistant', content: 'Answer' } }, + ], + leafId: 'a-2', + }, + }); + } + if (line.type === 'prompt') { + proc.writeEvent({ id: line.id, type: 'response', command: 'prompt', success: true }); + proc.writeEvent({ + id: line.id, + type: 'agent_end', + messages: [{ role: 'assistant', content: [{ type: 'text', text: 'Answer' }] }], + }); + } + }) + ), + }); + const adapter = new PiRuntimeAdapter({ rpcClient: client, sessionDir: '/tmp/pi' }); + + await adapter.ensureSession({ id: 's1', branchId: 'main', sessionPath: '/tmp/s.jsonl' }); + for await (const _event of adapter.run({ + sessionId: 's1', + branchId: 'main', + sessionPath: '/tmp/s.jsonl', + runId: 'r1', + content: 'hello', + })) { + // Drain the stream so post-run entry capture completes. + } + + await expect(adapter.getRunArtifacts({ sessionId: 's1', runId: 'r1' })).resolves.toMatchObject({ + runtimeSessionId: 'pi-1', + runtimeSessionPath: '/tmp/s.jsonl', + runtimeLeafId: 'a-2', + runtimeUserEntryId: 'u-2', + runtimeAssistantEntryId: 'a-2', + }); + }); + + it('prepares a child branch with Pi fork and records the new session file', async () => { + let stateCalls = 0; + const client = new PiRpcClient({ + spawnProcess: createSpawn(() => + new FakePiProcess((line, proc) => { + if (line.type === 'switch_session') { + proc.writeEvent({ id: line.id, type: 'response', command: 'switch_session', success: true }); + } + if (line.type === 'fork') { + proc.writeEvent({ + id: line.id, + type: 'response', + command: 'fork', + success: true, + data: { text: 'hello', cancelled: false }, + }); + } + if (line.type === 'get_state') { + stateCalls += 1; + proc.writeEvent({ + id: line.id, + type: 'response', + command: 'get_state', + success: true, + data: { + sessionId: `pi-${stateCalls}`, + sessionFile: stateCalls === 1 ? '/tmp/main.jsonl' : '/tmp/forked.jsonl', + }, + }); + } + if (line.type === 'get_entries') { + proc.writeEvent({ + id: line.id, + type: 'response', + command: 'get_entries', + success: true, + data: { entries: [], leafId: 'a-1' }, + }); + } + }) + ), + }); + const adapter = new PiRuntimeAdapter({ rpcClient: client, sessionDir: '/tmp/pi' }); + + await expect(adapter.prepareBranch({ + sessionId: 's1', + branchId: 'b1', + parentBranchId: 'main', + parentSessionPath: '/tmp/main.jsonl', + forkRuntimeEntryId: 'u-1', + })).resolves.toEqual({ + runtimeSessionId: 'pi-2', + runtimeSessionPath: '/tmp/forked.jsonl', + runtimeLeafId: 'a-1', + }); + }); + it('ends a cancelled run with a RUN_CANCELLED failure event', async () => { const client = new PiRpcClient({ spawnProcess: createSpawn(() => diff --git a/packages/shared/src/kernel/run-manager.test.ts b/packages/shared/src/kernel/run-manager.test.ts index 49ac515..696a314 100644 --- a/packages/shared/src/kernel/run-manager.test.ts +++ b/packages/shared/src/kernel/run-manager.test.ts @@ -8,6 +8,9 @@ import type { AgentRunInput, AgentRuntime, ApiResult, + RuntimeBranchPreparationInput, + RuntimeBranchState, + RuntimeRunArtifacts, RuntimeSession, ToolDefinition, } from '@finagent/core'; @@ -31,8 +34,11 @@ afterEach(async () => { await rm(dir, { recursive: true, force: true }); }); -function makeKernel(script: (input: AgentRunInput) => AsyncIterable) { - const runtime = new ScriptedRuntime(script); +function makeKernel( + script: (input: AgentRunInput) => AsyncIterable, + nativeBranch = false +) { + const runtime = new ScriptedRuntime(script, nativeBranch); const store = new JsonFileStore(dir); const sessions = new SessionManager({ sessions: new SessionRepository(store), @@ -52,20 +58,50 @@ function makeKernel(script: (input: AgentRunInput) => AsyncIterable) } class ScriptedRuntime implements AgentRuntime { - ensureSessionCalls: Array<{ id: string; sessionPath?: string }> = []; + ensureSessionCalls: Array<{ id: string; branchId?: string; sessionPath?: string }> = []; + prepareBranchCalls: RuntimeBranchPreparationInput[] = []; cancelCalls: Array<{ sessionId: string; runId: string }> = []; - constructor(private readonly script: (input: AgentRunInput) => AsyncIterable) {} + constructor( + private readonly script: (input: AgentRunInput) => AsyncIterable, + private readonly nativeBranch: boolean + ) {} async getTools(): Promise> { return { ok: true, data: [] }; } - async ensureSession(session: { id: string; title?: string; sessionPath?: string }): Promise { + async ensureSession(session: { + id: string; + title?: string; + branchId?: string; + sessionPath?: string; + }): Promise { this.ensureSessionCalls.push(session); return { sessionId: session.id, status: 'active' }; } + async prepareBranch(input: RuntimeBranchPreparationInput): Promise { + if (!this.nativeBranch) return {}; + this.prepareBranchCalls.push(input); + return { + runtimeSessionPath: `/runtime/${input.branchId}.jsonl`, + runtimeSessionId: `pi-${input.branchId}`, + runtimeLeafId: `leaf-${input.branchId}`, + }; + } + + async getRunArtifacts(input: { sessionId: string; runId: string }): Promise { + if (!this.nativeBranch) return undefined; + return { + runtimeSessionId: `pi-${input.sessionId}`, + runtimeSessionPath: `/runtime/${input.runId}.jsonl`, + runtimeLeafId: `leaf-${input.runId}`, + runtimeUserEntryId: `user-${input.runId}`, + runtimeAssistantEntryId: `assistant-${input.runId}`, + }; + } + async *run(input: AgentRunInput): AsyncIterable { yield* this.script(input); } diff --git a/packages/shared/src/kernel/run-manager.ts b/packages/shared/src/kernel/run-manager.ts index 5b4c128..53f5899 100644 --- a/packages/shared/src/kernel/run-manager.ts +++ b/packages/shared/src/kernel/run-manager.ts @@ -9,6 +9,7 @@ import type { Message, Run, RunContextSnapshot, + RuntimeRunArtifacts, SessionMeta, ToolCall, ToolCallRecord, @@ -128,11 +129,14 @@ export class RunManager { const visible = await this.sessions.listMessages(sessionId, sourceBranch.id); const userIndex = findUserMessageIndex(visible, sourceRun); const sourceUser = userIndex >= 0 ? visible[userIndex] : undefined; + const runtimeForkEntryId = sourceRun.runtimeUserEntryId ?? + (sourceUser ? await this.runtimeForkEntryForMessage(sessionId, sourceUser) : undefined); const child = await this.createDerivedBranch( sessionId, sourceBranch, userIndex > 0 ? visible[userIndex - 1].id : null, - `Retry · ${sourceRun.input.slice(0, 32)}` + `Retry · ${sourceRun.input.slice(0, 32)}`, + runtimeForkEntryId ); const session = await this.requireSession(sessionId); return this.startGeneration({ @@ -173,11 +177,14 @@ export class RunManager { if (!userMessage) { throw createCodeError('MESSAGE_NOT_FOUND', 'The user context for this answer was not found.'); } + const runtimeForkEntryId = sourceRun.runtimeUserEntryId ?? + await this.runtimeForkEntryForMessage(sessionId, userMessage); const child = await this.createDerivedBranch( sessionId, sourceBranch, userMessage.id, - `Regenerate · ${userMessage.content.slice(0, 28)}` + `Regenerate · ${userMessage.content.slice(0, 28)}`, + runtimeForkEntryId ); const session = await this.requireSession(sessionId); return this.startGeneration({ @@ -213,11 +220,13 @@ export class RunManager { if (!sourceMessage || sourceMessage.role !== 'user') { throw createCodeError('MESSAGE_NOT_FOUND', `User message ${userMessageId} was not found.`); } + const runtimeForkEntryId = await this.runtimeForkEntryForMessage(sessionId, sourceMessage); const child = await this.createDerivedBranch( sessionId, sourceBranch, sourceIndex > 0 ? visible[sourceIndex - 1].id : null, - `Edit · ${text.slice(0, 32)}` + `Edit · ${text.slice(0, 32)}`, + runtimeForkEntryId ); const session = await this.requireSession(sessionId); return this.startGeneration({ @@ -241,11 +250,16 @@ export class RunManager { if (!visible.some((message) => message.id === messageId)) { throw createCodeError('MESSAGE_NOT_FOUND', `Message ${messageId} was not found.`); } + const sourceMessage = visible.find((message) => message.id === messageId); + const runtimeForkEntryId = sourceMessage + ? await this.runtimeForkEntryForMessage(sessionId, sourceMessage) + : undefined; const child = await this.createDerivedBranch( sessionId, sourceBranch, messageId, - name?.trim() || `Fork · ${new Date(this.now()).toISOString()}` + name?.trim() || `Fork · ${new Date(this.now()).toISOString()}`, + runtimeForkEntryId ); await this.sessions.setActiveBranch(sessionId, child.id); return child; @@ -371,17 +385,33 @@ export class RunManager { return branch; } + /** Resolve the native Pi user-entry cursor for a Folio message. */ + private async runtimeForkEntryForMessage(sessionId: string, message: Message): Promise { + const run = message.runId ? await this.sessions.getRun(sessionId, message.runId) : null; + if (message.role === 'user') { + return run?.runtimeUserEntryId ?? message.runtimeEntryId; + } + // Pi forks before a user entry. When the UI forks from an assistant + // artifact, use the user generation that produced that artifact. + if (message.role === 'assistant') { + return run?.runtimeUserEntryId; + } + return message.runtimeEntryId; + } + private async createDerivedBranch( sessionId: string, parent: ConversationBranch, forkMessageId: string | null | undefined, - name: string + name: string, + runtimeForkEntryId?: string ): Promise { const child = await this.sessions.createBranch({ sessionId, name, parentBranchId: parent.id, forkMessageId, + runtimeForkEntryId, }); await this.sessions.setActiveBranch(sessionId, child.id); return child; @@ -399,10 +429,33 @@ export class RunManager { let sawTerminal = false; try { + let branch = run.branchId + ? await this.sessions.getBranch(run.sessionId, run.branchId) + : null; + let runtimeSessionPath = branch?.runtimeSessionPath ?? session.runtimeSessionPath; + + if (branch?.parentBranchId && !branch.runtimeSessionPath && this.runtime.prepareBranch) { + const parent = await this.sessions.getBranch(run.sessionId, branch.parentBranchId); + const prepared = await this.runtime.prepareBranch({ + sessionId: run.sessionId, + branchId: branch.id, + parentBranchId: branch.parentBranchId, + parentSessionPath: parent?.runtimeSessionPath ?? session.runtimeSessionPath, + forkMessageId: branch.forkMessageId, + forkRuntimeEntryId: branch.runtimeForkEntryId, + }); + branch = await this.sessions.updateBranch(run.sessionId, branch.id, { + runtimeLeafId: prepared.runtimeLeafId, + runtimeSessionPath: prepared.runtimeSessionPath, + }); + runtimeSessionPath = prepared.runtimeSessionPath ?? runtimeSessionPath; + } + await this.runtime.ensureSession({ id: run.sessionId, title: session.title, - sessionPath: session.runtimeSessionPath, + branchId: run.branchId, + sessionPath: runtimeSessionPath, recentSymbols: session.recentSymbols, }); @@ -410,6 +463,8 @@ export class RunManager { sessionId: run.sessionId, runId: run.id, content: run.input, + branchId: run.branchId, + sessionPath: runtimeSessionPath, workspaceContext, locale, })) { @@ -436,6 +491,16 @@ export class RunManager { const now = this.now(); const cancelled = Boolean(cancelRequested || (failure && failure.code === 'RUN_CANCELLED')); + let runtimeArtifacts: RuntimeRunArtifacts | undefined; + if (this.runtime.getRunArtifacts) { + try { + runtimeArtifacts = await this.runtime.getRunArtifacts({ sessionId: run.sessionId, runId: run.id }); + } catch { + // Runtime diagnostics must never turn an already-settled generation + // into a second failure. + } + } + if (cancelled) { run.status = 'cancelled'; run.answer = answer; @@ -448,10 +513,28 @@ export class RunManager { run.answer = answer; } run.completedAt = now; + if (runtimeArtifacts) { + run.runtimeSessionId = runtimeArtifacts.runtimeSessionId; + run.runtimeSessionPath = runtimeArtifacts.runtimeSessionPath; + run.runtimeLeafId = runtimeArtifacts.runtimeLeafId; + run.runtimeUserEntryId = runtimeArtifacts.runtimeUserEntryId; + run.runtimeAssistantEntryId = runtimeArtifacts.runtimeAssistantEntryId; + if (run.branchId) { + await this.sessions.updateBranch(run.sessionId, run.branchId, { + runtimeLeafId: runtimeArtifacts.runtimeLeafId, + runtimeSessionPath: runtimeArtifacts.runtimeSessionPath, + }); + } + } if (run.manifest) { run.manifest = { ...run.manifest, toolCallIds: toolCalls.map((toolCall) => toolCall.id), + runtimeSessionId: run.runtimeSessionId, + runtimeSessionPath: run.runtimeSessionPath, + runtimeLeafId: run.runtimeLeafId, + runtimeUserEntryId: run.runtimeUserEntryId, + runtimeAssistantEntryId: run.runtimeAssistantEntryId, }; } @@ -478,6 +561,7 @@ export class RunManager { operation: run.operation, sourceMessageId: run.sourceMessageId, contextSnapshot: run.contextSnapshot, + runtimeEntryId: run.runtimeAssistantEntryId, }; await this.sessions.appendMessage(run.sessionId, assistantMessage); } diff --git a/packages/shared/src/kernel/session-manager.ts b/packages/shared/src/kernel/session-manager.ts index 38c1d43..10d2ede 100644 --- a/packages/shared/src/kernel/session-manager.ts +++ b/packages/shared/src/kernel/session-manager.ts @@ -164,6 +164,25 @@ export class SessionManager { return branch; } + /** Persist runtime identity discovered after a native branch operation. */ + async updateBranch( + sessionId: string, + branchId: string, + patch: Partial + ): Promise { + const branch = await this.branches.get(sessionId, branchId); + if (!branch) return null; + const updated: ConversationBranch = { + ...branch, + ...patch, + id: branch.id, + sessionId: branch.sessionId, + updatedAt: this.now(), + }; + await this.branches.update(updated); + return updated; + } + /** Create an immutable child branch from the selected parent cursor. */ async createBranch(input: { sessionId: string; @@ -171,6 +190,8 @@ export class SessionManager { parentBranchId?: string; forkMessageId?: string | null; runtimeLeafId?: string; + runtimeSessionPath?: string; + runtimeForkEntryId?: string; }): Promise { const session = await this.sessions.get(input.sessionId); if (!session) throw new Error(`Session ${input.sessionId} was not found.`); @@ -185,6 +206,8 @@ export class SessionManager { parentBranchId: input.parentBranchId, forkMessageId: input.parentBranchId ? input.forkMessageId ?? null : undefined, runtimeLeafId: input.runtimeLeafId, + runtimeSessionPath: input.runtimeSessionPath, + runtimeForkEntryId: input.runtimeForkEntryId, }; if (branch.parentBranchId && !(await this.branches.get(input.sessionId, branch.parentBranchId))) { throw new Error(`Parent branch ${branch.parentBranchId} was not found.`); From 01f9712b2e5e22b2f650e5206ef89b4a23494d29 Mon Sep 17 00:00:00 2001 From: CC1227871 <2812624878@qq.com> Date: Fri, 11 Sep 2026 15:29:29 +0800 Subject: [PATCH 5/7] feat: add conversation branch switcher --- packages/i18n/src/locales/en-US/agent.ts | 4 ++ packages/i18n/src/locales/zh-CN/agent.ts | 4 ++ packages/ui/src/atoms/sessionAtoms.test.ts | 43 +++++++++++- packages/ui/src/atoms/sessionAtoms.ts | 66 ++++++++++++++++++- .../ui/src/components/agent/AgentPanel.tsx | 2 + .../src/components/agent/BranchSwitcher.tsx | 57 ++++++++++++++++ .../ui/src/components/kernel/KernelBridge.tsx | 8 ++- 7 files changed, 178 insertions(+), 6 deletions(-) create mode 100644 packages/ui/src/components/agent/BranchSwitcher.tsx diff --git a/packages/i18n/src/locales/en-US/agent.ts b/packages/i18n/src/locales/en-US/agent.ts index 9eab4f7..99b202b 100644 --- a/packages/i18n/src/locales/en-US/agent.ts +++ b/packages/i18n/src/locales/en-US/agent.ts @@ -85,6 +85,10 @@ export const agent = { chat: { noMessages: 'No messages yet. Start the conversation!', }, + branch: { + current: 'Branch', + switch: 'Switch conversation branch', + }, suggestions: { title: 'Try asking', research: [ diff --git a/packages/i18n/src/locales/zh-CN/agent.ts b/packages/i18n/src/locales/zh-CN/agent.ts index 0dc845c..7868fa7 100644 --- a/packages/i18n/src/locales/zh-CN/agent.ts +++ b/packages/i18n/src/locales/zh-CN/agent.ts @@ -85,6 +85,10 @@ export const agent = { chat: { noMessages: '还没有消息，开始对话吧！', }, + branch: { + current: '分支', + switch: '切换对话分支', + }, suggestions: { title: '试试问', research: [ diff --git a/packages/ui/src/atoms/sessionAtoms.test.ts b/packages/ui/src/atoms/sessionAtoms.test.ts index a0ac704..3fcca00 100644 --- a/packages/ui/src/atoms/sessionAtoms.test.ts +++ b/packages/ui/src/atoms/sessionAtoms.test.ts @@ -1,20 +1,25 @@ import { beforeEach, describe, expect, it } from 'bun:test'; import { createStore } from 'jotai'; import type { FinagentClient } from '../client'; -import type { Run, SessionMeta } from '@finagent/core'; +import type { ConversationBranch, Run, SessionMeta } from '@finagent/core'; import { activeMessagesAtom, activeSessionIdAtom, + branchesAtomFamily, createSessionAtom, hydrateSessionsAtom, + loadBranchesAtom, loadMessagesAtom, messagesAtomFamily, sessionsAtom, + switchBranchAtom, } from './sessionAtoms.ts'; let sessionCounter = 0; let savedSessions: SessionMeta[] = []; let savedMessages: Record = {}; +let savedBranches: Record = {}; +let savedBranchMessages: Record = {}; function makeSession(title: string): SessionMeta { sessionCounter += 1; @@ -45,8 +50,16 @@ function makeClient(): FinagentClient { ok: true as const, data: (savedMessages[sessionId] ?? []) as never[], }), - listBranches: async () => ({ ok: true as const, data: [] }), - setActiveBranch: async () => ({ ok: false as const, error: { code: 'TEST', message: 'no-op' } }), + listBranches: async (sessionId: string) => ({ ok: true as const, data: savedBranches[sessionId] ?? [] }), + setActiveBranch: async (sessionId: string, branchId: string) => { + const branch = (savedBranches[sessionId] ?? []).find((candidate) => candidate.id === branchId); + if (!branch) return { ok: false as const, error: { code: 'TEST', message: 'branch not found' } }; + savedSessions = savedSessions.map((session) => + session.id === sessionId ? { ...session, activeBranchId: branchId } : session + ); + savedMessages[sessionId] = savedBranchMessages[branchId] ?? []; + return { ok: true as const, data: branch }; + }, listRuns: async () => ({ ok: true as const, data: [] as Run[] }), startRun: async () => ({ ok: false as const, error: { code: 'TEST', message: 'no-op' } }), retryRun: async () => ({ ok: false as const, error: { code: 'TEST', message: 'no-op' } }), @@ -106,6 +119,8 @@ describe('session atoms', () => { sessionCounter = 0; savedSessions = []; savedMessages = {}; + savedBranches = {}; + savedBranchMessages = {}; }); it('hydrates sessions from the kernel and activates the first one', async () => { @@ -175,4 +190,26 @@ describe('session atoms', () => { expect(store.get(messagesAtomFamily('s1'))[0].content).toBe('from A'); expect(store.get(messagesAtomFamily('s2'))[0].content).toBe('from B'); }); + + it('loads and switches branches without mixing the visible transcript', async () => { + const store = createStore(); + savedSessions = [makeSession('Branching')]; + savedBranches.s1 = [ + { id: 'main', sessionId: 's1', name: 'Main', createdAt: 1, updatedAt: 1 }, + { id: 'alt', sessionId: 's1', name: 'Alternative', createdAt: 2, updatedAt: 2, parentBranchId: 'main', forkMessageId: null }, + ]; + savedBranchMessages.main = [{ id: 'main-user', role: 'user', content: 'main question', timestamp: 1 }]; + savedBranchMessages.alt = [{ id: 'alt-user', role: 'user', content: 'alternative question', timestamp: 2 }]; + savedMessages.s1 = savedBranchMessages.main; + const client = makeClient(); + + await store.set(hydrateSessionsAtom, client); + await store.set(loadBranchesAtom, client, 's1'); + await store.set(loadMessagesAtom, client, 's1'); + await store.set(switchBranchAtom, client, 's1', 'alt'); + + expect(store.get(branchesAtomFamily('s1'))).toHaveLength(2); + expect(store.get(sessionsAtom)[0].activeBranchId).toBe('alt'); + expect(store.get(activeMessagesAtom).map((message) => message.content)).toEqual(['alternative question']); + }); }); diff --git a/packages/ui/src/atoms/sessionAtoms.ts b/packages/ui/src/atoms/sessionAtoms.ts index 8a52015..dc3529c 100644 --- a/packages/ui/src/atoms/sessionAtoms.ts +++ b/packages/ui/src/atoms/sessionAtoms.ts @@ -1,6 +1,6 @@ import { atom } from 'jotai'; import { atomFamily } from 'jotai/utils'; -import type { Message, SessionMeta } from '@finagent/core'; +import type { ApiResult, ConversationBranch, Message, SessionMeta } from '@finagent/core'; import type { FinagentClient } from '../client'; // The kernel (main process) is the source of truth for sessions and messages; @@ -19,6 +19,24 @@ export const activeSessionAtom = atom((get) => { // Per-session message cache, loaded lazily from the kernel. export const messagesAtomFamily = atomFamily((_sessionId: string) => atom([])); +// Branch metadata is cached per session so switching sessions does not lose +// the last selected branch or make the renderer reconstruct lineage locally. +export const branchesAtomFamily = atomFamily((_sessionId: string) => atom([])); + +export const activeBranchesAtom = atom((get) => { + const activeId = get(activeSessionIdAtom); + return activeId ? get(branchesAtomFamily(activeId)) : []; +}); + +export const activeBranchAtom = atom((get) => { + const session = get(activeSessionAtom); + const branches = get(activeBranchesAtom); + if (!session) return null; + return branches.find((branch) => branch.id === session.activeBranchId) + ?? branches.find((branch) => !branch.parentBranchId) + ?? null; +}); + export const activeMessagesAtom = atom((get) => { const activeId = get(activeSessionIdAtom); if (!activeId) return []; @@ -59,6 +77,52 @@ export const loadMessagesAtom = atom( } ); +export const loadBranchesAtom = atom( + null, + async (_get, set, client: FinagentClient, sessionId: string) => { + const result = await client.kernel.listBranches(sessionId); + if (!result.ok) return result; + set(branchesAtomFamily(sessionId), result.data); + const active = result.data.find((branch) => branch.id === _get(activeSessionAtom)?.activeBranchId) + ?? result.data.find((branch) => !branch.parentBranchId); + if (active) { + set(sessionsAtom, (sessions) => sessions.map((session) => + session.id === sessionId && session.activeBranchId !== active.id + ? { ...session, activeBranchId: active.id } + : session + )); + } + return result; + } +); + +/** Switch the persisted branch and immediately reload its visible projection. */ +export const switchBranchAtom = atom( + null, + async (_get, set, client: FinagentClient, sessionId: string, branchId: string): Promise> => { + const result = await client.kernel.setActiveBranch(sessionId, branchId); + if (!result.ok) return result; + + set(sessionsAtom, (sessions) => sessions.map((session) => + session.id === sessionId ? { ...session, activeBranchId: branchId } : session + )); + set(branchesAtomFamily(sessionId), (branches) => branches.map((branch) => + branch.id === branchId ? result.data : branch + )); + + const messages = await client.kernel.getMessages(sessionId); + if (messages.ok) { + set(messagesAtomFamily(sessionId), messages.data); + set(loadedSessionIdsAtom, (loaded) => { + const next = new Set(loaded); + next.add(sessionId); + return next; + }); + } + return result; + } +); + export const createSessionAtom = atom( null, async (_get, set, client: FinagentClient, title?: string) => { diff --git a/packages/ui/src/components/agent/AgentPanel.tsx b/packages/ui/src/components/agent/AgentPanel.tsx index 9820dfb..b39b848 100644 --- a/packages/ui/src/components/agent/AgentPanel.tsx +++ b/packages/ui/src/components/agent/AgentPanel.tsx @@ -22,6 +22,7 @@ import { MessageList } from '../chat/MessageList'; import { MarkdownContent } from '../chat/MarkdownContent'; import { ModelSelector } from './ModelSelector'; import { ThinkingSelector } from './ThinkingSelector'; +import { BranchSwitcher } from './BranchSwitcher'; import { ToolActivity } from './ToolActivity'; import { ContextChip } from './ContextChip'; import { TraceInspector } from '../trace/TraceInspector'; @@ -251,6 +252,7 @@ export const AgentPanel: React.FC = () => {
+ {/* Scrollable body: tool activity, structured results, messages, live answer */}
diff --git a/packages/ui/src/components/agent/BranchSwitcher.tsx b/packages/ui/src/components/agent/BranchSwitcher.tsx new file mode 100644 index 0000000..dc7331a --- /dev/null +++ b/packages/ui/src/components/agent/BranchSwitcher.tsx @@ -0,0 +1,57 @@ +import React, { useState } from 'react'; +import { useTranslation } from 'react-i18next'; +import { GitBranch } from 'lucide-react'; +import { useAtomValue, useSetAtom } from 'jotai'; +import { activeBranchesAtom, activeBranchAtom, activeSessionIdAtom, switchBranchAtom } from '../../atoms'; +import { useFinagentClient } from '../../client'; + +interface BranchSwitcherProps { + disabled?: boolean; +} + +/** Compact branch selector kept beside the Copilot context, not in the chat transcript. */ +export const BranchSwitcher: React.FC = ({ disabled = false }) => { + const { t } = useTranslation(); + const client = useFinagentClient(); + const sessionId = useAtomValue(activeSessionIdAtom); + const branches = useAtomValue(activeBranchesAtom); + const activeBranch = useAtomValue(activeBranchAtom); + const switchBranch = useSetAtom(switchBranchAtom); + const [error, setError] = useState(null); + + if (!sessionId || !activeBranch) return null; + + const handleChange = async (event: React.ChangeEvent) => { + const branchId = event.target.value; + if (!branchId || branchId === activeBranch.id || disabled) return; + setError(null); + const result = await switchBranch(client, sessionId, branchId); + if (!result.ok) setError(result.error.message); + }; + + return ( +
+
+ + + +
+ {error &&
{error}
} +
+ ); +}; diff --git a/packages/ui/src/components/kernel/KernelBridge.tsx b/packages/ui/src/components/kernel/KernelBridge.tsx index b3e0667..ec1f3a5 100644 --- a/packages/ui/src/components/kernel/KernelBridge.tsx +++ b/packages/ui/src/components/kernel/KernelBridge.tsx @@ -3,6 +3,7 @@ import { useAtom, useSetAtom } from 'jotai'; import type { FinagentClient } from '../../client'; import { activeSessionIdAtom, + loadBranchesAtom, hydrateSessionsAtom, loadMessagesAtom, loadedSessionIdsAtom, @@ -19,6 +20,7 @@ import { applyAgentEventAtom } from '../../atoms/runAtoms'; export const KernelBridge: React.FC<{ client: FinagentClient }> = ({ client }) => { const hydrate = useSetAtom(hydrateSessionsAtom); const loadMessages = useSetAtom(loadMessagesAtom); + const loadBranches = useSetAtom(loadBranchesAtom); const applyEvent = useSetAtom(applyAgentEventAtom); const [activeSessionId] = useAtom(activeSessionIdAtom); const [loadedSessionIds] = useAtom(loadedSessionIdsAtom); @@ -32,10 +34,12 @@ export const KernelBridge: React.FC<{ client: FinagentClient }> = ({ client }) = }, [client, hydrate, applyEvent]); useEffect(() => { - if (activeSessionId && !loadedSessionIds.has(activeSessionId)) { + if (!activeSessionId) return; + void loadBranches(client, activeSessionId); + if (!loadedSessionIds.has(activeSessionId)) { void loadMessages(client, activeSessionId); } - }, [client, activeSessionId, loadedSessionIds, loadMessages]); + }, [client, activeSessionId, loadedSessionIds, loadBranches, loadMessages]); return null; }; From 1cc55bac44fb5fbcaad75e7633a9d22fe06fcb74 Mon Sep 17 00:00:00 2001 From: CC1227871 <2812624878@qq.com> Date: Fri, 11 Sep 2026 15:34:09 +0800 Subject: [PATCH 6/7] feat: add conversation message actions --- packages/i18n/src/locales/en-US/agent.ts | 8 ++ packages/i18n/src/locales/zh-CN/agent.ts | 8 ++ packages/ui/src/atoms/runAtoms.ts | 24 +++- packages/ui/src/atoms/sessionAtoms.ts | 44 +++++-- .../ui/src/components/agent/AgentPanel.tsx | 82 ++++++++++++- .../ui/src/components/chat/MessageList.tsx | 20 ++- packages/ui/src/components/chat/TurnCard.tsx | 114 +++++++++++++++++- .../ui/src/components/kernel/KernelBridge.tsx | 5 +- 8 files changed, 279 insertions(+), 26 deletions(-) diff --git a/packages/i18n/src/locales/en-US/agent.ts b/packages/i18n/src/locales/en-US/agent.ts index 99b202b..4229d51 100644 --- a/packages/i18n/src/locales/en-US/agent.ts +++ b/packages/i18n/src/locales/en-US/agent.ts @@ -89,6 +89,14 @@ export const agent = { current: 'Branch', switch: 'Switch conversation branch', }, + actions: { + edit: 'Edit', + regenerate: 'Regenerate', + retry: 'Retry', + fork: 'Fork', + save: 'Save and run', + cancel: 'Cancel', + }, suggestions: { title: 'Try asking', research: [ diff --git a/packages/i18n/src/locales/zh-CN/agent.ts b/packages/i18n/src/locales/zh-CN/agent.ts index 7868fa7..2dcc4db 100644 --- a/packages/i18n/src/locales/zh-CN/agent.ts +++ b/packages/i18n/src/locales/zh-CN/agent.ts @@ -89,6 +89,14 @@ export const agent = { current: '分支', switch: '切换对话分支', }, + actions: { + edit: '编辑', + regenerate: '重新生成', + retry: '重试', + fork: '创建分支', + save: '保存并运行', + cancel: '取消', + }, suggestions: { title: '试试问', research: [ diff --git a/packages/ui/src/atoms/runAtoms.ts b/packages/ui/src/atoms/runAtoms.ts index d0f2a21..04c4c6f 100644 --- a/packages/ui/src/atoms/runAtoms.ts +++ b/packages/ui/src/atoms/runAtoms.ts @@ -2,7 +2,7 @@ import { atom } from 'jotai'; import type { AgentEvent, ApiError, Message, Run, ToolCall, WorkspaceContext } from '@finagent/core'; import { isRuntimeInfraCode } from '@finagent/core'; import type { FinagentClient } from '../client'; -import { activeSessionIdAtom, messagesAtomFamily, sessionsAtom } from './sessionAtoms'; +import { activeSessionIdAtom, messagesAtomFamily, runsAtomFamily, sessionsAtom } from './sessionAtoms'; /** Live view of the currently executing run, streamed from kernel events. */ export interface RunView { @@ -61,9 +61,13 @@ export const applyAgentEventAtom = atom( if (event.type === 'run_started') { // Kernel persists the user message; surface it in the UI here. - if (event.payload.userMessageIsNew !== false) { + if (event.payload.userMessageIsNew !== false && !get(messages).some((message) => message.id === event.payload.userMessage.id)) { set(messages, [...get(messages), event.payload.userMessage]); } + set(runsAtomFamily(sessionId), (runs) => [ + event.payload.run, + ...runs.filter((candidate) => candidate.id !== event.payload.run.id), + ]); set(sessionsAtom, (sessions) => sessions.map((session) => session.id === sessionId ? { @@ -157,6 +161,11 @@ export const applyAgentEventAtom = atom( })), }; set(messages, [...get(messages), assistantMessage]); + set(runsAtomFamily(sessionId), (runs) => runs.map((candidate) => + candidate.id === event.runId + ? { ...candidate, status: 'completed' as const, completedAt: event.timestamp, answer: event.payload.answer, error: undefined } + : candidate + )); set(sessionsAtom, (sessions) => sessions.map((session) => session.id === sessionId ? { ...session, status: 'idle' as const, messageCount: session.messageCount + 1 } : session )); @@ -184,6 +193,17 @@ export const applyAgentEventAtom = atom( // renders a dedicated runtime banner with Retry + Diagnostics. Real // failures (tool errors, task failures) keep the existing message flow. const error = event.payload.error; + set(runsAtomFamily(sessionId), (runs) => runs.map((candidate) => + candidate.id === event.runId + ? { + ...candidate, + status: cancelled ? 'cancelled' as const : 'failed' as const, + completedAt: event.timestamp, + error, + answer: run.answer, + } + : candidate + )); if (isRuntimeInfraCode(error.code)) { set(lastRunSummaryAtom, (previous) => ({ runId: event.runId, diff --git a/packages/ui/src/atoms/sessionAtoms.ts b/packages/ui/src/atoms/sessionAtoms.ts index dc3529c..76b1959 100644 --- a/packages/ui/src/atoms/sessionAtoms.ts +++ b/packages/ui/src/atoms/sessionAtoms.ts @@ -1,6 +1,6 @@ import { atom } from 'jotai'; import { atomFamily } from 'jotai/utils'; -import type { ApiResult, ConversationBranch, Message, SessionMeta } from '@finagent/core'; +import type { ApiResult, ConversationBranch, Message, Run, SessionMeta } from '@finagent/core'; import type { FinagentClient } from '../client'; // The kernel (main process) is the source of truth for sessions and messages; @@ -22,12 +22,18 @@ export const messagesAtomFamily = atomFamily((_sessionId: string) => atom atom([])); +export const runsAtomFamily = atomFamily((_sessionId: string) => atom([])); export const activeBranchesAtom = atom((get) => { const activeId = get(activeSessionIdAtom); return activeId ? get(branchesAtomFamily(activeId)) : []; }); +export const activeRunsAtom = atom((get) => { + const activeId = get(activeSessionIdAtom); + return activeId ? get(runsAtomFamily(activeId)) : []; +}); + export const activeBranchAtom = atom((get) => { const session = get(activeSessionAtom); const branches = get(activeBranchesAtom); @@ -96,20 +102,28 @@ export const loadBranchesAtom = atom( } ); -/** Switch the persisted branch and immediately reload its visible projection. */ -export const switchBranchAtom = atom( +export const loadRunsAtom = atom( + null, + async (_get, set, client: FinagentClient, sessionId: string) => { + const result = await client.kernel.listRuns(sessionId); + if (result.ok) set(runsAtomFamily(sessionId), result.data); + return result; + } +); + +/** Refresh the local projection after an operation has already activated a branch in the kernel. */ +export const refreshBranchProjectionAtom = atom( null, async (_get, set, client: FinagentClient, sessionId: string, branchId: string): Promise> => { - const result = await client.kernel.setActiveBranch(sessionId, branchId); - if (!result.ok) return result; + const branches = await client.kernel.listBranches(sessionId); + if (!branches.ok) return branches; + const branch = branches.data.find((candidate) => candidate.id === branchId); + if (!branch) return { ok: false, error: { code: 'BRANCH_NOT_FOUND', message: `Branch ${branchId} was not found.` } }; + set(branchesAtomFamily(sessionId), branches.data); set(sessionsAtom, (sessions) => sessions.map((session) => session.id === sessionId ? { ...session, activeBranchId: branchId } : session )); - set(branchesAtomFamily(sessionId), (branches) => branches.map((branch) => - branch.id === branchId ? result.data : branch - )); - const messages = await client.kernel.getMessages(sessionId); if (messages.ok) { set(messagesAtomFamily(sessionId), messages.data); @@ -119,7 +133,17 @@ export const switchBranchAtom = atom( return next; }); } - return result; + return { ok: true, data: branch }; + } +); + +/** Switch the persisted branch and immediately reload its visible projection. */ +export const switchBranchAtom = atom( + null, + async (_get, set, client: FinagentClient, sessionId: string, branchId: string): Promise> => { + const result = await client.kernel.setActiveBranch(sessionId, branchId); + if (!result.ok) return result; + return set(refreshBranchProjectionAtom, client, sessionId, result.data.id); } ); diff --git a/packages/ui/src/components/agent/AgentPanel.tsx b/packages/ui/src/components/agent/AgentPanel.tsx index b39b848..8b6ab5e 100644 --- a/packages/ui/src/components/agent/AgentPanel.tsx +++ b/packages/ui/src/components/agent/AgentPanel.tsx @@ -2,9 +2,10 @@ import React, { useEffect, useRef, useState } from 'react'; import { useTranslation } from 'react-i18next'; import { ArrowUp, ChevronLeft, ChevronRight, Square, Sparkles } from 'lucide-react'; import { useAtom, useAtomValue, useSetAtom } from 'jotai'; -import type { ApiError, FolioTrace, PortfolioSnapshot, Quote, ToolCall } from '@finagent/core'; +import type { ApiError, FolioTrace, Message, PortfolioSnapshot, Quote, Run, ToolCall } from '@finagent/core'; import { activeMessagesAtom, + activeRunsAtom, activeSessionIdAtom, agentPanelVisibleAtom, cancelRunAtom, @@ -12,6 +13,7 @@ import { lastRunSummaryAtom, navSectionAtom, runViewAtom, + refreshBranchProjectionAtom, settingsTabAtom, workspaceContextAtom, type LastRunSummary, @@ -90,11 +92,13 @@ export const AgentPanel: React.FC = () => { const { t } = useTranslation(); const client = useFinagentClient(); const [messages] = useAtom(activeMessagesAtom); + const [runs] = useAtom(activeRunsAtom); const [activeSessionId] = useAtom(activeSessionIdAtom); const [runView] = useAtom(runViewAtom); const setAgentPanelVisible = useSetAtom(agentPanelVisibleAtom); const createSession = useSetAtom(createSessionAtom); const cancelRun = useSetAtom(cancelRunAtom); + const refreshBranchProjection = useSetAtom(refreshBranchProjectionAtom); const [lastRun, setLastRun] = useAtom(lastRunSummaryAtom); const workspaceContext = useAtomValue(workspaceContextAtom); @@ -174,13 +178,73 @@ export const AgentPanel: React.FC = () => { /** V8.1 §38: retry the last user message after an infra failure. */ const handleRetry = async () => { - const lastUser = [...messages].reverse().find((message) => message.role === 'user'); - if (!lastUser || !activeSessionId) return; + if (!lastRun || !activeSessionId || (lastRun.status !== 'failed' && lastRun.status !== 'cancelled')) return; setSendError(null); - const result = await client.kernel.startRun(activeSessionId, lastUser.content, workspaceContext); + const result = await client.kernel.retryRun(activeSessionId, lastRun.runId, workspaceContext); if (!result.ok) { setSendError(result.error.message); + return; } + await activateRunBranch(result.data); + }; + + const activateRunBranch = async (run: Run): Promise => { + if (!activeSessionId || !run.branchId) return true; + const result = await refreshBranchProjection(client, activeSessionId, run.branchId); + if (!result.ok) { + setSendError(result.error.message); + return false; + } + return true; + }; + + const handleRetryMessage = async (message: Message): Promise => { + if (!activeSessionId || !message.runId || isRunning) return false; + setSendError(null); + const result = await client.kernel.retryRun(activeSessionId, message.runId, workspaceContext); + if (!result.ok) { + setSendError(result.error.message); + return false; + } + return activateRunBranch(result.data); + }; + + const handleRegenerate = async (message: Message): Promise => { + if (!activeSessionId || isRunning) return false; + setSendError(null); + const result = await client.kernel.regenerateMessage(activeSessionId, message.id, workspaceContext); + if (!result.ok) { + setSendError(result.error.message); + return false; + } + return activateRunBranch(result.data); + }; + + const handleEdit = async (message: Message, content: string): Promise => { + if (!activeSessionId || isRunning) return false; + setSendError(null); + const result = await client.kernel.editMessage(activeSessionId, message.id, content, workspaceContext); + if (!result.ok) { + setSendError(result.error.message); + return false; + } + return activateRunBranch(result.data); + }; + + const handleFork = async (message: Message): Promise => { + if (!activeSessionId || isRunning) return false; + setSendError(null); + const result = await client.kernel.forkBranch(activeSessionId, message.id); + if (!result.ok) { + setSendError(result.error.message); + return false; + } + const projection = await refreshBranchProjection(client, activeSessionId, result.data.id); + if (!projection.ok) { + setSendError(projection.error.message); + return false; + } + return true; }; const handleStop = async () => { @@ -287,7 +351,15 @@ export const AgentPanel: React.FC = () => { }} /> )} - + {isRunning && } Promise; + onRegenerate?: (message: Message) => Promise; + onRetry?: (message: Message) => Promise; + onFork?: (message: Message) => Promise; } -export const MessageList: React.FC = ({ messages, isLoading }) => { +export const MessageList: React.FC = ({ messages, isLoading, runs = [], onEdit, onRegenerate, onRetry, onFork }) => { const { t } = useTranslation(); + const runById = new Map(runs.map((run) => [run.id, run])); return (
{/* Inside the 400px Copilot panel the shell already pads the body; no max-width rail. */}
{messages.map((message) => ( - + ))} {isLoading && (
diff --git a/packages/ui/src/components/chat/TurnCard.tsx b/packages/ui/src/components/chat/TurnCard.tsx index 3b31fdd..fede99e 100644 --- a/packages/ui/src/components/chat/TurnCard.tsx +++ b/packages/ui/src/components/chat/TurnCard.tsx @@ -1,21 +1,52 @@ -import React from 'react'; +import React, { useState } from 'react'; import { useTranslation } from 'react-i18next'; -import type { Message } from '@finagent/core'; +import type { Message, Run } from '@finagent/core'; import { MarkdownContent } from './MarkdownContent'; interface TurnCardProps { message: Message; + run?: Run; + onEdit?: (message: Message, content: string) => Promise; + onRegenerate?: (message: Message) => Promise; + onRetry?: (message: Message) => Promise; + onFork?: (message: Message) => Promise; } -export const TurnCard: React.FC = ({ message }) => { +export const TurnCard: React.FC = ({ message, run, onEdit, onRegenerate, onRetry, onFork }) => { const { t } = useTranslation(); const isUser = message.role === 'user'; const isTool = message.role === 'tool'; const toolCalls = message.toolCalls ?? []; + const [editing, setEditing] = useState(false); + const [draft, setDraft] = useState(message.content); + const [busy, setBusy] = useState(false); + + const invoke = async (action: (() => Promise) | undefined) => { + if (!action || busy) return; + setBusy(true); + try { + await action(); + } finally { + setBusy(false); + } + }; + + const saveEdit = async () => { + const content = draft.trim(); + if (!content || !onEdit) return; + await invoke(async () => { + const ok = await onEdit(message, content); + if (ok) setEditing(false); + return ok; + }); + }; + + const canRetry = run?.status === 'failed' || run?.status === 'cancelled'; + const canRegenerate = message.role === 'assistant' && run?.status === 'completed'; return (
= ({ message }) => { {t('agent.tool.label', { name: message.toolName })}
)} - {isUser ? ( + {editing ? ( +
+