WI-6 (#833): constrained recovery command + mosaic-context-refresh skill wrapper (#846)

packages/mosaic/framework/runtime/pi/mosaic-extension.ts

@@ -31,6 +31,14 @@ import { registerLeaseLifecycleHooks, type LeaseLifecyclePiApi } from './lease-l
 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
@@ -135,6 +143,78 @@ function checkPiMutatorGate(toolName: string): { block: true; reason: string } |
   };
 }
 
+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
 // ---------------------------------------------------------------------------
@@ -287,6 +367,32 @@ export default function register(pi: ExtensionAPI) {
   // 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();