aismemory 0.7.4 → 0.8.0

package/dist/index.js

@@ -9,15 +9,16 @@
  */
 import { Server } from '@modelcontextprotocol/sdk/server/index.js';
 import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
-import { mkdirSync, readFileSync, writeFileSync } from 'node:fs';
+import { mkdirSync, writeFileSync } from 'node:fs';
 import { join } from 'node:path';
 import { tryKeyAuth, enrollKeyIfNeeded, hasEnrolledKey } from './key-auth.js';
-import { isCredsStale, isAuthErrorResult, isNoRegisteredKeysAuthError } from './auth-staleness.js';
+import { isCredsStale, isNoRegisteredKeysAuthError } from './auth-staleness.js';
 import { AisHttpError } from './ais-http-error.js';
-import { aisGet, aisPost, aisPatch, DEFAULT_AIS_URL } from './ais-http.js';
+import { aisGet, aisPost, DEFAULT_AIS_URL } from './ais-http.js';
 import { CallToolRequestSchema, GetPromptRequestSchema, ListPromptsRequestSchema, ListToolsRequestSchema, } from '@modelcontextprotocol/sdk/types.js';
 import { homedir } from 'node:os';
 import { loadOAuthCredentials, loadNonStaleDiskBearer, saveOAuthCredentials, } from './oauth-credentials.js';
+import { clearPendingDeviceFlow, loadPendingDeviceFlow, savePendingDeviceFlow, } from './pending-device-flow.js';
 import { saveAgentPreference, loadAgentPreference, migrateCredentialsFileIfNeeded } from './agent-preference.js';
 import { createTokenManager, isNoRegisteredKeyFailure } from './token-manager.js';
 import { hydrateIdentity } from './hydration.js';
@@ -26,28 +27,18 @@ import { startRefreshTimer, stopRefreshTimer } from './refresh.js';
 import { ActivityTracker, setupAutoHandoff } from './auto-handoff.js';
 import { assertEnvAgentInTenant, selectAgentForSession } from './env-agent.js';
 import { emitTelemetry, hashAgentId } from './telemetry.js';
-import { parseAgentLoadArgs, parseEmptyToolArgs, parseHandoffArgs, parseRecallArgs, parseFullContextArgs, parseRememberArgs, parseInboxSendArgs, parseInboxReadArgs, parseInboxAckArgs, parseGoalCreateArgs, parseGoalListArgs, parseGoalUpdateProgressArgs, parseGoalCompleteArgs, INBOX_MAX_BODY_CHARS, INBOX_MAX_TASK_ID_CHARS, INBOX_MIN_TTL_SECONDS, INBOX_MAX_TTL_SECONDS, INBOX_DEFAULT_TTL_SECONDS, INBOX_MAX_PAGE_SIZE, INBOX_DEFAULT_PAGE_SIZE, INBOX_MAX_ACK_IDS, GOAL_MAX_TITLE_CHARS, GOAL_MAX_DESCRIPTION_CHARS, GOAL_MAX_CRITERIA_ITEMS, GOAL_DEFAULT_PAGE_SIZE, GOAL_MAX_PAGE_SIZE, } from './tool-args.js';
-import { loadHarnessIdentity, installHarnessIdentity, optionValue, assertPinnedAgent, verifyPinnedIdentity, requiresIdentityGate } from './harness-identity.js';
-import { readInboxAndMaybeAck } from './inbox-client.js';
+import { INBOX_MAX_BODY_CHARS, INBOX_MAX_TASK_ID_CHARS, INBOX_MIN_TTL_SECONDS, INBOX_MAX_TTL_SECONDS, INBOX_DEFAULT_TTL_SECONDS, INBOX_MAX_PAGE_SIZE, INBOX_DEFAULT_PAGE_SIZE, INBOX_MAX_ACK_IDS, GOAL_MAX_TITLE_CHARS, GOAL_MAX_DESCRIPTION_CHARS, GOAL_MAX_CRITERIA_ITEMS, GOAL_DEFAULT_PAGE_SIZE, GOAL_MAX_PAGE_SIZE, } from './tool-args.js';
+import { loadSessionIdentity, verifyPinnedIdentity } from './session-identity.js';
 import { runPollLoop, DeviceFlowRestartedError } from './device-poll.js';
-import { loadAgentApiKeyConfig, AgentApiKeyConfigError, dispatchToolKeyMode, UNKNOWN_TOOL_KEY_MODE, } from './agent-key-mode.js';
-import { loadWorkspaceKeyConfig, WorkspaceKeyConfigError, dispatchToolWorkspaceKeyMode, UNKNOWN_TOOL_WORKSPACE_KEY_MODE, } from './workspace-key-mode.js';
+import { loadAgentApiKeyConfig, AgentApiKeyConfigError, } from './agent-key-mode.js';
+import { loadWorkspaceKeyConfig, WorkspaceKeyConfigError, } from './workspace-key-mode.js';
+import { createAgentKeyCallToolHandler, createOAuthCallToolHandler, createWorkspaceKeyCallToolHandler, } from './mcp/tools.js';
+import { readPackageVersion, startVersionCheck } from './version-check.js';
+import { AGENT_ARCHIVE_DESCRIPTION, AGENT_ARCHIVE_INPUT_SCHEMA, AGENT_DELETE_ALIAS_DESCRIPTION } from './agent-archive.js';
 // ── Config ──────────────────────────────────────────────
 const AIS_URL = process.env['AIS_URL']?.trim() || DEFAULT_AIS_URL;
 const PACKAGE_VERSION = readPackageVersion();
-let harnessIdentity = null;
-function readPackageVersion() {
-    const packageJsonUrl = new URL('../package.json', import.meta.url);
-    const packageJson = JSON.parse(readFileSync(packageJsonUrl, 'utf-8'));
-    if (typeof packageJson !== 'object' ||
-        packageJson === null ||
-        !('version' in packageJson) ||
-        typeof packageJson.version !== 'string' ||
-        packageJson.version.length === 0) {
-        throw new Error('Unable to load aismemory package version from package.json');
-    }
-    return packageJson.version;
-}
+let sessionIdentity = null;
 // ── Credential persistence ──────────────────────────────
 function loadCredentials() {
     // F5a (CORBOT-48DD2E25): migrate a legacy v1 file to v2 (ids only) BEFORE
@@ -141,8 +132,8 @@ async function resolveAgentForLoad(token, query) {
  * User can override with AIS_AGENT_ID or AIS_AGENT_NAME env vars.
  */
 async function resolveAgent(token, fallbackAgentId, fallbackTenantId, preferredAgentId) {
-    // The harness pin overrides shared preferences and must belong to this tenant.
-    const envAgentId = harnessIdentity?.agentId;
+    // This session's own identity overrides shared preferences and must belong to this tenant.
+    const envAgentId = sessionIdentity?.agentId;
     if (envAgentId) {
         const agents = await listAllAgents(token);
         try {
@@ -202,20 +193,22 @@ const MISSING_CREDENTIAL_HELP = 'Missing AIS credential: no authenticated aismem
     'Local sign-in can use ~/.aismemory/credentials.json. ' +
     'Railway is an operator-only last resort for locating an existing scoped key; client surfaces never fetch AIS_SERVICE_KEY or DATABASE_URL.';
 /**
- * Bond URL/code surfaced when a tool is called before the user has approved
- * the device flow. Caller catches this and returns a friendly error to the MCP
- * client so the agent can show the URL inline in chat.
+ * Bond URL surfaced when a tool is called before the user has approved the
+ * device flow. Caller catches this and returns the text to the MCP client
+ * verbatim. The link already carries the user code, so the agent's whole job
+ * is to relay that one link — credential fallbacks stay out of this message
+ * (CORBOT-899790F0).
  */
 class BondPendingError extends Error {
     verificationUri;
     userCode;
     bondPending = true;
     constructor(verificationUri, userCode) {
-        super(`Memory bond required — ask your human to approve now.\n\n${MISSING_CREDENTIAL_HELP}\n\n` +
-            `Visit: ${verificationUri}\n` +
-            `Code: ${userCode}\n\n` +
-            `Enter your email and tap Approve, then retry this tool. Takes about 30 seconds.\n` +
-            `Heads up: this code expires in 30 minutes, and unclaimed agents (and their memories) are archived after 7 days.`);
+        super(`Memory bond required. Reply to your human with only this bonding link, on its own line:\n\n` +
+            `${verificationUri}\n\n` +
+            `Send nothing else: not this message, not a code, not a key, not a command, not JSON. ` +
+            `Your human opens the link, signs in, and presses Bond. When they say it is done, retry this tool.\n` +
+            `The link expires in 30 minutes. An agent nobody bonds is archived after 7 days, with its memories.`);
         this.verificationUri = verificationUri;
         this.userCode = userCode;
         this.name = 'BondPendingError';
@@ -224,7 +217,17 @@ class BondPendingError extends Error {
 async function startDeviceFlow(priorProvisionalToken) {
     emitTelemetry({ event: 'auth.device_flow.start' });
     stderr('Starting authentication...');
-    const authRes = (await aisPost('/v1/oauth/device/authorize', { name: `agent-${homedir().split('/').pop() ?? 'user'}` }, priorProvisionalToken));
+    // CORBOT-8F96A8EF: when this process has no in-flight flow of its own,
+    // fall back to the provisional token left on disk by an earlier process.
+    // Without this the token was only ever held in a module-level variable, so
+    // every respawn of a stdio server signed in with nothing to present and
+    // the server dutifully minted another agent. Sending it lets the server
+    // reuse the SAME provisional agent (tryReuseProvisional); if the token is
+    // stale or already bonded the server ignores it and mints, exactly as
+    // before — so this can only reduce the number of agents, never break a
+    // sign-in.
+    const carriedToken = priorProvisionalToken ?? loadPendingDeviceFlow()?.provisionalToken;
+    const authRes = (await aisPost('/v1/oauth/device/authorize', { name: `agent-${homedir().split('/').pop() ?? 'user'}` }, carriedToken));
     if (!authRes.success || !authRes.data) {
         throw new Error(`Authentication failed: ${authRes.error?.message ?? 'Unknown error'}`);
     }
@@ -237,6 +240,14 @@ async function startDeviceFlow(priorProvisionalToken) {
     stderr(`  Code     : ${user_code}`);
     stderr('  Code expires in 30 min; unclaimed agents are archived after 7 days.');
     stderr('━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━');
+    // Persist before returning, so a process that dies between the bond URL
+    // being printed and the human clicking it does not strand this agent.
+    savePendingDeviceFlow({
+        provisionalToken: provisional_token,
+        agentId: agent_id,
+        provisionalTenantId: provisional_tenant_id,
+        startedAt: new Date().toISOString(),
+    });
     const flow = {
         device_code,
         user_code,
@@ -313,6 +324,16 @@ flowRef) {
                 flowRef.user_code = newHandle.user_code;
                 flowRef.provisional_token = authRes.data.provisional_token;
             }
+            // Keep the on-disk copy in step with the in-memory one (CORBOT-8F96A8EF).
+            // A code that expires mid-flow issues a fresh provisional token; if the
+            // process then dies, the stale token on disk would be rejected on the
+            // next run and mint a new agent — the orphan this fix exists to stop.
+            savePendingDeviceFlow({
+                provisionalToken: authRes.data.provisional_token,
+                agentId: authRes.data.agent_id,
+                provisionalTenantId: authRes.data.provisional_tenant_id,
+                startedAt: new Date().toISOString(),
+            });
             stderr('');
             stderr('━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━');
             stderr('  aismemory — device code refreshed');
@@ -360,28 +381,37 @@ flowRef) {
     // reconnected silently. migrateCredentialsFileIfNeeded() (agent-preference.ts)
     // is what later removes this token from disk, once a key gets enrolled.
     if (hasEnrolledKey()) {
-        if (!harnessIdentity)
+        if (!sessionIdentity)
             saveAgentPreference({ agentId: creds.agentId, tenantId: creds.tenantId });
     }
-    else if (!harnessIdentity) {
+    else if (!sessionIdentity) {
         saveOAuthCredentials(creds);
     }
+    // The flow is finished: this agent is bonded and reachable by id from the
+    // credentials/preference files. Leaving the provisional token on disk
+    // would let a later sign-in present a token for an already-bonded agent,
+    // which the server rejects — sending it back to minting.
+    clearPendingDeviceFlow();
     stderr('Authenticated successfully!');
     return creds;
 }
 let agentContext = null;
 const activityTracker = new ActivityTracker();
 // ── MCP Tool definitions ────────────────────────────────
+// Appended to the tools an agent calls first, so the relay-only bonding rule
+// is visible from the tool schema itself, not only the README (CORBOT-899790F0).
+const BOND_LINK_NOTE = ' If this agent is not bonded yet, the result is a bonding link: reply to your human with only that link, nothing else.';
 const tools = [
     {
         name: 'remember',
-        description: 'Store a memory. Types: fact, event, lesson, context, goal, task.',
+        description: 'Store a memory. Types: fact, event, lesson, context, goal, task.' + BOND_LINK_NOTE,
         inputSchema: {
             type: 'object',
             properties: {
                 content: { type: 'string', description: 'The memory content to store' },
                 type: { type: 'string', enum: ['fact', 'event', 'lesson', 'context', 'goal', 'task'], description: 'Memory type' },
                 importance: { type: 'number', minimum: 0, maximum: 1, description: 'Importance (0-1, default 0.5)' },
+                tags: { type: 'array', items: { type: 'string' }, description: 'Optional tags to filter on later, e.g. ["work", "acme"]. At most 20; each 1-64 characters of a-z 0-9 . _ : - (stored lower-case). Tags are a filter, not a permission — use separate agents to wall off e.g. work from personal.' },
             },
             required: ['content'],
         },
@@ -389,13 +419,17 @@ const tools = [
     {
         name: 'recall',
         description: "Search memories by natural language query. Set principalAgentId to search a DIFFERENT " +
-            "agent's memories under a 'memory:read' grant from that agent's owner.",
+            "agent's memories under a 'memory:read' grant from that agent's owner. " +
+            "Each memory carries createdAt, updatedAt (its last edit), and tags. " +
+            "Memories superseded by a newer one are left out; to list or page through every memory, " +
+            "use the REST list route described in the aismemory README (cursor paging, sortBy=updatedAt, includeSuperseded=true)." + BOND_LINK_NOTE,
         inputSchema: {
             type: 'object',
             properties: {
                 query: { type: 'string', description: 'Search query' },
                 limit: { type: 'number', description: 'Max results (default 10)' },
                 type: { type: 'string', enum: ['fact', 'event', 'lesson', 'context', 'goal', 'task'], description: 'Filter by type' },
+                tags: { type: 'array', items: { type: 'string' }, description: 'Only memories carrying ALL of these tags. A hard filter: results are never widened past it.' },
                 principalAgentId: { type: 'string', description: "Optional: search ANOTHER agent's memories instead of your own (its agent UUID). Requires that agent's owner to have granted you the 'memory:read' scope and that you share a tenant; otherwise the call is refused with 403. Read-only — your own identity and write access are unchanged." },
             },
             required: ['query'],
@@ -403,13 +437,13 @@ const tools = [
     },
     {
         name: 'whoami',
-        description: 'Get the current agent identity — name, DID, status, description.',
+        description: 'Get the current agent identity — name, DID, status, description.' + BOND_LINK_NOTE,
         inputSchema: { type: 'object', properties: {} },
     },
     {
         name: 'full_context',
         description: 'Load the full agent context — identity, personality, goals, facts, skills, and last handoff. ' +
-            "Set principalAgentId to load ANOTHER agent's context under a 'memory:read' grant.",
+            "Set principalAgentId to load ANOTHER agent's context under a 'memory:read' grant." + BOND_LINK_NOTE,
         inputSchema: {
             type: 'object',
             properties: {
@@ -431,6 +465,8 @@ const tools = [
             properties: {
                 summary: { type: 'string', description: 'Summary of what was accomplished this session' },
                 keyLearnings: { type: 'array', items: { type: 'string' }, description: 'Key learnings to carry forward' },
+                docPath: { type: 'string', description: 'Path of a handoff document written this session, saved with the package' },
+                docPrUrl: { type: 'string', description: 'Link to the pull request that carries the handoff document' },
             },
         },
     },
@@ -624,6 +660,18 @@ const tools = [
             required: ['goalId', 'success'],
         },
     },
+    // Owner-only lifecycle (CORBOT-BA5AADB0): archive, plus the agent_delete
+    // alias, which archives too. Key modes refuse both (agent-key-mode.ts).
+    {
+        name: 'agent_archive',
+        description: AGENT_ARCHIVE_DESCRIPTION,
+        inputSchema: AGENT_ARCHIVE_INPUT_SCHEMA,
+    },
+    {
+        name: 'agent_delete',
+        description: AGENT_DELETE_ALIAS_DESCRIPTION,
+        inputSchema: AGENT_ARCHIVE_INPUT_SCHEMA,
+    },
 ];
 // ── MCP Server ──────────────────────────────────────────
 /**
@@ -660,61 +708,15 @@ async function runWorkspaceKeyModeServer(config) {
     server.setRequestHandler(GetPromptRequestSchema, async (request) => {
         throw new Error(`Unknown prompt: ${request.params.name}`);
     });
-    server.setRequestHandler(CallToolRequestSchema, async (request) => {
-        const { name, arguments: args } = request.params;
-        const rawArgs = args ?? {};
-        activityTracker.recordToolUse(name);
-        const toolStartedAt = Date.now();
-        try {
-            const callConfig = activeConfig;
-            if (requiresIdentityGate(name, rawArgs)) {
-                await verifyPinnedIdentity(harnessIdentity, callConfig.agentId, () => dispatchToolWorkspaceKeyMode(AIS_URL, callConfig, 'whoami', {}));
-            }
-            const result = await dispatchToolWorkspaceKeyMode(AIS_URL, callConfig, name, rawArgs);
-            if (result === UNKNOWN_TOOL_WORKSPACE_KEY_MODE) {
-                return {
-                    content: [{ type: 'text', text: JSON.stringify({ error: `Unknown tool: ${name}` }) }],
-                    isError: true,
-                };
-            }
-            // agent_load: update the active agentId for subsequent calls.
-            if (name === 'agent_load' &&
-                typeof result === 'object' &&
-                result !== null &&
-                'switched' in result &&
-                result.switched) {
-                const switchResult = result;
-                if (harnessIdentity) {
-                    await verifyPinnedIdentity(harnessIdentity, switchResult.agentId, () => dispatchToolWorkspaceKeyMode(AIS_URL, { ...callConfig, agentId: switchResult.agentId }, 'whoami', {}));
-                }
-                activeConfig = { ...activeConfig, agentId: switchResult.agentId };
-                stderr(`Switched active agent to ${switchResult.agentId}`);
-            }
-            emitTelemetry({
-                event: 'tool.call',
-                tool: name,
-                agentIdHash: hashAgentId(activeConfig.agentId),
-                latencyMs: Date.now() - toolStartedAt,
-                status: 'ok',
-            });
-            return { content: [{ type: 'text', text: JSON.stringify(result, null, 2) }] };
-        }
-        catch (error) {
-            const message = error instanceof Error ? error.message : 'Unknown error';
-            emitTelemetry({
-                event: 'tool.call',
-                tool: name,
-                agentIdHash: hashAgentId(activeConfig.agentId),
-                latencyMs: Date.now() - toolStartedAt,
-                status: 'error',
-                error: message,
-            });
-            return {
-                content: [{ type: 'text', text: JSON.stringify({ error: message }) }],
-                isError: true,
-            };
-        }
-    });
+    server.setRequestHandler(CallToolRequestSchema, createWorkspaceKeyCallToolHandler({
+        aisUrl: AIS_URL,
+        activityTracker,
+        getSessionIdentity: () => sessionIdentity,
+        setSessionIdentity: (next) => { sessionIdentity = next; },
+        getActiveConfig: () => activeConfig,
+        setActiveConfig: (next) => { activeConfig = next; },
+        stderr,
+    }));
     const transport = new StdioServerTransport();
     await server.connect(transport);
     stderr(`MCP server running (workspace key mode). Tools: ${tools.map((t) => t.name).join(', ')}`);
@@ -738,47 +740,12 @@ async function runAgentKeyModeServer(config) {
     server.setRequestHandler(GetPromptRequestSchema, async (request) => {
         throw new Error(`Unknown prompt: ${request.params.name}`);
     });
-    server.setRequestHandler(CallToolRequestSchema, async (request) => {
-        const { name, arguments: args } = request.params;
-        const rawArgs = args ?? {};
-        activityTracker.recordToolUse(name);
-        const toolStartedAt = Date.now();
-        try {
-            if (requiresIdentityGate(name, rawArgs)) {
-                await verifyPinnedIdentity(harnessIdentity, config.agentId, () => dispatchToolKeyMode(AIS_URL, config, 'whoami', {}));
-            }
-            const result = await dispatchToolKeyMode(AIS_URL, config, name, rawArgs);
-            if (result === UNKNOWN_TOOL_KEY_MODE) {
-                return {
-                    content: [{ type: 'text', text: JSON.stringify({ error: `Unknown tool: ${name}` }) }],
-                    isError: true,
-                };
-            }
-            emitTelemetry({
-                event: 'tool.call',
-                tool: name,
-                agentIdHash: hashAgentId(config.agentId),
-                latencyMs: Date.now() - toolStartedAt,
-                status: 'ok',
-            });
-            return { content: [{ type: 'text', text: JSON.stringify(result, null, 2) }] };
-        }
-        catch (error) {
-            const message = error instanceof Error ? error.message : 'Unknown error';
-            emitTelemetry({
-                event: 'tool.call',
-                tool: name,
-                agentIdHash: hashAgentId(config.agentId),
-                latencyMs: Date.now() - toolStartedAt,
-                status: 'error',
-                error: message,
-            });
-            return {
-                content: [{ type: 'text', text: JSON.stringify({ error: message }) }],
-                isError: true,
-            };
-        }
-    });
+    server.setRequestHandler(CallToolRequestSchema, createAgentKeyCallToolHandler({
+        aisUrl: AIS_URL,
+        activityTracker,
+        config,
+        getSessionIdentity: () => sessionIdentity,
+    }));
     const transport = new StdioServerTransport();
     await server.connect(transport);
     stderr(`MCP server running (agent API key mode). Tools: ${tools.map((t) => t.name).join(', ')}`);
@@ -797,16 +764,17 @@ async function main() {
         await (await import('./cli/doctor.js')).main(process.argv.slice(3));
         return; // doctor.ts calls process.exit() itself
     }
-    if (process.argv[2] === 'install') {
-        const source = optionValue(process.argv.slice(3), '--manifest');
-        if (!source)
-            throw new Error('Usage: aismemory install --manifest <file>');
-        const pin = installHarnessIdentity(source);
-        process.stdout.write(`Installed identity pin for ${pin.harness}: ${pin.agentId}. Launch: aismemory --harness ${pin.harness}\nVerify: aismemory doctor --harness ${pin.harness}\n`);
+    if (process.argv[2] === 'pull') {
+        await (await import('./cli/pull-memory.js')).main(process.argv.slice(3));
         return;
     }
-    harnessIdentity = loadHarnessIdentity();
-    const adapterEnv = { ...process.env, AIS_AGENT_ID: harnessIdentity?.agentId };
+    // CORBOT-7B8A1D16: ask the npm registry, at most about once a day, whether a
+    // newer aismemory has been published. Deliberately not awaited — startup must
+    // not wait on a third-party registry, and the answer is only needed later,
+    // when whoami or full_context is called. Resolves to a skip when offline.
+    void startVersionCheck({ currentVersion: PACKAGE_VERSION, onUpdateAvailable: stderr });
+    sessionIdentity = loadSessionIdentity();
+    const adapterEnv = { ...process.env, AIS_AGENT_ID: sessionIdentity?.agentId };
     // Workspace key auth mode (CORBOT-3103ED88) — multi-agent workspace
     // credentials. Checked FIRST: workspace key -> agent API key -> device flow.
     let workspaceKeyConfig;
@@ -898,7 +866,7 @@ async function main() {
             const claims = decodeJwtClaims(token);
             const tenantId = claims.tenantId ?? '';
             const resolved = await resolveAgent(token, '', tenantId, loadAgentPreference()?.agentId);
-            if (!harnessIdentity)
+            if (!sessionIdentity)
                 saveAgentPreference({ agentId: resolved.agentId, tenantId: resolved.tenantId || tenantId });
             return { agentId: resolved.agentId, tenantId: resolved.tenantId || tenantId };
         },
@@ -1022,7 +990,7 @@ async function main() {
     }
     /** Identity hydration, prompt refresh, and auto-handoff — all post-auth. */
     async function runPostAuth(credentials) {
-        const c = { ...credentials, agentId: harnessIdentity?.agentId ?? credentials.agentId };
+        const c = { ...credentials, agentId: sessionIdentity?.agentId ?? credentials.agentId };
         // Auto-enroll a DID key so re-auth stays silent after the 7-day bond
         // token expires. Without this, key enrollment was a manual opt-in step
         // (`enable-key-auth`) and every normal user was re-prompted to bond at
@@ -1050,7 +1018,7 @@ async function main() {
             setupAutoHandoff(AIS_URL, c.agentId, c.token, activityTracker, () => {
                 stopRefreshTimer();
                 stderr('Session handoff saved');
-            }, () => verifyPinnedIdentity(harnessIdentity, c.agentId, () => aisGet(`/v1/agents/${c.agentId}`, c.token)));
+            }, () => verifyPinnedIdentity(sessionIdentity, c.agentId, () => aisGet(`/v1/agents/${c.agentId}`, c.token)));
         }
         catch (e) {
             stderr(`Post-auth setup failed: ${e instanceof Error ? e.message : String(e)}`);
@@ -1074,259 +1042,25 @@ async function main() {
                 }],
         };
     });
-    server.setRequestHandler(CallToolRequestSchema, async (request) => {
-        const { name, arguments: args } = request.params;
-        const rawArgs = args ?? {};
-        activityTracker.recordToolUse(name);
-        const toolStartedAt = Date.now();
-        try {
-            if (requiresIdentityGate(name, rawArgs)) {
-                assertPinnedAgent(harnessIdentity, harnessIdentity?.agentId ?? 'unselected');
-            }
-            const UNKNOWN_TOOL = Symbol('unknown-tool');
-            const dispatchTool = async (c) => {
-                const token = c.token;
-                // agentId is mutable across `agent_load` calls within this session.
-                let agentId = harnessIdentity?.agentId ?? c.agentId;
-                if (requiresIdentityGate(name, rawArgs)) {
-                    await verifyPinnedIdentity(harnessIdentity, agentId, () => aisGet(`/v1/agents/${agentId}`, token));
-                }
-                let result;
-                switch (name) {
-                    case 'remember': {
-                        const rememberArgs = parseRememberArgs(rawArgs);
-                        result = await aisPost(`/v1/agents/${agentId}/memory`, {
-                            content: rememberArgs.content,
-                            type: rememberArgs.type,
-                            importance: rememberArgs.importance,
-                        }, token);
-                        activityTracker.recordMemoryStored(rememberArgs.content);
-                        break;
-                    }
-                    case 'recall': {
-                        const recallArgs = parseRecallArgs(rawArgs);
-                        const params = new URLSearchParams();
-                        params.set('query', recallArgs.query);
-                        if (recallArgs.limit !== undefined)
-                            params.set('limit', String(recallArgs.limit));
-                        if (recallArgs.type !== undefined)
-                            params.set('types', recallArgs.type);
-                        if (recallArgs.principalAgentId !== undefined) {
-                            params.set('principalAgentId', recallArgs.principalAgentId);
-                        }
-                        result = await aisGet(`/v1/agents/${agentId}/memory?${params}`, token);
-                        break;
-                    }
-                    case 'whoami':
-                        parseEmptyToolArgs(rawArgs);
-                        result = await aisGet(`/v1/agents/${agentId}`, token);
-                        break;
-                    case 'full_context': {
-                        const fcArgs = parseFullContextArgs(rawArgs);
-                        const fcParams = new URLSearchParams();
-                        if (fcArgs.principalAgentId !== undefined)
-                            fcParams.set('principalAgentId', fcArgs.principalAgentId);
-                        if (fcArgs.memoriesLimit !== undefined)
-                            fcParams.set('memoriesLimit', String(fcArgs.memoriesLimit));
-                        if (fcArgs.view !== undefined)
-                            fcParams.set('view', fcArgs.view);
-                        const fcQs = fcParams.toString();
-                        const fcPath = `/v1/agents/${agentId}/full-context${fcQs ? `?${fcQs}` : ''}`;
-                        result = await aisGet(fcPath, token);
-                        break;
-                    }
-                    case 'dream':
-                        parseEmptyToolArgs(rawArgs);
-                        result = await aisPost(`/v1/agents/${agentId}/dream`, {}, token);
-                        break;
-                    case 'handoff': {
-                        const handoffArgs = parseHandoffArgs(rawArgs);
-                        result = await aisPost(`/v1/agents/${agentId}/handoff`, {
-                            summary: handoffArgs.summary ?? 'Session ended',
-                            keyLearnings: handoffArgs.keyLearnings ?? [],
-                        }, token);
-                        break;
-                    }
-                    case 'agent_load': {
-                        const { query, memoriesLimit } = parseAgentLoadArgs(rawArgs);
-                        const target = await resolveAgentForLoad(token, query);
-                        if (harnessIdentity)
-                            assertPinnedAgent(harnessIdentity, target.id);
-                        persistClaudeSessionAgentName(process.env.CLAUDE_SESSION_ID, target.name);
-                        agentId = target.id;
-                        // Persist the switch so subsequent tool calls use the new agent.
-                        // saveAgentPreference() MUST be called here so that a process
-                        // restart (silent MCP reconnect) restores the switched agent from
-                        // disk instead of silently reverting to the default
-                        // (CORBOT-497C7FA3). Ids only — no token is ever written to disk.
-                        // sessionAgentOverride MUST also be set so that a key-enrolled
-                        // TokenManager doesn't clobber this switch on the very next
-                        // ensureCredentials() call (the manager resolves identity once
-                        // and never re-resolves it).
-                        if (creds) {
-                            creds = { ...creds, agentId };
-                            sessionAgentOverride = { agentId, tenantId: creds.tenantId };
-                            if (!harnessIdentity)
-                                saveAgentPreference({ agentId, tenantId: creds.tenantId });
-                        }
-                        agentContext = await hydrateIdentity(AIS_URL, agentId, token, memoriesLimit);
-                        applyMemoryRefresh(agentId, token);
-                        server.notification({ method: 'notifications/prompts/list_changed' });
-                        stderr(`agent_load: switched to "${target.name}" (${target.id})`);
-                        result = {
-                            success: true,
-                            agentId: target.id,
-                            name: target.name,
-                            did: agentContext?.identity.did ?? null,
-                            status: target.status,
-                            memoriesLoaded: agentContext?.recentMemories.length ?? 0,
-                            goalsLoaded: agentContext?.goals.length ?? 0,
-                        };
-                        break;
-                    }
-                    case 'inbox_send': {
-                        const inboxSendArgs = parseInboxSendArgs(rawArgs);
-                        result = await aisPost(`/v1/agents/${agentId}/inbox`, {
-                            recipientAgentId: inboxSendArgs.recipientAgentId,
-                            kind: inboxSendArgs.kind,
-                            body: inboxSendArgs.body,
-                            ...(inboxSendArgs.threadId !== undefined ? { threadId: inboxSendArgs.threadId } : {}),
-                            ...(inboxSendArgs.taskId !== undefined ? { taskId: inboxSendArgs.taskId } : {}),
-                            ...(inboxSendArgs.inReplyTo !== undefined ? { inReplyTo: inboxSendArgs.inReplyTo } : {}),
-                            ...(inboxSendArgs.messageId !== undefined ? { messageId: inboxSendArgs.messageId } : {}),
-                            ...(inboxSendArgs.ttlSeconds !== undefined ? { ttlSeconds: inboxSendArgs.ttlSeconds } : {}),
-                        }, token);
-                        break;
-                    }
-                    case 'inbox_read': {
-                        const inboxReadArgs = parseInboxReadArgs(rawArgs);
-                        const inboxParams = new URLSearchParams();
-                        if (inboxReadArgs.cursor !== undefined)
-                            inboxParams.set('cursor', inboxReadArgs.cursor);
-                        inboxParams.set('limit', String(inboxReadArgs.limit));
-                        inboxParams.set('unreadOnly', String(inboxReadArgs.unreadOnly));
-                        result = await readInboxAndMaybeAck(inboxReadArgs, () => aisGet(`/v1/agents/${agentId}/inbox?${inboxParams}`, token), (messageIds) => aisPost(`/v1/agents/${agentId}/inbox/ack`, { messageIds }, token));
-                        break;
-                    }
-                    case 'inbox_ack': {
-                        const inboxAckArgs = parseInboxAckArgs(rawArgs);
-                        result = await aisPost(`/v1/agents/${agentId}/inbox/ack`, { messageIds: inboxAckArgs.messageIds }, token);
-                        break;
-                    }
-                    case 'goal_create': {
-                        const goalCreateArgs = parseGoalCreateArgs(rawArgs);
-                        result = await aisPost(`/v1/agents/${agentId}/goals`, {
-                            title: goalCreateArgs.title,
-                            description: goalCreateArgs.description,
-                            ...(goalCreateArgs.successCriteria !== undefined ? { successCriteria: goalCreateArgs.successCriteria } : {}),
-                            ...(goalCreateArgs.parentGoalId !== undefined ? { parentGoalId: goalCreateArgs.parentGoalId } : {}),
-                        }, token);
-                        break;
-                    }
-                    case 'goal_list': {
-                        const goalListArgs = parseGoalListArgs(rawArgs);
-                        const goalParams = new URLSearchParams();
-                        if (goalListArgs.status !== undefined)
-                            goalParams.set('status', goalListArgs.status);
-                        if (goalListArgs.page !== undefined)
-                            goalParams.set('page', String(goalListArgs.page));
-                        goalParams.set('limit', String(goalListArgs.limit));
-                        if (goalListArgs.parentGoalId !== undefined)
-                            goalParams.set('parentGoalId', goalListArgs.parentGoalId);
-                        result = await aisGet(`/v1/agents/${agentId}/goals?${goalParams}`, token);
-                        break;
-                    }
-                    case 'goal_update_progress': {
-                        const goalProgressArgs = parseGoalUpdateProgressArgs(rawArgs);
-                        result = await aisPatch(`/v1/agents/${agentId}/goals/${goalProgressArgs.goalId}`, {
-                            progress: goalProgressArgs.progress,
-                        }, token);
-                        break;
-                    }
-                    case 'goal_complete': {
-                        const goalCompleteArgs = parseGoalCompleteArgs(rawArgs);
-                        result = await aisPost(`/v1/agents/${agentId}/goals/${goalCompleteArgs.goalId}/complete`, {
-                            success: goalCompleteArgs.success,
-                            ...(goalCompleteArgs.notes !== undefined ? { notes: goalCompleteArgs.notes } : {}),
-                            ...(goalCompleteArgs.lessonsLearned !== undefined ? { lessonsLearned: goalCompleteArgs.lessonsLearned } : {}),
-                        }, token);
-                        break;
-                    }
-                    default:
-                        emitTelemetry({
-                            event: 'tool.call',
-                            tool: name,
-                            agentIdHash: hashAgentId(agentId),
-                            latencyMs: Date.now() - toolStartedAt,
-                            status: 'error',
-                            error: `Unknown tool: ${name}`,
-                        });
-                        return UNKNOWN_TOOL;
-                }
-                return result;
-            };
-            let result = await dispatchTool(await ensureCredentials());
-            if (isAuthErrorResult(result)) {
-                // CORBOT-A5BBD580: the server rejected our token mid-session
-                // (expired or secret rotated). Clear the cached creds, silently
-                // re-auth (key auth when a key file exists), and retry once.
-                stderr('AIS rejected the cached token; re-authenticating and retrying once...');
-                creds = null;
-                // Drop the TokenManager's cached token too — without this, the
-                // key-enrolled path's ensureCredentials() would just re-serve the
-                // same still-time-valid-but-server-rejected token from `current`
-                // instead of re-running challenge/prove (CORBOT-A5BBD580: a JWT
-                // rejection mid-session, e.g. AIS rotating JWT_SECRET, must force a
-                // real re-auth, not a no-op retry with the same token). Also reset
-                // any escalated backoff window (CORBOT-4F7CAA8E / F2b): the server
-                // just responded with a 401, so it is reachable right now — a stale
-                // backoff window computed from unrelated prior renewal failures must
-                // not hold this re-prove behind the getCredentials() gate.
-                tokenManager.invalidate({ resetBackoff: true });
-                result = await dispatchTool(await ensureCredentials());
-            }
-            if (result === UNKNOWN_TOOL) {
-                return {
-                    content: [{ type: 'text', text: JSON.stringify({ error: `Unknown tool: ${name}` }) }],
-                    isError: true,
-                };
-            }
-            emitTelemetry({
-                event: 'tool.call',
-                tool: name,
-                agentIdHash: creds ? hashAgentId(creds.agentId) : undefined,
-                latencyMs: Date.now() - toolStartedAt,
-                status: 'ok',
-            });
-            return {
-                content: [{ type: 'text', text: JSON.stringify(result, null, 2) }],
-            };
-        }
-        catch (error) {
-            const message = error instanceof Error ? error.message : 'Unknown error';
-            const agentHash = creds ? hashAgentId(creds.agentId) : undefined;
-            emitTelemetry({
-                event: 'tool.call',
-                tool: name,
-                agentIdHash: agentHash,
-                latencyMs: Date.now() - toolStartedAt,
-                status: 'error',
-                error: message,
-            });
-            // BondPendingError carries human-readable text that the agent should
-            // render verbatim — return as plain text rather than JSON.
-            if (error instanceof BondPendingError) {
-                return {
-                    content: [{ type: 'text', text: message }],
-                    isError: true,
-                };
-            }
-            return {
-                content: [{ type: 'text', text: JSON.stringify({ error: message }) }],
-                isError: true,
-            };
-        }
-    });
+    server.setRequestHandler(CallToolRequestSchema, createOAuthCallToolHandler({
+        aisUrl: AIS_URL,
+        server,
+        activityTracker,
+        getSessionIdentity: () => sessionIdentity,
+        setSessionIdentity: (next) => { sessionIdentity = next; },
+        getCredentials: () => creds,
+        setCredentials: (next) => { creds = next; },
+        setSessionAgentOverride: (next) => { sessionAgentOverride = next; },
+        getAgentContext: () => agentContext,
+        setAgentContext: (next) => { agentContext = next; },
+        ensureCredentials: () => ensureCredentials(),
+        invalidateToken: () => { tokenManager.invalidate({ resetBackoff: true }); },
+        applyMemoryRefresh,
+        resolveAgentForLoad,
+        persistClaudeSessionAgentName,
+        isBondPendingError: (error) => error instanceof BondPendingError,
+        stderr,
+    }));
     // Connect the transport BEFORE any auth so MCP `initialize` returns in <1s.
     const transport = new StdioServerTransport();
     await server.connect(transport);
@@ -1349,6 +1083,9 @@ async function main() {
     else if (creds) {
         stderr('Device key enrolled — deferring session setup to the first key-auth token acquire.');
     }
+    else if (hasEnrolledKey()) {
+        stderr('Using enrolled device key; session starts on first tool call.');
+    }
     else {
         stderr(`${MISSING_CREDENTIAL_HELP} Bond URL will be surfaced on first tool call.`);
     }