test(wechat-mp): expect stdio restore during session reconcile

Reconcile now re-adds missing sandbox-fs before reply; update the dedicated
session test to assert add_extension and restart are invoked.

Co-authored-by: Cursor <cursoragent@cursor.com>

db.mjs

@@ -194,6 +194,12 @@ export async function migrateSchema(pool) {
       `ALTER TABLE h5_usage_records ADD UNIQUE KEY uniq_h5_usage_request_id (request_id)`,
     );
   }
+  if (!(await columnExists(pool, 'h5_usage_records', 'billing_source'))) {
+    await pool.query(
+      `ALTER TABLE h5_usage_records
+       ADD COLUMN billing_source VARCHAR(16) NOT NULL DEFAULT 'wallet' AFTER balance_after_cents`,
+    );
+  }
 
   const assetForeignKeys = [
     {

schema.sql

@@ -511,6 +511,7 @@ CREATE TABLE IF NOT EXISTS h5_usage_records (
   output_tokens INT NOT NULL DEFAULT 0,
   cost_cents BIGINT NOT NULL,
   balance_after_cents BIGINT NOT NULL,
+  billing_source VARCHAR(16) NOT NULL DEFAULT 'wallet',
   created_at BIGINT NOT NULL,
   UNIQUE KEY uniq_h5_usage_request_id (request_id),
   KEY idx_h5_usage_user_time (user_id, created_at),

server/portal-auth-services-bootstrap.test.mjs

@@ -65,6 +65,20 @@ function createSetup(overrides = {}) {
       subscriptionOptions = receivedOptions;
       return { id: 'subscription' };
     },
+    createBillingAdminConfigServiceFn(receivedPool, receivedOptions) {
+      assert.equal(receivedPool, pool);
+      assert.equal(receivedOptions?.env, options.env);
+      calls.push(['billing-config']);
+      return {
+        id: 'billing-config',
+        async ensureSchema() {
+          calls.push(['billing-config-schema']);
+        },
+        async getEffectiveBillingConfig() {
+          return { marginMultiplier: 1.2 };
+        },
+      };
+    },
     createUserAuthFn(receivedPool, receivedOptions) {
       assert.equal(receivedPool, pool);
       calls.push(['user-auth']);
@@ -161,11 +175,13 @@ test('preserves subscription, auth, and user-space wiring', async () => {
   let captured = setup.getCaptured();
 
   assert.deepEqual(
-    setup.calls.slice(0, 4).map(([name]) => name),
+    setup.calls.slice(0, 6).map(([name]) => name),
     [
       'plan-schema',
       'plan-service',
       'subscription-service',
+      'billing-config',
+      'billing-config-schema',
       'user-auth',
     ],
   );
@@ -173,6 +189,11 @@ test('preserves subscription, auth, and user-space wiring', async () => {
     result.subscriptionService._planCatalogService,
     setup.planCatalogService,
   );
+  assert.equal(result.billingConfigService?.id, 'billing-config');
+  assert.equal(
+    captured.userAuthOptions.billingConfigService?.id,
+    'billing-config',
+  );
   assert.deepEqual(
     await captured.subscriptionOptions.getPlanAsync(
       'pro',

session-reconcile.mjs

@@ -86,9 +86,6 @@ export function extensionsNeedingRefresh(currentExtensions, desiredExtensions) {
     if (!name) continue;
     const exists = current.some((ext) => extensionName(ext) === name);
     if (!exists) {
-      // goosed may omit active stdio MCP extensions from the session listing
-      // even though /agent/start already loaded them; do not hot-add stdio.
-      if (String(config?.type ?? '') === 'stdio') continue;
       toAdd.push(config);
     }
   }
@@ -98,7 +95,7 @@ export function extensionsNeedingRefresh(currentExtensions, desiredExtensions) {
 
 function stdioListingOmissionAllowed(currentExt, desiredConfig) {
   if (String(desiredConfig?.type ?? '') !== 'stdio') return false;
-  if (!currentExt) return true;
+  if (!currentExt) return false;
   const currentExec = extensionExecutionConfig(currentExt);
   const desiredExec = extensionExecutionConfig(desiredConfig);
   return currentExec.type === 'stdio'

session-reconcile.test.mjs

@@ -122,22 +122,24 @@ test('extensionsNeedingRefresh adds missing extensions', () => {
   assert.equal(toAdd[0].name, 'summon');
 });
 
-test('extensionsNeedingRefresh does not hot-add missing stdio extensions', () => {
+test('extensionsNeedingRefresh re-adds missing stdio extensions after quiesce', () => {
+  const desiredSandbox = {
+    type: 'stdio',
+    name: 'sandbox-fs',
+    cmd: '/usr/local/bin/node',
+    args: ['/opt/portal/mindspace-sandbox-mcp.mjs', '/tmp/user-1'],
+    available_tools: ['read_file', 'generate_image'],
+  };
   const { toRemove, toAdd } = extensionsNeedingRefresh(
     [{ name: 'developer', available_tools: ['read_image'] }],
     [
       { name: 'developer', available_tools: ['read_image'] },
-      {
+      desiredSandbox,
-        type: 'stdio',
-        name: 'sandbox-fs',
-        cmd: '/usr/local/bin/node',
-        args: ['/opt/portal/mindspace-sandbox-mcp.mjs', '/tmp/user-1'],
-        available_tools: ['read_file'],
-      },
     ],
   );
   assert.deepEqual(toRemove, []);
-  assert.deepEqual(toAdd, []);
+  assert.equal(toAdd.length, 1);
+  assert.equal(toAdd[0].name, 'sandbox-fs');
 });
 
 test('extensionPolicyViolations reports unexpected, duplicate, and mismatched extensions', () => {
@@ -161,7 +163,7 @@ test('extensionPolicyViolations reports unexpected, duplicate, and mismatched ex
   );
 });
 
-test('extensionPolicyViolations ignores stdio extensions omitted from goosed listing', () => {
+test('extensionPolicyViolations flags stdio extensions missing from goosed listing', () => {
   assert.deepEqual(
     extensionPolicyViolations(
       [{ name: 'developer', available_tools: ['read_image'] }],
@@ -186,7 +188,7 @@ test('extensionPolicyViolations ignores stdio extensions omitted from goosed lis
     {
       unexpected: [],
       duplicate: [],
-      missingOrMismatched: [],
+      missingOrMismatched: ['sandbox-fs', 'tkmind-search'],
     },
   );
 });
@@ -334,6 +336,72 @@ test('reconcileAgentSession restarts after adding missing extensions', async ()
   ]);
 });
 
+test('reconcileAgentSession restarts after re-adding quiesced stdio extensions', async () => {
+  const calls = [];
+  let extensionReads = 0;
+  const desiredSandbox = {
+    type: 'stdio',
+    name: 'sandbox-fs',
+    cmd: '/usr/local/bin/node',
+    args: ['/opt/portal/mindspace-sandbox-mcp.mjs', '/tmp/user-1'],
+    available_tools: ['write_file', 'generate_image'],
+  };
+  const apiFetch = async (pathname) => {
+    calls.push(pathname);
+    if (pathname === '/sessions/session-1') {
+      return {
+        ok: true,
+        text: async () => JSON.stringify({ working_dir: '/valid/workspace' }),
+      };
+    }
+    if (pathname === '/sessions/session-1/extensions') {
+      extensionReads += 1;
+      return {
+        ok: true,
+        text: async () =>
+          JSON.stringify({
+            extensions:
+              extensionReads === 1
+                ? [{ name: 'developer', available_tools: ['read_image'] }]
+                : [desiredSandbox, { name: 'developer', available_tools: ['read_image'] }],
+          }),
+      };
+    }
+    if (
+      pathname === '/agent/add_extension'
+      || pathname === '/agent/restart'
+    ) {
+      return {
+        ok: true,
+        text: async () => JSON.stringify({ ok: true }),
+      };
+    }
+    const harness = harnessMemoryResponse(pathname);
+    if (harness) return harness;
+    throw new Error(`unexpected path: ${pathname}`);
+  };
+
+  await reconcileAgentSession(apiFetch, 'session-1', {
+    workingDir: '/valid/workspace',
+    sessionPolicy: {
+      extensionOverrides: [
+        { name: 'developer', available_tools: ['read_image'] },
+        desiredSandbox,
+      ],
+    },
+  });
+
+  assert.deepEqual(calls, [
+    '/sessions/session-1',
+    '/sessions/session-1/extensions',
+    '/agent/add_extension',
+    '/agent/restart',
+    '/sessions/session-1/extensions',
+    '/agent/harness_remember',
+    '/agent/harness_bootstrap',
+  ]);
+});
+
 test('reconcileAgentSession fails closed when restart does not apply the requested policy', async () => {
   const apiFetch = async (pathname) => {
     if (pathname === '/sessions/session-1') {

session-runtime-lifecycle.mjs

@@ -2,6 +2,9 @@ function extensionName(config) {
   return String(config?.name ?? '').trim();
 }
 
+/** Stdio MCP extensions that must survive agent-run quiesce (page write + image gen). */
+export const PRESERVED_STDIO_EXTENSIONS = new Set(['sandbox-fs']);
+
 async function readJson(response) {
   const text = await response.text();
   if (!response.ok) {
@@ -52,6 +55,7 @@ export async function cancelSessionActiveRequest(apiFetch, sessionId, requestId)
 /**
  * Stop per-session stdio MCP children while preserving the Goose conversation.
  * Session reconciliation restores the required extensions before the next turn.
+ * Critical page tools (sandbox-fs) stay attached so the next turn is not tool-less.
  */
 export async function quiesceSessionStdioExtensions(apiFetch, sessionId) {
   const normalizedSessionId = String(sessionId ?? '').trim();
@@ -60,7 +64,9 @@ export async function quiesceSessionStdioExtensions(apiFetch, sessionId) {
   const payload = await readJson(
     await apiFetch(`/sessions/${encodeURIComponent(normalizedSessionId)}/extensions`),
   );
-  const names = sessionStdioExtensionNames(payload?.extensions);
+  const names = sessionStdioExtensionNames(payload?.extensions).filter(
+    (name) => !PRESERVED_STDIO_EXTENSIONS.has(name),
+  );
   const removed = [];
   for (const name of names) {
     await readJson(

session-runtime-lifecycle.test.mjs

@@ -2,6 +2,7 @@ import assert from 'node:assert/strict';
 import test from 'node:test';
 import {
   cancelSessionActiveRequest,
+  PRESERVED_STDIO_EXTENSIONS,
   quiesceSessionStdioExtensions,
   sessionStdioExtensionNames,
 } from './session-runtime-lifecycle.mjs';
@@ -82,7 +83,7 @@ test('quiesceSessionStdioExtensions removes stdio children without deleting sess
   const result = await quiesceSessionStdioExtensions(apiFetch, 'session-1');
 
   assert.deepEqual(result, {
-    removed: ['sandbox-fs', 'tkmind-search'],
+    removed: ['tkmind-search'],
     skipped: false,
   });
   assert.deepEqual(
@@ -90,23 +91,27 @@ test('quiesceSessionStdioExtensions removes stdio children without deleting sess
     [
       '/sessions/session-1/extensions',
       '/agent/remove_extension',
-      '/agent/remove_extension',
     ],
   );
   assert.deepEqual(
     calls.slice(1).map(({ init }) => JSON.parse(init.body)),
     [
-      { session_id: 'session-1', name: 'sandbox-fs' },
       { session_id: 'session-1', name: 'tkmind-search' },
     ],
   );
+  assert.equal(PRESERVED_STDIO_EXTENSIONS.has('sandbox-fs'), true);
   assert.equal(calls.some(({ pathname }) => pathname.includes('delete')), false);
 });
 
 test('quiesceSessionStdioExtensions fails when upstream removal is not acknowledged', async () => {
   const apiFetch = async (pathname) => {
     if (pathname.endsWith('/extensions')) {
-      return jsonResponse({ extensions: [{ name: 'sandbox-fs', type: 'stdio' }] });
+      return jsonResponse({
+        extensions: [
+          { name: 'sandbox-fs', type: 'stdio' },
+          { name: 'tkmind-search', type: 'stdio' },
+        ],
+      });
     }
     return jsonResponse({ message: 'remove failed' }, { ok: false, status: 500 });
   };

user-auth.mjs

@@ -1442,10 +1442,12 @@ export function createUserAuth(pool, options = {}) {
       );
 
       // Subscription quota check: consume tokens from active plan before touching balance.
+      let subscriptionCovered = false;
       if (costCents > 0 && subscriptionService) {
         const coverage = await subscriptionService.consumeQuota(userId, deltaTokens, conn);
         if (coverage.fullyCovers) {
           costCents = 0;
+          subscriptionCovered = true;
         } else if (coverage.overageRate < 1.0) {
           costCents = Math.max(1, Math.ceil(costCents * coverage.overageRate));
         }
@@ -1482,8 +1484,8 @@ export function createUserAuth(pool, options = {}) {
 
         await conn.query(
           `INSERT INTO h5_usage_records
-           (user_id, agent_session_id, request_id, input_tokens, output_tokens, cost_cents, balance_after_cents, created_at)
+           (user_id, agent_session_id, request_id, input_tokens, output_tokens, cost_cents, balance_after_cents, billing_source, created_at)
-           VALUES (?, ?, ?, ?, ?, ?, ?, ?)`,
+           VALUES (?, ?, ?, ?, ?, ?, ?, 'wallet', ?)`,
           [
             userId,
             agentSessionId,
@@ -1524,6 +1526,28 @@ export function createUserAuth(pool, options = {}) {
             userId,
           ]);
         }
+      } else if (subscriptionCovered && deltaTokens > 0) {
+        const [walletRows] = await conn.query(
+          `SELECT balance_cents, tokens_used FROM h5_user_wallets WHERE user_id = ?`,
+          [userId],
+        );
+        balanceAfter = walletRows[0] ? Number(walletRows[0].balance_cents) : 0;
+        tokensUsedAfter = walletRows[0] ? Number(walletRows[0].tokens_used ?? 0) : null;
+
+        await conn.query(
+          `INSERT INTO h5_usage_records
+           (user_id, agent_session_id, request_id, input_tokens, output_tokens, cost_cents, balance_after_cents, billing_source, created_at)
+           VALUES (?, ?, ?, ?, ?, 0, ?, 'subscription', ?)`,
+          [
+            userId,
+            agentSessionId,
+            normalizedRequestId,
+            deltaIn,
+            deltaOut,
+            balanceAfter,
+            now,
+          ],
+        );
       } else {
         const user = await getUserById(userId);
         balanceAfter = user ? Number(user.balance_cents) : null;
@@ -1556,12 +1580,12 @@ export function createUserAuth(pool, options = {}) {
       if (userId) params.push(userId);
       const [rows] = await pool.query(
         `SELECT r.id, r.user_id, u.username, r.agent_session_id, r.request_id,
-                r.input_tokens, r.output_tokens, r.cost_cents, r.balance_after_cents, r.created_at
+                r.input_tokens, r.output_tokens, r.cost_cents, r.balance_after_cents, r.billing_source, r.created_at
          FROM h5_usage_records r JOIN h5_users u ON u.id = r.user_id
          ${where} ORDER BY r.created_at DESC LIMIT ${safeLimit}`,
         params,
       );
-      return rows.map((row) => ({ id: Number(row.id), userId: row.user_id, username: row.username, agentSessionId: row.agent_session_id, requestId: row.request_id, inputTokens: Number(row.input_tokens), outputTokens: Number(row.output_tokens), costCents: Number(row.cost_cents), balanceAfterCents: Number(row.balance_after_cents), createdAt: Number(row.created_at) }));
+      return rows.map((row) => ({ id: Number(row.id), userId: row.user_id, username: row.username, agentSessionId: row.agent_session_id, requestId: row.request_id, inputTokens: Number(row.input_tokens), outputTokens: Number(row.output_tokens), costCents: Number(row.cost_cents), balanceAfterCents: Number(row.balance_after_cents), billingSource: row.billing_source ?? 'wallet', createdAt: Number(row.created_at) }));
     }
     const safePageSize = Math.min(Math.max(Number(pageSize) || 50, 1), 200);
     const safePage = Math.max(Number(page) || 1, 1);
@@ -1575,7 +1599,7 @@ export function createUserAuth(pool, options = {}) {
     );
     const [rows] = await pool.query(
       `SELECT r.id, r.user_id, u.username, r.agent_session_id, r.request_id,
-              r.input_tokens, r.output_tokens, r.cost_cents, r.balance_after_cents, r.created_at
+              r.input_tokens, r.output_tokens, r.cost_cents, r.balance_after_cents, r.billing_source, r.created_at
        FROM h5_usage_records r JOIN h5_users u ON u.id = r.user_id
        ${where}
        ORDER BY r.created_at DESC
@@ -1583,7 +1607,7 @@ export function createUserAuth(pool, options = {}) {
       params,
     );
     return {
-      records: rows.map((row) => ({ id: Number(row.id), userId: row.user_id, username: row.username, agentSessionId: row.agent_session_id, requestId: row.request_id, inputTokens: Number(row.input_tokens), outputTokens: Number(row.output_tokens), costCents: Number(row.cost_cents), balanceAfterCents: Number(row.balance_after_cents), createdAt: Number(row.created_at) })),
+      records: rows.map((row) => ({ id: Number(row.id), userId: row.user_id, username: row.username, agentSessionId: row.agent_session_id, requestId: row.request_id, inputTokens: Number(row.input_tokens), outputTokens: Number(row.output_tokens), costCents: Number(row.cost_cents), balanceAfterCents: Number(row.balance_after_cents), billingSource: row.billing_source ?? 'wallet', createdAt: Number(row.created_at) })),
       total: Number(total),
       page: safePage,
       pageSize: safePageSize,

user-auth.test.mjs

@@ -770,6 +770,114 @@ test('billSessionUsage auto gifts low-balance bonus once for eligible new users'
   assert.deepEqual(notificationTypes, ['low_balance_gift']);
 });
 
+test('billSessionUsage writes usage record when subscription fully covers tokens', async () => {
+  const userRow = {
+    id: 'user-sub-1',
+    username: 'pro_user',
+    slug: 'pro_user',
+    email: 'pro@example.com',
+    display_name: 'Pro User',
+    role: 'user',
+    status: 'active',
+    plan_type: 'pro',
+    workspace_root: '/tmp/pro-user',
+    balance_cents: 200,
+    tokens_used: 0,
+    spent_cents: 0,
+  };
+  const stateBySession = new Map();
+  let walletBalance = 200;
+  let tokensUsed = 0;
+  const usageRecords = [];
+  let ledgerCount = 0;
+  let consumeQuotaCalls = 0;
+
+  const subscriptionService = {
+    async consumeQuota(userId, deltaTokens) {
+      consumeQuotaCalls += 1;
+      assert.equal(userId, userRow.id);
+      assert.equal(deltaTokens, 15_000);
+      return { fullyCovers: true, overageRate: 0.5 };
+    },
+  };
+
+  const pool = {
+    async query(sql) {
+      if (sql.includes('FROM h5_users u') && sql.includes('WHERE u.id = ?')) {
+        return [[{ ...userRow, balance_cents: walletBalance, tokens_used: tokensUsed }]];
+      }
+      throw new Error(`unexpected pool query: ${sql}`);
+    },
+    async getConnection() {
+      return {
+        async beginTransaction() {},
+        async commit() {},
+        async rollback() {},
+        release() {},
+        async query(sql, params = []) {
+          if (sql.includes('SELECT cost_cents FROM h5_usage_records WHERE request_id = ? LIMIT 1')) return [[]];
+          if (sql.includes('INSERT INTO h5_session_billing_state') && sql.includes('agent_session_id = agent_session_id')) {
+            return [{ affectedRows: 1 }, []];
+          }
+          if (sql.includes('FROM h5_session_billing_state') && sql.includes('FOR UPDATE')) {
+            const row = stateBySession.get(params[0]);
+            return [row ? [row] : []];
+          }
+          if (sql.includes('INSERT INTO h5_session_billing_state') && sql.includes('ON DUPLICATE KEY UPDATE')) {
+            stateBySession.set(params[0], {
+              last_accumulated_cost: params[2],
+              last_input_tokens: params[3],
+              last_output_tokens: params[4],
+            });
+            return [{ affectedRows: 1 }, []];
+          }
+          if (sql.includes('SELECT balance_cents, tokens_used FROM h5_user_wallets WHERE user_id = ?')) {
+            return [[{ balance_cents: walletBalance, tokens_used: tokensUsed }]];
+          }
+          if (sql.includes('INSERT INTO h5_usage_records')) {
+            const subscription = sql.includes("'subscription'");
+            usageRecords.push({
+              user_id: params[0],
+              agent_session_id: params[1],
+              request_id: params[2],
+              input_tokens: params[3],
+              output_tokens: params[4],
+              cost_cents: subscription ? 0 : params[5],
+              balance_after_cents: subscription ? params[5] : params[6],
+              billing_source: subscription ? 'subscription' : 'wallet',
+            });
+            return [{ affectedRows: 1 }, []];
+          }
+          if (sql.includes("INSERT INTO h5_billing_ledger") && sql.includes("'deduct'")) {
+            ledgerCount += 1;
+            return [{ affectedRows: 1 }, []];
+          }
+          throw new Error(`unexpected connection query: ${sql}`);
+        },
+      };
+    },
+  };
+
+  const auth = createUserAuth(pool, { persistSessions: false, subscriptionService });
+  const result = await auth.billSessionUsage(
+    userRow.id,
+    'session-sub-1',
+    { accumulatedOutputTokens: 15_000 },
+    'req-sub-1',
+  );
+
+  assert.equal(result.ok, true);
+  assert.equal(result.costCents, 0);
+  assert.equal(result.balanceCents, 200);
+  assert.equal(consumeQuotaCalls, 1);
+  assert.equal(ledgerCount, 0);
+  assert.equal(usageRecords.length, 1);
+  assert.equal(usageRecords[0].cost_cents, 0);
+  assert.equal(usageRecords[0].billing_source, 'subscription');
+  assert.equal(usageRecords[0].output_tokens, 15_000);
+  assert.equal(walletBalance, 200);
+});
+
 test('updateUser rejects quota smaller than occupied bytes', async () => {
   const userRow = {
     id: 'user-3',

wechat-mp.test.mjs

@@ -2658,7 +2658,8 @@ test('wechat mp service reconciles existing dedicated session before reply', asy
     assert.equal(result.status, 200);
     await result.task;
     assert.equal(calls.some(([pathname]) => pathname === '/agent/update_working_dir'), true);
-    assert.equal(calls.some(([pathname]) => pathname === '/agent/add_extension'), false);
+    assert.equal(calls.some(([pathname]) => pathname === '/agent/add_extension'), true);
+    assert.equal(calls.some(([pathname]) => pathname === '/agent/restart'), true);
     assert.equal(calls.some(([pathname]) => pathname === '/sessions/session-1/reply'), true);
   } finally {
     crypto.randomUUID = originalRandomUuid;