518 lines
19 KiB
TypeScript
518 lines
19 KiB
TypeScript
/**
|
|
* mosaic-extension.ts — Pi Extension for Mosaic Framework
|
|
*
|
|
* Integrates the Mosaic agent framework into Pi sessions launched via `mosaic pi`.
|
|
* Handles:
|
|
* 1. Session start — run repo hooks, detect active mission, display status
|
|
* 2. Session end — run repo hooks, clean up session lock
|
|
* 3. Mission context — inject active mission state into conversation
|
|
* 4. Memory routing — remind agent to use ~/.config/mosaic/memory/
|
|
*/
|
|
|
|
import type { ExtensionAPI, ExtensionContext } from '@earendil-works/pi-coding-agent';
|
|
import { Type } from 'typebox';
|
|
import {
|
|
existsSync,
|
|
readFileSync,
|
|
writeFileSync,
|
|
unlinkSync,
|
|
mkdirSync,
|
|
renameSync,
|
|
} from 'node:fs';
|
|
import { join, basename } from 'node:path';
|
|
import { homedir } from 'node:os';
|
|
import { execSync, spawnSync } from 'node:child_process';
|
|
import { registerLeaseLifecycleHooks, type LeaseLifecyclePiApi } from './lease-lifecycle.js';
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Config
|
|
// ---------------------------------------------------------------------------
|
|
|
|
const MOSAIC_HOME = process.env['MOSAIC_HOME'] ?? join(homedir(), '.config', 'mosaic');
|
|
const MUTATOR_GATE = join(MOSAIC_HOME, 'tools', 'lease-broker', 'mutator-gate.py');
|
|
const LEASE_REVOKER = join(MOSAIC_HOME, 'tools', 'lease-broker', 'revoke-lease.py');
|
|
const RECOVERY_COMMAND = join(MOSAIC_HOME, 'tools', 'lease-broker', 'recover-context.py');
|
|
const RECEIPT_OBSERVER_CLIENT = join(
|
|
MOSAIC_HOME,
|
|
'tools',
|
|
'lease-broker',
|
|
'receipt-observer-client.py',
|
|
);
|
|
const RECOVERY_TOOL = 'mosaic_context_recover';
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Helpers
|
|
// ---------------------------------------------------------------------------
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Native heartbeat (fleet R14/R15)
|
|
// ---------------------------------------------------------------------------
|
|
// When this agent runs under the Mosaic fleet (MOSAIC_AGENT_NAME set), the
|
|
// extension writes its OWN heartbeat in the same .hb contract `fleet ps` reads
|
|
// (ts/pid/status[/model]) and touches a `.hb.native` precedence marker so the
|
|
// shell sidecar defers. Native HB knows the real turn state (busy/ok), so it is
|
|
// more accurate than the pane-PID-only sidecar fallback.
|
|
const HB_AGENT_NAME = process.env['MOSAIC_AGENT_NAME'] ?? '';
|
|
const HB_RUN_DIR = process.env['MOSAIC_HEARTBEAT_RUN_DIR'] ?? join(MOSAIC_HOME, 'fleet', 'run');
|
|
const HB_INTERVAL_MS = (() => {
|
|
const s = Number.parseInt(process.env['MOSAIC_HEARTBEAT_INTERVAL'] ?? '', 10);
|
|
return Number.isFinite(s) && s > 0 ? s * 1000 : 15_000;
|
|
})();
|
|
|
|
function nativeHbEnabled(): boolean {
|
|
return HB_AGENT_NAME.length > 0;
|
|
}
|
|
|
|
function readModelId(ctx: ExtensionContext): string | null {
|
|
const m = ctx.model as unknown as { id?: string; name?: string } | undefined;
|
|
return m?.id ?? m?.name ?? null;
|
|
}
|
|
|
|
function writeNativeHeartbeat(status: 'ok' | 'busy', model: string | null): void {
|
|
if (!nativeHbEnabled()) return;
|
|
try {
|
|
mkdirSync(HB_RUN_DIR, { recursive: true });
|
|
const hb = join(HB_RUN_DIR, `${HB_AGENT_NAME}.hb`);
|
|
const lines = [`ts=${nowIso()}`, `pid=${process.pid}`, `status=${status}`];
|
|
if (model) lines.push(`model=${model}`);
|
|
const tmp = `${hb}.tmp.${process.pid}`;
|
|
writeFileSync(tmp, lines.join('\n') + '\n');
|
|
renameSync(tmp, hb); // atomic replace — fleet ps never reads a partial file
|
|
// Precedence marker: tells the shell sidecar that native HB is authoritative.
|
|
writeFileSync(join(HB_RUN_DIR, `${HB_AGENT_NAME}.hb.native`), nowIso() + '\n');
|
|
} catch {
|
|
// Best-effort: never let heartbeat I/O disrupt the Pi session.
|
|
}
|
|
}
|
|
|
|
function clearNativeMarker(): void {
|
|
if (!nativeHbEnabled()) return;
|
|
try {
|
|
const m = join(HB_RUN_DIR, `${HB_AGENT_NAME}.hb.native`);
|
|
if (existsSync(m)) unlinkSync(m); // native stopping — let the sidecar take over
|
|
} catch {
|
|
/* ignore */
|
|
}
|
|
}
|
|
|
|
function safeRead(filePath: string): string | null {
|
|
try {
|
|
return readFileSync(filePath, 'utf-8');
|
|
} catch {
|
|
return null;
|
|
}
|
|
}
|
|
|
|
function safeJsonRead(filePath: string): Record<string, unknown> | null {
|
|
const raw = safeRead(filePath);
|
|
if (!raw) return null;
|
|
try {
|
|
return JSON.parse(raw) as Record<string, unknown>;
|
|
} catch {
|
|
return null;
|
|
}
|
|
}
|
|
|
|
function nowIso(): string {
|
|
return new Date().toISOString().replace(/\.\d{3}Z$/, 'Z');
|
|
}
|
|
|
|
function runPiLeaseRevoker(args: string[]): boolean {
|
|
const result = spawnSync('python3', [LEASE_REVOKER, ...args], {
|
|
encoding: 'utf8',
|
|
timeout: 2_000,
|
|
env: process.env,
|
|
});
|
|
return result.status === 0;
|
|
}
|
|
|
|
function checkPiMutatorGate(toolName: string): { block: true; reason: string } | undefined {
|
|
const result = spawnSync('python3', [MUTATOR_GATE, '--runtime', 'pi'], {
|
|
input: `${JSON.stringify({ tool_name: toolName })}\n`,
|
|
encoding: 'utf8',
|
|
timeout: 2_000,
|
|
env: process.env,
|
|
});
|
|
if (result.status === 0) return undefined;
|
|
const detail = String(result.stderr ?? '')
|
|
.trim()
|
|
.split('\n')[0];
|
|
return {
|
|
block: true,
|
|
reason: detail || 'BLOCKED: Mosaic mutator gate is unavailable or the lease is UNVERIFIED.',
|
|
};
|
|
}
|
|
|
|
function checkPiRecoveryGate(): { block: true; reason: string } | undefined {
|
|
return checkPiMutatorGate(RECOVERY_TOOL);
|
|
}
|
|
|
|
function assistantMessageText(message: unknown): string | undefined {
|
|
if (typeof message !== 'object' || message === null) return undefined;
|
|
const value = message as { role?: unknown; content?: unknown };
|
|
if (value.role !== 'assistant') return undefined;
|
|
if (typeof value.content === 'string') return value.content;
|
|
if (!Array.isArray(value.content)) return undefined;
|
|
const text: string[] = [];
|
|
for (const part of value.content) {
|
|
if (typeof part !== 'object' || part === null) return undefined;
|
|
const typed = part as { type?: unknown; text?: unknown };
|
|
if (typed.type !== 'text' || typeof typed.text !== 'string') return undefined;
|
|
text.push(typed.text);
|
|
}
|
|
return text.join('');
|
|
}
|
|
|
|
function recordPiMessageEnd(message: unknown): void {
|
|
const latestAssistantMessage = assistantMessageText(message);
|
|
if (latestAssistantMessage === undefined) return;
|
|
// This sends finalized Pi message_end content only to the daemon-owned
|
|
// authenticated observer transport, never to the public broker request API.
|
|
spawnSync('python3', [RECEIPT_OBSERVER_CLIENT, '--runtime', 'pi'], {
|
|
input: `${JSON.stringify({ latest_assistant_message: latestAssistantMessage })}\n`,
|
|
encoding: 'utf8',
|
|
timeout: 2_000,
|
|
env: process.env,
|
|
});
|
|
}
|
|
|
|
function runPiRecoveryCommand(params: {
|
|
phase: 'begin' | 'complete';
|
|
construction?: string;
|
|
compactionEpoch?: number;
|
|
requestEpoch?: number;
|
|
}): { content: Array<{ type: 'text'; text: string }> } {
|
|
const args = [RECOVERY_COMMAND, params.phase];
|
|
if (params.phase === 'begin') {
|
|
if (
|
|
typeof params.construction !== 'string' ||
|
|
!Number.isInteger(params.compactionEpoch) ||
|
|
!Number.isInteger(params.requestEpoch) ||
|
|
params.compactionEpoch < 0 ||
|
|
params.requestEpoch < 0
|
|
) {
|
|
return {
|
|
content: [
|
|
{ type: 'text', text: 'Recovery begin requires construction and non-negative epochs.' },
|
|
],
|
|
};
|
|
}
|
|
args.push(
|
|
'--construction',
|
|
params.construction,
|
|
'--compaction-epoch',
|
|
String(params.compactionEpoch),
|
|
'--request-epoch',
|
|
String(params.requestEpoch),
|
|
);
|
|
}
|
|
const result = spawnSync('python3', args, {
|
|
encoding: 'utf8',
|
|
timeout: 3_000,
|
|
env: process.env,
|
|
});
|
|
const output = result.status === 0 ? String(result.stdout ?? '') : String(result.stderr ?? '');
|
|
return { content: [{ type: 'text', text: output || 'Constrained recovery refused.' }] };
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Mission detection
|
|
// ---------------------------------------------------------------------------
|
|
|
|
interface ActiveMission {
|
|
name: string;
|
|
id: string;
|
|
status: string;
|
|
milestonesTotal: number;
|
|
milestonesCompleted: number;
|
|
}
|
|
|
|
function detectMission(cwd: string): ActiveMission | null {
|
|
const missionFile = join(cwd, '.mosaic', 'orchestrator', 'mission.json');
|
|
const data = safeJsonRead(missionFile);
|
|
if (!data) return null;
|
|
|
|
const status = String(data.status ?? 'inactive');
|
|
if (status !== 'active' && status !== 'paused') return null;
|
|
|
|
const milestones = Array.isArray(data.milestones) ? data.milestones : [];
|
|
const completed = milestones.filter(
|
|
(m: unknown) =>
|
|
typeof m === 'object' && m !== null && (m as Record<string, unknown>).status === 'completed',
|
|
).length;
|
|
|
|
return {
|
|
name: String(data.name ?? 'unnamed'),
|
|
id: String(data.mission_id ?? ''),
|
|
status,
|
|
milestonesTotal: milestones.length,
|
|
milestonesCompleted: completed,
|
|
};
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Session lock management
|
|
// ---------------------------------------------------------------------------
|
|
|
|
function sessionLockPath(cwd: string): string {
|
|
return join(cwd, '.mosaic', 'orchestrator', 'session.lock');
|
|
}
|
|
|
|
function writeSessionLock(cwd: string): void {
|
|
const lockDir = join(cwd, '.mosaic', 'orchestrator');
|
|
if (!existsSync(lockDir)) return; // Only write lock if orchestrator dir exists
|
|
|
|
const lock = {
|
|
session_id: `pi-${new Date().toISOString().replace(/[:.]/g, '-')}-${process.pid}`,
|
|
runtime: 'pi',
|
|
pid: process.pid,
|
|
started_at: nowIso(),
|
|
project_path: cwd,
|
|
milestone_id: '',
|
|
};
|
|
|
|
try {
|
|
writeFileSync(sessionLockPath(cwd), JSON.stringify(lock, null, 2) + '\n', 'utf-8');
|
|
} catch {
|
|
// Non-fatal — orchestrator dir may not be writable
|
|
}
|
|
}
|
|
|
|
function cleanSessionLock(cwd: string): void {
|
|
try {
|
|
const lockFile = sessionLockPath(cwd);
|
|
if (existsSync(lockFile)) {
|
|
unlinkSync(lockFile);
|
|
}
|
|
} catch {
|
|
// Non-fatal
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Repo hooks
|
|
// ---------------------------------------------------------------------------
|
|
|
|
function runRepoHook(cwd: string, hookName: string): void {
|
|
const script = join(cwd, 'scripts', 'agent', `${hookName}.sh`);
|
|
if (!existsSync(script)) return;
|
|
|
|
try {
|
|
spawnSync('bash', [script], {
|
|
cwd,
|
|
stdio: 'pipe',
|
|
timeout: 30_000,
|
|
env: { ...process.env, MOSAIC_RUNTIME: 'pi' },
|
|
});
|
|
} catch {
|
|
// Non-fatal
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Build mission summary for notifications
|
|
// ---------------------------------------------------------------------------
|
|
|
|
function buildMissionSummary(cwd: string, mission: ActiveMission): string {
|
|
const lines: string[] = [
|
|
`Mission: ${mission.name}`,
|
|
`Status: ${mission.status} | Milestones: ${mission.milestonesCompleted}/${mission.milestonesTotal}`,
|
|
];
|
|
|
|
// Task counts
|
|
const tasksFile = join(cwd, 'docs', 'TASKS.md');
|
|
const tasksContent = safeRead(tasksFile);
|
|
if (tasksContent) {
|
|
const tableRows = tasksContent
|
|
.split('\n')
|
|
.filter((l) => l.startsWith('|') && !l.includes('---'));
|
|
const total = Math.max(0, tableRows.length - 1); // minus header
|
|
const done = (tasksContent.match(/\|\s*done\s*\|/gi) ?? []).length;
|
|
lines.push(`Tasks: ${done} done / ${total} total`);
|
|
}
|
|
|
|
// Latest scratchpad
|
|
try {
|
|
const spDir = join(cwd, 'docs', 'scratchpads');
|
|
if (existsSync(spDir)) {
|
|
const files = execSync(`ls -t "${spDir}"/*.md 2>/dev/null | head -1`, {
|
|
encoding: 'utf-8',
|
|
timeout: 5000,
|
|
}).trim();
|
|
if (files) lines.push(`Scratchpad: ${basename(files)}`);
|
|
}
|
|
} catch {
|
|
// Non-fatal
|
|
}
|
|
|
|
lines.push('', 'Read ORCHESTRATOR-PROTOCOL.md + TASKS.md before proceeding.');
|
|
return lines.join('\n');
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Extension registration
|
|
// ---------------------------------------------------------------------------
|
|
|
|
export default function register(pi: ExtensionAPI) {
|
|
let sessionCwd = process.cwd();
|
|
let hbStatus: 'ok' | 'busy' = 'ok';
|
|
let hbModel: string | null = null;
|
|
let hbTimer: ReturnType<typeof setInterval> | null = null;
|
|
|
|
// ── Compaction observers and same-PID generation rollover ─────────────
|
|
registerLeaseLifecycleHooks(pi as unknown as LeaseLifecyclePiApi, runPiLeaseRevoker);
|
|
|
|
// ── Whole mutator-class authorization gate ────────────────────────────
|
|
// Every Pi tool, including unknown/custom tools, reaches the broker-backed
|
|
// class gate before execution. Broker/script failure blocks fail-closed.
|
|
pi.on('tool_call', async (event) => checkPiMutatorGate(event.toolName));
|
|
|
|
// Pi records only a finalized assistant entry at message_end. It never uses
|
|
// after_provider_response, which occurs before stream consumption.
|
|
pi.on('message_end', async (event) => {
|
|
recordPiMessageEnd((event as unknown as { message?: unknown }).message);
|
|
});
|
|
|
|
// The recovery custom tool is the only Pi invocation that maps to the
|
|
// broker's exempt RECOVERY_TOOL identity. It is not a Bash exception.
|
|
pi.registerTool({
|
|
name: RECOVERY_TOOL,
|
|
label: 'Mosaic Context Recovery',
|
|
description:
|
|
'Run the constrained broker-backed context recovery flow. This is the sole ungated mutator.',
|
|
parameters: Type.Object({
|
|
phase: Type.Union([Type.Literal('begin'), Type.Literal('complete')]),
|
|
construction: Type.Optional(Type.String()),
|
|
compactionEpoch: Type.Optional(Type.Integer({ minimum: 0 })),
|
|
requestEpoch: Type.Optional(Type.Integer({ minimum: 0 })),
|
|
}),
|
|
async execute(_toolCallId, params) {
|
|
const blocked = checkPiRecoveryGate();
|
|
if (blocked !== undefined) return { content: [{ type: 'text', text: blocked.reason }] };
|
|
return runPiRecoveryCommand(params);
|
|
},
|
|
});
|
|
|
|
// ── Session Start ─────────────────────────────────────────────────────
|
|
pi.on('session_start', async (_event, ctx) => {
|
|
sessionCwd = process.cwd();
|
|
|
|
// Run repo session-start hook
|
|
runRepoHook(sessionCwd, 'session-start');
|
|
|
|
// Detect active mission
|
|
const mission = detectMission(sessionCwd);
|
|
|
|
if (mission) {
|
|
// Write session lock for orchestrator awareness
|
|
writeSessionLock(sessionCwd);
|
|
|
|
const summary = buildMissionSummary(sessionCwd, mission);
|
|
ctx.ui.notify(`🎯 Active Mosaic Mission\n${summary}`, 'info');
|
|
} else {
|
|
ctx.ui.notify('Mosaic framework loaded', 'info');
|
|
}
|
|
|
|
// Native heartbeat: write immediately, then on an interval. Idle = 'ok';
|
|
// turn_start/turn_end flip the status so `fleet ps` reflects real activity.
|
|
if (nativeHbEnabled()) {
|
|
hbModel = readModelId(ctx);
|
|
writeNativeHeartbeat('ok', hbModel);
|
|
hbTimer = setInterval(() => writeNativeHeartbeat(hbStatus, hbModel), HB_INTERVAL_MS);
|
|
if (typeof hbTimer.unref === 'function') hbTimer.unref();
|
|
}
|
|
});
|
|
|
|
// ── Turn lifecycle → accurate busy/ok heartbeat ───────────────────────
|
|
pi.on('turn_start', async (_event, ctx) => {
|
|
hbStatus = 'busy';
|
|
hbModel = readModelId(ctx) ?? hbModel;
|
|
writeNativeHeartbeat('busy', hbModel);
|
|
});
|
|
pi.on('turn_end', async (_event, ctx) => {
|
|
hbStatus = 'ok';
|
|
hbModel = readModelId(ctx) ?? hbModel;
|
|
writeNativeHeartbeat('ok', hbModel);
|
|
});
|
|
|
|
// ── Session Shutdown ──────────────────────────────────────────────────
|
|
// (The pi API event is 'session_shutdown'; the prior 'session_end' handler
|
|
// never fired — fixed here so repo hooks + lock cleanup actually run.)
|
|
pi.on('session_shutdown', async (_event, _ctx) => {
|
|
if (hbTimer) {
|
|
clearInterval(hbTimer);
|
|
hbTimer = null;
|
|
}
|
|
clearNativeMarker();
|
|
|
|
// Run repo session-end hook
|
|
runRepoHook(sessionCwd, 'session-end');
|
|
|
|
// Clean up session lock
|
|
cleanSessionLock(sessionCwd);
|
|
});
|
|
|
|
// ── Register /mosaic-status command ───────────────────────────────────
|
|
pi.registerCommand('mosaic-status', {
|
|
description: 'Show Mosaic mission status for the current project',
|
|
handler: async (_args, ctx) => {
|
|
const mission = detectMission(sessionCwd);
|
|
if (!mission) {
|
|
ctx.ui.notify('No active Mosaic mission in this project.', 'info');
|
|
return;
|
|
}
|
|
const summary = buildMissionSummary(sessionCwd, mission);
|
|
ctx.ui.notify(`🎯 Mission Status\n${summary}`, 'info');
|
|
},
|
|
});
|
|
|
|
// ── Register /mosaic-memory command ───────────────────────────────────
|
|
pi.registerCommand('mosaic-memory', {
|
|
description: 'Show Mosaic memory directory path and contents',
|
|
handler: async (_args, ctx) => {
|
|
const memDir = join(MOSAIC_HOME, 'memory');
|
|
if (!existsSync(memDir)) {
|
|
ctx.ui.notify(`Memory directory: ${memDir} (empty)`, 'info');
|
|
return;
|
|
}
|
|
try {
|
|
const files = execSync(`ls -la "${memDir}" 2>/dev/null`, {
|
|
encoding: 'utf-8',
|
|
timeout: 5000,
|
|
}).trim();
|
|
ctx.ui.notify(`Memory directory: ${memDir}\n${files}`, 'info');
|
|
} catch {
|
|
ctx.ui.notify(`Memory directory: ${memDir}`, 'info');
|
|
}
|
|
},
|
|
});
|
|
|
|
// ── Register mosaic_mission_status tool (model-callable) ──────────────
|
|
// R14 "proper tool usage": give the agent a first-class tool to load its
|
|
// active Mosaic mission, milestone progress, task counts, and latest
|
|
// scratchpad — so it self-orients on in-flight work before planning,
|
|
// instead of shelling out or guessing. Mirrors the /mosaic-status command
|
|
// but returns the summary as tool output the LLM can read.
|
|
pi.registerTool({
|
|
name: 'mosaic_mission_status',
|
|
label: 'Mosaic Mission Status',
|
|
description:
|
|
'Return the active Mosaic mission, milestone progress, task counts, and latest scratchpad for the current project. Returns a note when no mission is active.',
|
|
promptSnippet: 'Read the active Mosaic mission + task state for the current project',
|
|
promptGuidelines: [
|
|
'Use mosaic_mission_status at the start of a session or task to load the active mission, milestone progress, and open tasks before planning work.',
|
|
],
|
|
parameters: Type.Object({}),
|
|
async execute(_toolCallId, _params, _signal, _onUpdate, _ctx) {
|
|
const mission = detectMission(sessionCwd);
|
|
const text = mission
|
|
? buildMissionSummary(sessionCwd, mission)
|
|
: 'No active Mosaic mission in this project.';
|
|
return {
|
|
content: [{ type: 'text', text }],
|
|
details: mission ? { ...mission } : { active: false },
|
|
};
|
|
},
|
|
});
|
|
}
|