refactor(session): scope pi sessions to agents instead of tasks

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`.
This commit is contained in:
ZacharyZhang-NY committed 2026-04-17 15:42:13 -04:00
1 parent 5a3647e98e
commit 92795f7bfb
10 files changed
+127 -16

No files matched your search

+4 -4
View File
@@ -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<AgentRecord>(
"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<AgentRecord>(
"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,
+2 -2
View File
@@ -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;
+12
View File
@@ -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) {
+1 -1
View File
@@ -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) => ({
+1 -1
View File
@@ -63,7 +63,7 @@ export function buildSnapshotMethods(
...companyArgs,
).map((record) => ({ ...record, isPrimary: toBoolean(record.isPrimary) })) as WorkspaceRecord[],
agents: readAll<SqliteRow>(
`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<GoalRecord>(
+65
View File
@@ -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;