From 92795f7bfbca77d95722639936fbd8852de232b8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Fri, 17 Apr 2026 15:42:13 -0400 Subject: [PATCH] refactor(session): scope pi sessions to agents instead of tasks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Sessions were keyed by (task, agent). Every task switch started a fresh transcript — the agent lost everything it had learned on the previous task. That breaks the product premise: a CEO should remember its past decisions, an engineer should carry context between bugs. Move pi's AgentMessage transcript to be agent-scoped. One session per agent (per company). Tasks enter the ongoing conversation as user turns; rotation/compaction boundaries follow the agent's long-running session. Additive schema only — no row rewrites on the live production DB. - Migration v34: `ALTER TABLE agents ADD COLUMN session_id TEXT` + `idx_agents_session_id`. `ensureOptionalColumns` mirror for fresh DBs. - `AgentRecord.sessionId?: string | null` in the read model; included in every full-agent SELECT (entity-crud, snapshot, agents-hierarchy chain-of-command / direct-reports / heartbeat-eligible). - `db.ensureAgentSession({ agentId, companyId, adapterType })` — idempotent: returns existing `agents.session_id` if set, else inserts a new `task_sessions` row with an empty `task_id` sentinel (agent- scoped, task-id is now metadata) and stores the id on the agent. - `db.setAgentSessionId(agentId, sessionId)` — used by rotation. - `db.updateTasksActiveSessionForAgent(agentId, sessionId)` — mirrors the agent's current session onto every task the agent owns, so existing per-task readers (`tasks.active_session_id`) still work. - `turn-context.resolveSession`: pulls the agent's session via `ensureAgentSession` instead of `getActiveTaskSession(task.id)`, then points the current task at it so UI continues to resolve. - `session-management-service.rotateSession`: new session carries `taskId:""`, then `setAgentSessionId` + `updateTasksActiveSessionForAgent` so every task the agent owns follows the rotation boundary. - `session-entries.listActive`: `ORDER BY created_at, sequence` — the created_at timestamp is authoritative; sequence stays as same-ms tiebreaker. Forward-safety for future merge scenarios. Existing agents' first post-upgrade run creates a fresh agent-scoped session. Their pre-migration per-task transcripts remain in the DB as historical records but are not loaded for new runs. This is a one-time transcript reset, accepted over a risky lossy merge on the live DB. Documented in a code comment on `ensureAgentSession`. --- src/main/database-base.ts | 3 ++ src/main/database/agents-hierarchy.ts | 8 ++-- src/main/database/entity-crud.ts | 4 +- src/main/database/schema.ts | 12 +++++ src/main/database/session-entries.ts | 2 +- src/main/database/snapshot.ts | 2 +- src/main/database/tasks-crud.ts | 65 ++++++++++++++++++++++++++ src/main/orchestration/turn-context.ts | 19 ++++++-- src/main/session-management-service.ts | 19 ++++++-- src/shared/types/records-core.ts | 9 ++++ 10 files changed, 127 insertions(+), 16 deletions(-) diff --git a/src/main/database-base.ts b/src/main/database-base.ts index 4f8fc7bb..facfb82f 100644 --- a/src/main/database-base.ts +++ b/src/main/database-base.ts @@ -260,6 +260,9 @@ export class AppDatabaseBase { incrementSessionUsage(sessionId: string, inputTokens: number, outputTokens: number, runId: string) { return this.tasks.incrementSessionUsage(sessionId, inputTokens, outputTokens, runId); } archiveSession(sessionId: string, reason: string, handoffMarkdown: string) { return this.tasks.archiveSession(sessionId, reason, handoffMarkdown); } createNewSession(input: { taskId: string; companyId: string; agentId: string; adapterType: string; sessionParamsJson: string | null; sessionDisplayId: string | null }) { return this.tasks.createNewSession(input); } + ensureAgentSession(input: { agentId: string; companyId: string; adapterType: string }) { return this.tasks.ensureAgentSession(input); } + setAgentSessionId(agentId: string, sessionId: string) { return this.tasks.setAgentSessionId(agentId, sessionId); } + updateTasksActiveSessionForAgent(agentId: string, sessionId: string) { return this.tasks.updateTasksActiveSessionForAgent(agentId, sessionId); } getRecentRunsForTask(taskId: string, limit: number) { return this.tasks.getRecentRunsForTask(taskId, limit); } getRecentActivityForTask(taskId: string, limit: number) { return this.tasks.getRecentActivityForTask(taskId, limit); } checkoutTask(taskId: string, runId: string) { return this.tasks.checkoutTask(taskId, runId); } diff --git a/src/main/database/agents-hierarchy.ts b/src/main/database/agents-hierarchy.ts index 0e2ce7ab..fea27214 100644 --- a/src/main/database/agents-hierarchy.ts +++ b/src/main/database/agents-hierarchy.ts @@ -66,7 +66,7 @@ export function buildAgentHierarchyMethods( model, capabilities, role_level as roleLevel, budget_monthly_usd as budgetMonthlyUsd, spent_monthly_usd as spentMonthlyUsd, metadata_json as metadataJson, heartbeat_enabled as heartbeatEnabled, heartbeat_interval_sec as heartbeatIntervalSec, - last_heartbeat_at as lastHeartbeatAt, created_at as createdAt, updated_at as updatedAt, + last_heartbeat_at as lastHeartbeatAt, session_id as sessionId, created_at as createdAt, updated_at as updatedAt, 0 as depth FROM agents WHERE id = ? UNION ALL @@ -75,7 +75,7 @@ export function buildAgentHierarchyMethods( a.model, a.capabilities, a.role_level, a.budget_monthly_usd, a.spent_monthly_usd, a.metadata_json, a.heartbeat_enabled, a.heartbeat_interval_sec, - a.last_heartbeat_at, a.created_at, a.updated_at, + a.last_heartbeat_at, a.session_id, a.created_at, a.updated_at, c.depth + 1 FROM agents a JOIN chain c ON a.id = c.reportsTo WHERE c.depth < 20 @@ -92,7 +92,7 @@ export function buildAgentHierarchyMethods( getDirectReports(agentId: string, companyId: string): AgentRecord[] { return readAll( - "select id, company_id as companyId, name, role, title, department, status, reports_to as reportsTo, connector_id as connectorId, workspace_id as workspaceId, model, capabilities, role_level as roleLevel, budget_monthly_usd as budgetMonthlyUsd, spent_monthly_usd as spentMonthlyUsd, metadata_json as metadataJson, heartbeat_enabled as heartbeatEnabled, heartbeat_interval_sec as heartbeatIntervalSec, last_heartbeat_at as lastHeartbeatAt, created_at as createdAt, updated_at as updatedAt from agents where company_id = ? and reports_to = ?", + "select id, company_id as companyId, name, role, title, department, status, reports_to as reportsTo, connector_id as connectorId, workspace_id as workspaceId, model, capabilities, role_level as roleLevel, budget_monthly_usd as budgetMonthlyUsd, spent_monthly_usd as spentMonthlyUsd, metadata_json as metadataJson, heartbeat_enabled as heartbeatEnabled, heartbeat_interval_sec as heartbeatIntervalSec, last_heartbeat_at as lastHeartbeatAt, session_id as sessionId, created_at as createdAt, updated_at as updatedAt from agents where company_id = ? and reports_to = ?", companyId, agentId, ).map((record) => ({ @@ -201,7 +201,7 @@ export function buildAgentHierarchyMethods( listHeartbeatEligibleAgents(companyId: string): AgentRecord[] { return readAll( - "select id, company_id as companyId, name, role, title, department, status, reports_to as reportsTo, connector_id as connectorId, workspace_id as workspaceId, model, capabilities, role_level as roleLevel, budget_monthly_usd as budgetMonthlyUsd, spent_monthly_usd as spentMonthlyUsd, metadata_json as metadataJson, heartbeat_enabled as heartbeatEnabled, heartbeat_interval_sec as heartbeatIntervalSec, last_heartbeat_at as lastHeartbeatAt, created_at as createdAt, updated_at as updatedAt from agents where company_id = ? and heartbeat_enabled = 1 and status not in ('terminated', 'paused', 'error', 'pending_approval')", + "select id, company_id as companyId, name, role, title, department, status, reports_to as reportsTo, connector_id as connectorId, workspace_id as workspaceId, model, capabilities, role_level as roleLevel, budget_monthly_usd as budgetMonthlyUsd, spent_monthly_usd as spentMonthlyUsd, metadata_json as metadataJson, heartbeat_enabled as heartbeatEnabled, heartbeat_interval_sec as heartbeatIntervalSec, last_heartbeat_at as lastHeartbeatAt, session_id as sessionId, created_at as createdAt, updated_at as updatedAt from agents where company_id = ? and heartbeat_enabled = 1 and status not in ('terminated', 'paused', 'error', 'pending_approval')", companyId, ).map((record) => ({ ...record, diff --git a/src/main/database/entity-crud.ts b/src/main/database/entity-crud.ts index 1c9f2bf2..471dcb35 100644 --- a/src/main/database/entity-crud.ts +++ b/src/main/database/entity-crud.ts @@ -64,7 +64,7 @@ export function buildEntityCrudMethods( getAgent(id: string): AgentRecord { const row = getDb() .prepare( - "select id, company_id as companyId, name, role, title, department, status, reports_to as reportsTo, connector_id as connectorId, workspace_id as workspaceId, model, capabilities, role_level as roleLevel, budget_monthly_usd as budgetMonthlyUsd, spent_monthly_usd as spentMonthlyUsd, metadata_json as metadataJson, heartbeat_enabled as heartbeatEnabled, heartbeat_interval_sec as heartbeatIntervalSec, last_heartbeat_at as lastHeartbeatAt, created_at as createdAt, updated_at as updatedAt from agents where id = ?", + "select id, company_id as companyId, name, role, title, department, status, reports_to as reportsTo, connector_id as connectorId, workspace_id as workspaceId, model, capabilities, role_level as roleLevel, budget_monthly_usd as budgetMonthlyUsd, spent_monthly_usd as spentMonthlyUsd, metadata_json as metadataJson, heartbeat_enabled as heartbeatEnabled, heartbeat_interval_sec as heartbeatIntervalSec, last_heartbeat_at as lastHeartbeatAt, session_id as sessionId, created_at as createdAt, updated_at as updatedAt from agents where id = ?", ) .get(id) as AgentRecord | undefined; if (!row) throw new Error(`Agent not found: ${id}`); @@ -75,7 +75,7 @@ export function buildEntityCrudMethods( findAgent(id: string): AgentRecord | null { const row = getDb() .prepare( - "select id, company_id as companyId, name, role, title, department, status, reports_to as reportsTo, connector_id as connectorId, workspace_id as workspaceId, model, capabilities, role_level as roleLevel, budget_monthly_usd as budgetMonthlyUsd, spent_monthly_usd as spentMonthlyUsd, metadata_json as metadataJson, heartbeat_enabled as heartbeatEnabled, heartbeat_interval_sec as heartbeatIntervalSec, last_heartbeat_at as lastHeartbeatAt, created_at as createdAt, updated_at as updatedAt from agents where id = ?", + "select id, company_id as companyId, name, role, title, department, status, reports_to as reportsTo, connector_id as connectorId, workspace_id as workspaceId, model, capabilities, role_level as roleLevel, budget_monthly_usd as budgetMonthlyUsd, spent_monthly_usd as spentMonthlyUsd, metadata_json as metadataJson, heartbeat_enabled as heartbeatEnabled, heartbeat_interval_sec as heartbeatIntervalSec, last_heartbeat_at as lastHeartbeatAt, session_id as sessionId, created_at as createdAt, updated_at as updatedAt from agents where id = ?", ) .get(id) as AgentRecord | undefined; if (!row) return null; diff --git a/src/main/database/schema.ts b/src/main/database/schema.ts index 134b3fb9..82569e8c 100644 --- a/src/main/database/schema.ts +++ b/src/main/database/schema.ts @@ -30,6 +30,8 @@ export function ensureOptionalColumns(db: DatabaseSync) { addColumnIfMissing("agents", "heartbeat_enabled", "alter table agents add column heartbeat_enabled integer not null default 0"); addColumnIfMissing("agents", "heartbeat_interval_sec", "alter table agents add column heartbeat_interval_sec integer not null default 120"); addColumnIfMissing("agents", "last_heartbeat_at", "alter table agents add column last_heartbeat_at text"); + addColumnIfMissing("agents", "session_id", "alter table agents add column session_id text"); + db.exec("create index if not exists idx_agents_session_id on agents(session_id)"); addColumnIfMissing("companies", "auto_approve_hires", "alter table companies add column auto_approve_hires integer not null default 0"); addColumnIfMissing("companies", "review_deliverables", "alter table companies add column review_deliverables integer not null default 0"); addColumnIfMissing("companies", "execution_mode", "alter table companies add column execution_mode text not null default 'ai_ceo'"); @@ -501,6 +503,16 @@ export function runMigrations( ); CREATE INDEX IF NOT EXISTS idx_credentials_status ON credentials(status)`, }, + { + // Per-agent session scope. Agents carry a single active session across + // tasks — switching tasks becomes a "now work on X" user turn on the + // agent's ongoing transcript. Additive-only: existing task_sessions + // rows remain as historical records; new runs use the agent-scoped + // session created lazily on first post-upgrade run. + version: 34, name: "add_agents_session_id", + sql: `ALTER TABLE agents ADD COLUMN session_id TEXT; + CREATE INDEX IF NOT EXISTS idx_agents_session_id ON agents(session_id)`, + }, ]; for (const migration of migrations) { diff --git a/src/main/database/session-entries.ts b/src/main/database/session-entries.ts index ae27f3be..40b5415d 100644 --- a/src/main/database/session-entries.ts +++ b/src/main/database/session-entries.ts @@ -82,7 +82,7 @@ export function buildSessionEntryMethods( is_superseded: number; created_at: string; }>( - "select id, task_session_id, sequence, kind, role, payload_json, is_superseded, created_at from session_entries where task_session_id = ? and is_superseded = 0 order by sequence asc", + "select id, task_session_id, sequence, kind, role, payload_json, is_superseded, created_at from session_entries where task_session_id = ? and is_superseded = 0 order by created_at asc, sequence asc", taskSessionId, ); return rows.map((row) => ({ diff --git a/src/main/database/snapshot.ts b/src/main/database/snapshot.ts index c6012985..7a94aee7 100644 --- a/src/main/database/snapshot.ts +++ b/src/main/database/snapshot.ts @@ -63,7 +63,7 @@ export function buildSnapshotMethods( ...companyArgs, ).map((record) => ({ ...record, isPrimary: toBoolean(record.isPrimary) })) as WorkspaceRecord[], agents: readAll( - `select id, company_id as companyId, name, role, title, department, status, reports_to as reportsTo, connector_id as connectorId, workspace_id as workspaceId, model, capabilities, role_level as roleLevel, budget_monthly_usd as budgetMonthlyUsd, spent_monthly_usd as spentMonthlyUsd, metadata_json as metadataJson, heartbeat_enabled as heartbeatEnabled, heartbeat_interval_sec as heartbeatIntervalSec, last_heartbeat_at as lastHeartbeatAt, created_at as createdAt, updated_at as updatedAt from agents${companyWhere} order by created_at asc`, + `select id, company_id as companyId, name, role, title, department, status, reports_to as reportsTo, connector_id as connectorId, workspace_id as workspaceId, model, capabilities, role_level as roleLevel, budget_monthly_usd as budgetMonthlyUsd, spent_monthly_usd as spentMonthlyUsd, metadata_json as metadataJson, heartbeat_enabled as heartbeatEnabled, heartbeat_interval_sec as heartbeatIntervalSec, last_heartbeat_at as lastHeartbeatAt, session_id as sessionId, created_at as createdAt, updated_at as updatedAt from agents${companyWhere} order by created_at asc`, ...companyArgs, ).map((record: SqliteRow) => ({ ...record, heartbeatEnabled: toBoolean(record.heartbeatEnabled) })) as AgentRecord[], goals: readAll( diff --git a/src/main/database/tasks-crud.ts b/src/main/database/tasks-crud.ts index 7604bdaf..3606b80d 100644 --- a/src/main/database/tasks-crud.ts +++ b/src/main/database/tasks-crud.ts @@ -440,6 +440,71 @@ export function buildTaskCrudMethods( .run(ts, reason, handoffMarkdown, ts, sessionId); }, + /** + * Returns the agent's active session id, creating one lazily on first + * call. Session is scoped to the agent (not to any particular task); + * switching tasks emits a user turn on the ongoing transcript rather + * than spawning a new session. + * + * Existing agents with pre-migration per-task sessions start fresh + * after this upgrade: their first run creates a new agent-scoped + * session. The old task_sessions + session_entries rows remain in the + * DB as historical records but are not loaded for new runs. This is a + * one-time transcript reset, accepted over a risky merge on a live + * production DB. + */ + ensureAgentSession(input: { + agentId: string; + companyId: string; + adapterType: string; + }): string { + const agentRow = getDb() + .prepare("SELECT session_id FROM agents WHERE id = ?") + .get(input.agentId) as { session_id: string | null } | undefined; + if (agentRow?.session_id) return agentRow.session_id; + + const sessionId = randomUUID(); + const ts = nowIso(); + transaction(() => { + // task_id is left as empty string — the session is agent-scoped; + // the column persists for backwards-compatible schema but no + // longer keys the session to a particular task. + getDb() + .prepare( + `INSERT INTO task_sessions (id, company_id, task_id, agent_id, adapter_type, + session_params_json, session_display_id, run_count, total_input_tokens, + total_output_tokens, created_at, updated_at) + VALUES (?, ?, '', ?, ?, NULL, NULL, 0, 0, 0, ?, ?)`, + ) + .run(sessionId, input.companyId, input.agentId, input.adapterType, ts, ts); + getDb() + .prepare("UPDATE agents SET session_id = ?, updated_at = ? WHERE id = ?") + .run(sessionId, ts, input.agentId); + }); + return sessionId; + }, + + setAgentSessionId(agentId: string, sessionId: string): void { + const ts = nowIso(); + getDb() + .prepare("UPDATE agents SET session_id = ?, updated_at = ? WHERE id = ?") + .run(sessionId, ts, agentId); + }, + + /** + * Point every task currently assigned to an agent at the agent's new + * active session. Called after session rotation so UI paths reading + * per-task active_session_id still see the agent's current session. + */ + updateTasksActiveSessionForAgent(agentId: string, sessionId: string): void { + const ts = nowIso(); + getDb() + .prepare( + "UPDATE tasks SET active_session_id = ?, updated_at = ? WHERE assignee_agent_id = ?", + ) + .run(sessionId, ts, agentId); + }, + createNewSession(input: { taskId: string; companyId: string; diff --git a/src/main/orchestration/turn-context.ts b/src/main/orchestration/turn-context.ts index 336f6e91..826bc4c4 100644 --- a/src/main/orchestration/turn-context.ts +++ b/src/main/orchestration/turn-context.ts @@ -127,11 +127,24 @@ function resolveEnabledConnectorModels( function resolveSession( db: AppDatabase, + agent: RunDispatchParams["agent"], task: RunDispatchParams["task"], + companyId: string, connectorId: string, sessionHandoff: string | null, deps: Pick, ) { + // Per-agent session: one ongoing transcript per agent; switching tasks + // becomes a "now work on X" user turn rather than a fresh session. + const agentSessionId = db.ensureAgentSession({ + agentId: agent.id, + companyId, + adapterType: connectorId, + }); + // Mirror the agent's active session onto the task pointer so UI paths + // that read per-task active_session_id keep working. + db.updateTasksActiveSessionForAgent(agent.id, agentSessionId); + const activeSession = db.getActiveTaskSession(task.id); let taskSession: { id: string | null; @@ -145,7 +158,7 @@ function resolveSession( : null, sessionDisplayId: activeSession.sessionDisplayId, } - : { id: null, ...db.getTaskSession(task.id) }; + : { id: agentSessionId, sessionParams: null, sessionDisplayId: null }; let finalHandoff = sessionHandoff; if (deps.evaluatePreRunSessionRotation && activeSession) { @@ -155,7 +168,7 @@ function resolveSession( if (!finalHandoff && rotationResult.handoff) { finalHandoff = rotationResult.handoff; } - deps.logger.info(`[session] Pre-run compaction for task ${task.id}: session rotated before execution`); + deps.logger.info(`[session] Pre-run compaction for agent ${agent.id}: session rotated before execution`); } } @@ -260,7 +273,7 @@ export function buildRunPayload(params: RunDispatchParams, deps: TurnContextDeps const reviewTasks = db.listReviewTasksForManager(agent.id, companyId); const recentComments = assembleRecentComments(assignedTasks, reviewTasks, snapshot); const initialHandoff = resolveSessionHandoff(task.id, deps); - const { taskSession, sessionHandoff } = resolveSession(db, task, connector.id, initialHandoff, deps); + const { taskSession, sessionHandoff } = resolveSession(db, agent, task, companyId, connector.id, initialHandoff, deps); const skills = resolveSkills(companyId, agent.id, connector.id, agent.name, deps); const systemPrompt = apiUrl ? buildTaskPrompt({ diff --git a/src/main/session-management-service.ts b/src/main/session-management-service.ts index 442261d0..1af93674 100644 --- a/src/main/session-management-service.ts +++ b/src/main/session-management-service.ts @@ -128,19 +128,23 @@ export function createSessionManagementService({ db, logger }: SessionManagement return result; } - // Archive current session, generate handoff, create new session, update task pointer + // Archive current session, generate handoff, create new agent-scoped + // session, update the agent pointer and any task active-session pointers + // still referencing the old session. function rotateSession(session: TaskSessionRecord): string { const reason = "threshold_exceeded"; const handoff = generateHandoff(session.taskId, session); - logger.info(`[session] Rotating session ${session.id} for task ${session.taskId}: ${reason}`); + logger.info(`[session] Rotating session ${session.id} for agent ${session.agentId}: ${reason}`); // 1. Archive old session db.archiveSession(session.id, reason, handoff); - // 2. Create new session row with fresh counters, carry forward session_params + // 2. Create new session row with fresh counters, carry forward + // session_params. task_id is empty because the session is scoped + // to the agent, not to any particular task. const newSessionId = db.createNewSession({ - taskId: session.taskId, + taskId: "", companyId: session.companyId, agentId: session.agentId, adapterType: session.adapterType, @@ -148,7 +152,12 @@ export function createSessionManagementService({ db, logger }: SessionManagement sessionDisplayId: session.sessionDisplayId, }); - logger.info(`[session] New session ${newSessionId} created for task ${session.taskId}`); + // 3. Point the agent and any tasks still holding the old session id + // at the freshly-minted session. + db.setAgentSessionId(session.agentId, newSessionId); + db.updateTasksActiveSessionForAgent(session.agentId, newSessionId); + + logger.info(`[session] New session ${newSessionId} created for agent ${session.agentId}`); return newSessionId; } diff --git a/src/shared/types/records-core.ts b/src/shared/types/records-core.ts index 935b33f8..6e85b6ac 100644 --- a/src/shared/types/records-core.ts +++ b/src/shared/types/records-core.ts @@ -95,6 +95,15 @@ export interface AgentRecord { heartbeatEnabled: boolean; heartbeatIntervalSec: number; lastHeartbeatAt: string | null; + /** + * Agent-scoped active session id. One active session per agent (per + * company) — switching tasks appends a "now work on X" user turn to the + * agent's ongoing transcript rather than starting a fresh session. + * Optional on the read model for the same reason `roleLevel` is: existing + * callers constructing an AgentRecord literally (tests, fixtures) do not + * need to set it, and pre-migration rows legitimately have NULL. + */ + sessionId?: string | null; createdAt: string; updatedAt: string; }