diff --git a/.claude-plugin/marketplace.json b/.claude-plugin/marketplace.json index 1070950e9..11b63198a 100644 --- a/.claude-plugin/marketplace.json +++ b/.claude-plugin/marketplace.json @@ -11,7 +11,7 @@ "name": "memwal", "source": "./packages/mcp/plugin", "description": "Automatic Walrus Memory — proactive recall and durable-fact saving via the MemWal MCP + lifecycle hooks.", - "version": "0.0.13" + "version": "0.0.14" } ] } diff --git a/.cursor-plugin/marketplace.json b/.cursor-plugin/marketplace.json index a4b700e2b..35ec36556 100644 --- a/.cursor-plugin/marketplace.json +++ b/.cursor-plugin/marketplace.json @@ -11,7 +11,7 @@ "name": "memwal", "source": "./packages/mcp/plugin", "description": "Automatic Walrus Memory — proactive recall and durable-fact saving via the MemWal MCP + lifecycle hooks.", - "version": "0.0.13" + "version": "0.0.14" } ] } diff --git a/.github/ISSUE_TEMPLATE/bug.yml b/.github/ISSUE_TEMPLATE/bug.yml index 1a2d78b7b..46c92ebef 100644 --- a/.github/ISSUE_TEMPLATE/bug.yml +++ b/.github/ISSUE_TEMPLATE/bug.yml @@ -10,7 +10,7 @@ body: Do not paste delegate private keys, mnemonics, or other secrets. - If this is a security finding, use a [private advisory](https://github.com/MystenLabs/MemWal/security/advisories/new) instead of this form. + If this is a security finding, email [security@mystenlabs.com](mailto:security@mystenlabs.com). Do not use this form. See our [security policy](https://github.com/MystenLabs/MemWal/security/policy). - type: dropdown id: surface diff --git a/.github/ISSUE_TEMPLATE/config.yml b/.github/ISSUE_TEMPLATE/config.yml index a08bdb25e..933094dc6 100644 --- a/.github/ISSUE_TEMPLATE/config.yml +++ b/.github/ISSUE_TEMPLATE/config.yml @@ -10,5 +10,5 @@ contact_links: url: https://discord.gg/walrusprotocol about: Questions and builder support that are not a bug or a feature request. - name: Report a security vulnerability - url: https://github.com/MystenLabs/MemWal/security/advisories/new - about: Private advisory. Do not file a public issue for security findings. + url: https://github.com/MystenLabs/MemWal/security/policy + about: Email security@mystenlabs.com. Do not report security issues through GitHub or Discord. diff --git a/.github/workflows/check-mcp-pin.yml b/.github/workflows/check-mcp-pin.yml new file mode 100644 index 000000000..18da13e92 --- /dev/null +++ b/.github/workflows/check-mcp-pin.yml @@ -0,0 +1,64 @@ +name: Check MCP launcher pin + +# The plugin launcher installs whatever `mcpPackageVersion` in +# packages/mcp/plugin/plugin.json names, so a pin npm does not carry breaks every +# fresh plugin install until someone edits the manifest. Whether a version exists +# is a fact about npm rather than about this repository, so no unit test can hold +# it — trusted-launcher.test.mjs checks the shape of the pin and its relationship +# to `version`, and this checks that the thing actually exists. +# +# Its own workflow rather than a job in test.yml for the same reason as +# check-model-ids: the failure can arrive with no commit attached. The `dev` +# dist-tag moves whenever a push to `dev` touches packages/mcp, and that number is +# never committed back, so the manifest drifts on its own. test.yml has no +# schedule trigger. +# +# The registry endpoint is public, so this needs no token and no environment. + +on: + pull_request: + # Only when the pin, the check, or the launcher that reads the pin moves. + # Every other PR would be taking a network dependency on npm for nothing. + paths: + - "packages/mcp/plugin/plugin.json" + - "packages/mcp/plugin/scripts/lib/mcp-launch.mjs" + - "scripts/check-mcp-pin.mjs" + - ".github/workflows/check-mcp-pin.yml" + push: + branches: [main, staging, dev] + paths: + - "packages/mcp/plugin/plugin.json" + - "scripts/check-mcp-pin.mjs" + - ".github/workflows/check-mcp-pin.yml" + schedule: + # Catches the drift nobody pushed: a dist-tag that moved under a manifest + # that did not. Reads nothing shared, so it needs no bench-account offset. + - cron: "0 9 * * 1" + workflow_dispatch: + +concurrency: + group: check-mcp-pin-${{ github.ref }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + +permissions: + contents: read + +jobs: + mcp-pin: + name: Launcher pin is published + runs-on: ubuntu-latest + timeout-minutes: 5 + + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Setup Node + uses: actions/setup-node@v4 + with: + node-version: "22" + + # No install step: the check is dependency-free and talks to the public + # registry over fetch. + - name: Check the launcher pin against npm + run: node scripts/check-mcp-pin.mjs diff --git a/README.md b/README.md index 405e92d7f..1eac4c024 100644 --- a/README.md +++ b/README.md @@ -57,6 +57,8 @@ await memwal.restore("demo"); ## Contributing +Report security issues by email to [security@mystenlabs.com](mailto:security@mystenlabs.com). Do not report security issues through GitHub or Discord. See [SECURITY.md](SECURITY.md). + We want to be explicit about this while Walrus Memory is in beta: feedback, bug reports, docs fixes, examples, and implementation contributions are all welcome. diff --git a/SECURITY.md b/SECURITY.md new file mode 100644 index 000000000..01d09092e --- /dev/null +++ b/SECURITY.md @@ -0,0 +1,10 @@ +# Security Policy + +We appreciate your participation in keeping Walrus Memory secure. This document describes how to report +security vulnerabilities. + +## Reporting a Vulnerability + +Report security issues by email to +[security@mystenlabs.com](mailto:security@mystenlabs.com). +**Do not report security issues through GitHub or Discord.** diff --git a/apps/app/public/skills/setup b/apps/app/public/skills/setup index edecce6a6..a9433bb10 100644 --- a/apps/app/public/skills/setup +++ b/apps/app/public/skills/setup @@ -26,14 +26,13 @@ user explicitly asks for developer integration. ## Setup Rules -1. Identify the AI client first. A client that cannot run a local MCP command - but can attach two custom connector headers uses the Remote MCP section - below. Every other client (Claude Desktop, Claude Code, Cursor, Codex, - other local apps) uses the Local MCP Server section. ChatGPT web - connectors are not supported: chatgpt.com exposes only a single bearer - field and cannot supply the required `x-memwal-account-id` header. The - ChatGPT desktop app is a different product: it ships the Codex CLI and - reads `~/.codex`, so treat it as Codex and use the Codex section below. +1. Identify the AI client and mode before checking local prerequisites. Regular + Claude Desktop chat and Claude web use **Claude OAuth Connector** below. + The **Code tab inside Claude Desktop** uses the Claude Code plugin path; + it is not the same setup as regular Desktop chat. If the mode is unclear, + ask which tab the user wants to use. Other local clients use Local MCP + Server. Clients using explicit account/key headers use Remote MCP + (header-capable clients). The ChatGPT desktop app uses the Codex section. 2. Decide whether you have local shell and filesystem access. 3. If you can edit local user config files, do the setup yourself after showing the config you will add. @@ -65,7 +64,38 @@ user explicitly asks for developer integration. still writing `MEMORY.md`. If Claude's built-in memory is switched on and competes, tell the user they can turn it off in account settings. -## Requirements +## Claude OAuth Connector (Claude Desktop chat and Claude web) + +Use this path for regular Claude chat when custom connectors are available. +Production supports OAuth at `https://relayer.memory.walrus.xyz/api/mcp`. +No local Node.js, `claude` CLI, config-file edits, or manually copied delegate +key are needed. Skip the local requirements and terminal login sections. + +1. Ask the user to open Claude's **Settings → Connectors → Add custom connector** + and enter `https://relayer.memory.walrus.xyz/api/mcp`. +2. Follow the browser consent flow with a Sui wallet that owns the Walrus Memory + account. If needed, the flow guides account creation. Explain the scopes + shown on the consent screen: `memwal:read` reads memories, `memwal:write` + saves them, and `offline_access` allows token refresh. The relayer holds an + encrypted delegate private key for this flow. Let the user review and + approve the wallet authorization themselves. +3. Return to Claude, enable the connector for the conversation if needed, and + verify memory tools by saving and recalling a neutral setup-check fact. + If authorization expires, reconnect through Claude's connector settings; + do not prescribe a local `npx` login or require local-only login tools. + +This provides memory tools in chat, not the Code plugin's lifecycle hooks. +Do not claim it installs the plugin in the Code tab. If the user wants local +credentials instead, offer the explicit Claude Desktop local MCP option below. +If custom connectors are unavailable, report that limitation before offering +another route; do not silently switch to local installation. + +Full consent and disconnect details: +https://docs.wal.app/walrus-memory/mcp/claude-connector +Removing the connector alone does not remove the onchain delegate; the user +can revoke that delegate from the Walrus Memory dashboard. + +## Requirements (local MCP and header authentication only) For local MCP clients, check Node.js: @@ -87,7 +117,7 @@ wallet. The sign-in flow writes persistent credentials to: Those credentials are what survive MCP server restarts. -## Recommended Login +## Recommended Login (local MCP and header authentication only) When you can run a real terminal command, prefer this first. Run it from the user's home directory when possible, especially if the current project is a @@ -166,10 +196,11 @@ and merge: } ``` -### Claude Desktop +### Claude Desktop local MCP (optional) -Use the config file below. Claude Desktop's Settings > Connectors UI only -supports OAuth remote servers and does not work for Walrus Memory. +For regular chat, prefer the Claude OAuth Connector section above. Use this +local alternative only when the user chooses it; it requires Node.js 20+. +This config does not install the plugin for the Code tab. Edit: @@ -239,11 +270,27 @@ bundles the MCP server plus `SessionStart` / `UserPromptSubmit` / `PostToolUse` hooks. Those hooks make the agent prefer `memwal_*` over Claude Code's built-in `MEMORY.md`. MCP-only leaves auto-remember and auto-recall unreliable. -If you have a terminal, run the CLI yourself. Do not ask the user to type -`/plugin` commands unless you cannot run `claude plugin`. +This includes the Code tab inside Claude Desktop. Claude's current Desktop +plugin UI can install plugins without a terminal: open the `+` menu beside the +prompt, choose **Plugins → Add plugin**, select MemWal from a configured +marketplace, then use **Manage plugins** to confirm it is enabled. If the MemWal +marketplace is not listed, add it with the terminal commands below. Do not paste +`/plugin` commands into chat as plain text; that can make the agent edit settings +without installing the plugin. + +Having the desktop app installed does not guarantee that the `claude` CLI is +available in a shell. For the terminal route, check `node --version`, +`npm --version`, and `claude --version` first. If a prerequisite is missing, +explain the exact blocker and point to https://nodejs.org/ or +https://code.claude.com/docs/en/setup. Do not create `enabledPlugins` entries +as a substitute for a verified install. If the user only needs regular chat, +offer the OAuth connector instead; it does not add Code hooks. + +Run the following as separate commands in a terminal, stopping on the first +failed command: ```bash -claude plugin marketplace add MystenLabs/MemWal +claude plugin marketplace add https://github.com/MystenLabs/MemWal.git claude plugin install memwal@memwal-plugins -s user claude plugin list claude mcp list @@ -253,14 +300,25 @@ claude mcp list enabled. `claude mcp list` should show a connected `memwal` server. If either command fails, stop and report the exact error. -The slash commands below are only for a human inside an interactive Claude -Code session: +The slash commands are only for an interactive Claude Code session that +actually exposes `/plugin`; they are not chat instructions. The Desktop UI path +above is preferred when available. If `/plugin` is unavailable, use the +terminal route instead: ```text -/plugin marketplace add MystenLabs/MemWal +/plugin marketplace add https://github.com/MystenLabs/MemWal.git /plugin install memwal@memwal-plugins ``` +A delegate key is not required to install the plugin. After installation, +fully quit and reopen Claude Code, then ask it to run `memwal_login` and follow +the browser sign-in. Verify the plugin is enabled, the MCP server is connected, +and a memory can be saved and recalled before reporting success. + +If adding the HTTPS source reports an existing marketplace with a different +source, inspect `claude plugin marketplace list` and explain the conflict; +do not remove an existing marketplace or configuration without review. + Preserve an existing user-scope `memwal` MCP server (for example a staging or pinned version). The plugin adds hooks plus its own server; do not remove the user server to "clean up". @@ -320,13 +378,22 @@ Add an MCP server with: - command: `npx` - args: `["-y", "@mysten-incubation/memwal-mcp"]` +## Permissions during local setup + +Explain shell-command approvals and the plugin/config files being changed. +These instructions do not require reading Desktop files or controlling Apple +Music. If an unrelated OS permission prompt appears, stop and identify the +exact command and prompt before advising the user to grant access. Do not +request broad filesystem or automation access as a generic installation fix. +Do not attribute a prompt to Claude or MemWal without evidence. + ## Restart Fully quit and reopen the AI client after editing MCP config. On macOS, use `Cmd+Q`; closing the window is not enough. The first launch may take 5-10 seconds while `npx` fetches the package. -## Verify Tools +## Verify Tools (local MCP) Fastest check: ask the client to call `memwal_health`. A healthy reply proves the server is reachable. If health succeeds but memory tools return 401, @@ -378,7 +445,8 @@ If the tools support a `namespace` argument, use a setup-only namespace such as Some clients connect to a remote MCP server URL instead of running a local command. Use this path only when the client can attach BOTH custom request headers shown below. A client that exposes only a single bearer/token field -cannot use this path. +cannot use this header-authentication path. Claude Desktop/web should use the +OAuth connector above, which does not require custom headers. ChatGPT web connectors are not supported. The chatgpt.com connector UI exposes only one bearer slot, so the required `x-memwal-account-id` header cannot be @@ -418,8 +486,8 @@ cat ~/.memwal/credentials.json settings, never into the chat. You (the agent) must not read or print this file. Treat the bearer token like an API key and never save it in a repo file. -If the client cannot attach both custom headers, remote MCP is not available -for it today. Tell the user which clients are supported (Claude Desktop, Claude +If the client cannot attach both custom headers, this header-authentication +path is unavailable. For Claude Desktop/web, use OAuth instead. Tell the user which clients are supported (Claude Desktop, Claude Code, Cursor, Codex) and offer to set one of those up instead. Do not re-run this skill for the same client. @@ -434,7 +502,7 @@ this skill for the same client. | `codex plugin marketplace add` prints `already added from a different source` | A stale clone is left in `~/.codex/.tmp/marketplaces/memwal-plugins/`. Note the command still exits 0, so a scripted install does not notice. Remove that folder and the empty `~/.codex/plugins/cache/memwal-plugins/`, then add the marketplace again. | | Cursor shows no Walrus Memory plugin | Confirm the folder `~/.cursor/plugins/local/memwal` exists and holds `.cursor-plugin/plugin.json`, then fully quit and reopen Cursor. There is no `cursor plugin` CLI, so verify from the logs instead: the newest folder under `~/Library/Application Support/Cursor/logs/` on macOS gets an `mcp-server-plugin-memwal-memwal.log` once the plugin loads. A `mcp-server-user-memwal.log` next to it means a duplicate manual entry is still in `~/.cursor/mcp.json`. | | Memory tools fail with `MCP rate limit: ip_active_cap` (HTTP 429), sometimes followed by 503 | Too many concurrent memwal-mcp sessions from one machine. The usual causes are a duplicate server, where a plugin install sits next to a leftover manual `memwal` entry in the client's MCP config, and stale `memwal-mcp` processes left behind by earlier sessions. Remove the duplicate entry first. Then list the leftovers with `pgrep -fl memwal-mcp` and ask the user which clients they still want running before ending any process. Restart the client afterwards. | -| `claude plugin` commands are not recognized | This Claude Code build has no plugin CLI. Update Claude Code, or use the MCP-only fallback and tell the user hooks are missing. | +| `claude` is missing or `claude plugin` is not recognized | Check CLI availability and version; the Desktop app alone does not guarantee a shell CLI. Install/update Claude Code for the plugin route, or offer OAuth for regular chat. Do not claim a config edit completed installation. | | Only `memwal_login` works | Credentials are missing. Run `memwal_login` or `npx -y @mysten-incubation/memwal-mcp login --prod`. | | Memory tools return 401 | The delegate key may be stale or revoked. Run `npx -y @mysten-incubation/memwal-mcp login --prod` again. | | `memwal_recall` returns nothing although memories were saved before | Run `memwal_restore ` to rebuild the search index from Walrus, then retry the recall. | @@ -442,6 +510,12 @@ this skill for the same client. ## Final Report +For OAuth, report connector authorization and save/recall verification status. +Do not require a local restart, credentials file, plugin, or Node.js. If the +user has not completed consent or tools have not been verified, say what is +still pending. Automatic save/recall in regular chat is best-effort, not the +Code plugin's lifecycle-hook behavior. + Keep the final response short and make the restart the first thing the user sees whenever a restart is required. diff --git a/apps/app/src/pages/ConnectMcp.test.tsx b/apps/app/src/pages/ConnectMcp.test.tsx index ad569a68f..294631260 100644 --- a/apps/app/src/pages/ConnectMcp.test.tsx +++ b/apps/app/src/pages/ConnectMcp.test.tsx @@ -116,10 +116,16 @@ describe('MCP sign-in hand-off', () => { expect(screen.getByText(/Nothing answered on your computer/i)).toBeInTheDocument() expect(screen.getByText(/stopped waiting after this tab was already open/i)).toBeInTheDocument() expect(screen.queryByText(/opened too late/i)).not.toBeInTheDocument() - expect(screen.getByText(/Sign in again and open the new link straight away/i)).toBeInTheDocument() - expect(screen.getByText(/left running through the wallet prompt/i)).toBeInTheDocument() - expect(screen.getByText(/unused key from this attempt is already on your account/i)).toBeInTheDocument() expect(screen.queryByText(/usually works/i)).not.toBeInTheDocument() + // The key is registered and saved by the client that started this + // sign-in, so a restart reclaims it. Signing in again cannot register + // it twice, and removing it from the dashboard throws it away. + expect(screen.getByText(/Restart your MCP client within 24 hours/i)).toBeInTheDocument() + expect(screen.getByText(/only if you mean to abandon it/i)).toBeInTheDocument() + expect( + screen.queryByText(/Remove it from the dashboard if you are not using it/i), + ).not.toBeInTheDocument() + expect(screen.queryByText(/Sign in again and open the new link straight away/i)).not.toBeInTheDocument() }) it('explains an expired link when preflight cannot reach the listener', async () => { diff --git a/apps/app/src/pages/ConnectMcp.tsx b/apps/app/src/pages/ConnectMcp.tsx index 99ddf63e4..9e6edc4b4 100644 --- a/apps/app/src/pages/ConnectMcp.tsx +++ b/apps/app/src/pages/ConnectMcp.tsx @@ -620,10 +620,11 @@ function SuccessCard({ : 'The app on your computer rejected this hand-off. That usually means this tab is leftover from a sign-in that already finished, or the request did not match what the app expected.'}

- Sign in again and open the new link straight away. A - retry only helps once the MCP client is left running through the wallet - prompt. The unused key from this attempt is already on your account. - Remove it from the dashboard if you are not using it.{' '} + Restart your MCP client within 24 hours. On start it + finds the key from this attempt and signs you in, with no second wallet + prompt. Do not sign in again first: this key is already registered to + your account, and it cannot be registered twice. Remove it from the + dashboard only if you mean to abandon it.{' '} {config.docsUrl && ( { + capturedWhere = clause; + return Promise.resolve([]); +}); +const from = vi.fn(() => ({ where })); +const select = vi.fn(() => ({ from })); +const values = vi.fn(() => Promise.resolve([])); +const insert = vi.fn(() => ({ values })); + +vi.mock("server-only", () => ({})); +vi.mock("postgres", () => ({ default: () => ({}) })); +vi.mock("drizzle-orm/postgres-js", () => ({ + drizzle: () => ({ select, insert }), +})); + +beforeEach(() => { + vi.clearAllMocks(); + capturedWhere = undefined; +}); + +describe("voteMessage DB-layer scoping", () => { + it("scopes the existing-vote lookup to both chatId and messageId", async () => { + const { voteMessage } = await import("./queries"); + + await voteMessage({ + chatId: CHAT_VAL, + messageId: MESSAGE_VAL, + type: "up", + }); + + expect(where).toHaveBeenCalledTimes(1); + expect(capturedWhere).toBeDefined(); + + const { sql, params } = new PgDialect().sqlToQuery(capturedWhere as never); + + expect(sql).toMatch(/"chatId"/); + expect(sql).toMatch(/"messageId"/); + expect(params).toContain(CHAT_VAL); + expect(params).toContain(MESSAGE_VAL); + }); +}); diff --git a/apps/status/server.mjs b/apps/status/server.mjs index 47ec71793..f8111a76f 100644 --- a/apps/status/server.mjs +++ b/apps/status/server.mjs @@ -354,8 +354,13 @@ async function probeRelayer(name, rawBase, target) { } } - const writesPaused = isRecord(health) && health.writes === 'paused' - const reportedOk = isRecord(health) && health.status === 'ok' && !writesPaused + const writes = isRecord(health) ? health.writes : undefined + // /health stays HTTP 200 + status:"ok" through a Walrus outage so CI's + // wait-for-relayer gate does not hang. The write-path signal is `writes`. + // Allowlist: a fourth writes value must not silently become operational + // the way degraded did before this PR. Missing writes (old relayer) is ok. + const reportedOk = + isRecord(health) && health.status === 'ok' && (writes === 'ok' || writes === undefined) const status = response.ok ? (reportedOk ? 'operational' : 'degraded') : 'outage' return { diff --git a/apps/status/src/App.tsx b/apps/status/src/App.tsx index 959680876..47be15fc4 100644 --- a/apps/status/src/App.tsx +++ b/apps/status/src/App.tsx @@ -102,6 +102,7 @@ interface ComponentRow { uptimeLabel: string history: HistoryBucket[] writesPaused?: boolean + writesDegraded?: boolean } interface IncidentDay { @@ -166,9 +167,10 @@ function getOverallStatus(snapshot: StatusSnapshot | null, loadState: LoadState) return 'monitoring' } -function getStatusTitle(status: StatusKind, writesPaused = false) { +function getStatusTitle(status: StatusKind, writesPaused = false, writesDegraded = false) { if (status === 'operational') return 'All Systems Operational' if (status === 'degraded' && writesPaused) return 'Writes Paused' + if (status === 'degraded' && writesDegraded) return 'Writes Degraded' if (status === 'degraded') return 'Degraded Performance' if (status === 'outage') return 'Service Disruption' return 'Checking System Status' @@ -261,6 +263,7 @@ function buildRows(snapshot: StatusSnapshot | null, loadState: LoadState): Compo uptimeLabel: formatUptime(history), history: normalizeBuckets(history, component.status), writesPaused: component.health?.writes === 'paused', + writesDegraded: component.health?.writes === 'degraded', } }) } @@ -401,8 +404,20 @@ function calendarRangeLabel(months: CalendarMonth[]) { return `${months[0].label} to ${months[months.length - 1].label}` } -function StatusPill({ status, writesPaused }: { status: StatusKind; writesPaused?: boolean }) { - const label = writesPaused ? 'Writes Paused' : statusLabel[status] +function StatusPill({ + status, + writesPaused, + writesDegraded, +}: { + status: StatusKind + writesPaused?: boolean + writesDegraded?: boolean +}) { + const label = writesPaused + ? 'Writes Paused' + : writesDegraded + ? 'Writes Degraded' + : statusLabel[status] return {label} } @@ -469,7 +484,11 @@ function ComponentStatusRow({ row }: { row: ComponentRow }) {

{row.name}

- +
@@ -1054,6 +1073,7 @@ export default function App() { const overallStatus = getOverallStatus(snapshot, loadState) const writesPaused = (snapshot?.components ?? []).some((c) => c.health?.writes === 'paused') + const writesDegraded = (snapshot?.components ?? []).some((c) => c.health?.writes === 'degraded') const rows = useMemo(() => buildRows(snapshot, loadState), [snapshot, loadState]) const uptimeRows = useMemo(() => rows.filter((row) => row.status !== 'monitoring'), [rows]) const productionHistory = snapshot?.histories?.['relayer-production'] @@ -1073,7 +1093,7 @@ export default function App() { {route === 'current' && ( <>
-

{getStatusTitle(overallStatus, writesPaused)}

+

{getStatusTitle(overallStatus, writesPaused, writesDegraded)}

{(error || componentError || snapshot?.database?.error) && ( diff --git a/docs/mcp/changelog.mdx b/docs/mcp/changelog.mdx index cd8cc258a..26132c73b 100644 --- a/docs/mcp/changelog.mdx +++ b/docs/mcp/changelog.mdx @@ -28,9 +28,35 @@ questions: - What changed in the MemWal MCP changelog? - When was the automatic memory plugin added to MemWal MCP? answer: >- - The latest MCP package release is 0.0.13. A `memwal_remember`, `memwal_remember_bulk` or `memwal_analyze` whose reply never arrives is no longer reported as safe to retry: the relayer accepts those with HTTP 202 and finishes them in a durable queue, so the write may already have landed, and `/api/remember/bulk` has no idempotency key — repeating it stores a second paid copy. The tool now says the write may have completed and points at `memwal_recall` to check before re-saving, while a lost read still says plainly that retrying is safe. When the relayer rejects the saved delegate key, tool calls now get an auth error pointing at `memwal_login` instead of waiting minutes for a retry hint that cannot work. It writes the credentials file by creating a new 0600 file and renaming it into place, so a sign-in never puts the delegate private key into a credentials.json that a manual chmod or a restored backup left world-readable. A completed sign-in is now confirmed with a notification and a one-shot banner naming the account and the resolved credentials path, and the bridge keeps reading stdin after an in-session login instead of going deaf. Unrecognised command-line options now warn instead of being silently ignored, `--help` lists the network presets and the URLs each resolves to, `memwal_health` names the relayer the client dialled, and `memwal_restore` reports `failed` and retries the same page when truncation is a download or embed blip, instead of always telling the agent to raise `limit`. Plugin launch configs (`.mcp.json`, Cursor/Codex copies, and the Codex fallback installer) now pin `@mysten-incubation/memwal-mcp@0.0.13` so npx cannot keep a cached 0.0.5. Version 0.0.12 forwards the MCP client's initialize.clientInfo to the relayer so sidecar logs can name the coding agent (Claude Code, Codex, Cursor, and others) on each session and tool call, and it resolves the credential directory on every access so MEMWAL_CREDS_DIR can override it. + The latest published MCP package on npm `latest` is 0.0.13. This branch's in-tree plugin launcher installs `@mysten-incubation/memwal-mcp@0.0.14-dev.0` because npm has no `0.0.14` yet — pinning the unpublished version would fail to resolve. 0.0.14 (unreleased) stops cold start from advertising tools an older relayer does not serve, bounds stalled relayer calls, honours `retry_after` on 429, and returns `memwal_remember` at accept by default (`MEMWAL_MCP_REMEMBER_WAIT_MS=0`; set `90000` to wait for `blob_id`). 0.0.13 stopped reporting a timed-out write as safe to retry, answers a rejected delegate key with `memwal_login` instead of a dropped-connection story, backs off on handshake 429, and pinned plugin launchers to `@0.0.13` so npx cannot keep a cached 0.0.5. --- +## 0.0.14 + +Unreleased package version. The plugin's launcher installs `@mysten-incubation/memwal-mcp@0.0.14-dev.0` (the published `dev` dist-tag), named by `mcpPackageVersion` in `plugin/plugin.json`, until `0.0.14` exists on npm. Drop that key in the same release that publishes it, so the launcher goes back to the plugin's own version, or plugin installs stay on the dev build. + +### Added + +- Opt-in Streamable HTTP transport for the stdio bridge. Set `MEMWAL_MCP_TRANSPORT=http` (aliases `streamable`, `streamable-http`) to dial the relayer's single `/api/mcp` endpoint, where a call is answered on the same request instead of being split across a POST and an SSE stream, so there is no idle watchdog. The default stays `sse` and unrecognised values fall back to it. Reconnect replay is not transport-aware yet: on Streamable a disconnect mid-send can look sent, so a replayed write may duplicate. + +### Fixed + +- The plugin no longer starts the MCP server through `npx @mysten-incubation/memwal-mcp@`. npx resolves a package name against the directory the MCP client was started in — the project the user has open — so a project carrying an installed package of that name, claiming the pinned version, was run instead of the published one; a pinned `npx …@0.0.14` command was reproduced running a project's own binary while offline. The version pin was no defence, because the planted package simply claims that version. Every launch site (`.mcp.json`, the Cursor and Codex copies, and the Codex fallback installer) now runs `plugin/scripts/launch_mcp.mjs`, which installs the pinned version once under `~/.memwal/runtime/memwal-mcp@` and launches that absolute entry point with the current node binary. Nothing on that path consults the project's `node_modules`, a PATH-relative bin shim, or the project's `.npmrc`, and the launcher fails instead of falling back to the name when the trusted install cannot be produced. Exploiting the old behaviour required write access to installed package and bin files inside the project, so a plain repository clone or a `package.json` alone was never enough. `MEMWAL_MCP_RUNTIME_DIR` relocates the trusted directory and is rejected unless it is an absolute path. (WALM-640) +- An unfinished write is no longer reported with a `blob_id`. `persist_uploaded_state` stamps a real `blob_id` and status `uploaded` the moment the blob is on Walrus, but the job is not done — SetMetadataAndTransfer has not handed the blob object to the user and `insert_vector` has not indexed it, so the fact is not recallable and a later failure can still lose it. `uploaded` was already counted as in flight; the render was not, and printed the `blob_id` of any row that had one. A live dev batch read `0/6 saved, 6 still uploading` with three of those rows carrying a `blob_id=`, contradicting the same reply's rule that only a `blob_id` means the fact is stored — an agent that believes it tells the user a fact is saved while it is still in flight. Only a `done` row shows a `blob_id` now, which is what the single-job path always did. Same guard in `memwal_remember_bulk` and `memwal_analyze`, whose timeout rows render from the same template. (#921 follow-up) +- The cold-start tool list no longer advertises a tool the relayer may not serve. The bridge ships on npm and updates itself while a relayer ships per environment, so 0.0.14-dev.0 dialled prod and staging still on 0.0.13: cold start named `memwal_remember_status`, which neither registers, and the pending-write wording sent the agent to go call it — one live run spent 90.67s there before erroring. Cold start is now a floor rather than a forecast (`BASELINE_RELAYER_TOOLS`): it carries only what the oldest supported relayer serves and its descriptions name nothing outside it, while newer tools still reach the client a beat later on the relayer's own `tools/list`. Initialize `instructions` use the same floor (they used to name `memwal_remember_status` before any `tools/list`). And a call for a tool outside that floor — including during the cold-start window, before any upstream `tools/list` has been seen — is now answered locally and at once — naming the tools that do exist and saying plainly that nothing ran — instead of being forwarded into a wait that only ends at the orphan deadline. (#928) +- A tool call that times out or cannot reach the relayer now says why and what to do, in three lines: `Cause`, `Relayer health` (checked against `/health` on the spot) and `Next step`. It used to read `Tool error: This operation was aborted`. A recall the relayer cut short names the step it was stuck in: credential check, embedding, vector search, Walrus download or SEAL decrypt. Writes are never told a retry is safe. (WALM-396) +- A call whose reply never arrives is answered with the relayer's health. The bridge checks `/health` first and says whether the relayer is up with this one call stuck, unhealthy, down, or not resolvable at the configured URL, instead of a bare "did not answer". A call being answered this way is never re-sent, so a reconnect landing while the check runs cannot run a write a second time. A lost `memwal_recall` reply is now answered after 2 minutes instead of 4. (WALM-396) +- Bound every relayer call the tools make. The pinned SDK aborts a request only when the caller passes a signal, which none of the memory methods do, so a stalled socket kept a tool running with no ceiling — `memwal_remember` was observed still going past 120s against a 90s budget. Accepts are bounded at 15s (`MEMWAL_MCP_ACCEPT_DEADLINE_MS`), waits at their own budget plus grace. The request is not cancelled — the SDK exposes no way to pass a signal — but the agent is no longer held by it. +- Honour the relayer's `retry_after` instead of dropping the write. Once the per-delegate-key budget (60 weighted requests/minute) is spent the relayer answers 429 with a cooldown, and nothing backed off: the fact was never written and the agent saw only an opaque tool error. A short cooldown is now absorbed; a long one is reported with the wait named, stating plainly that the fact was NOT saved and pointing at the cheaper shape — one `memwal_remember_bulk` rather than N single calls, one `memwal_remember_status(job_ids)` rather than N status calls. Only rejections that provably never reached the handler retry, so `/api/remember/bulk`, which carries no idempotency key, cannot be duplicated by a retry. +- `memwal_remember` sends a content-derived idempotency key, so the retry its own timeout message invites really does attach to the job already in flight instead of storing a second paid copy. The key is computed by the tool rather than relied on from the SDK, whose published build mints a random UUID per client instance. +- `memwal_remember_status` accepts `job_ids` to settle a whole batch in one call, and reports a mixed batch honestly — a still-uploading row no longer renders the poll timeout as `error=`, which read as a failed write. Settling in one request also matters against the rate limit: 20 ids cost one request, not twenty. +- The bridge's cold-start tool list no longer disagrees with the sidecar's. `memwal_remember_status` advertised only `job_id`, required, under `additionalProperties: false`, so the batch call the tools themselves instruct was rejected until `tools/list_changed` arrived; the `waitMs` ceiling advertised 60000 after the sidecar lowered it to 45000, which came back as an MCP validation error; and `memwal_remember_bulk` still carried its pre-queue description. Tests now pin the parts an agent acts on. +- Persist the delegate keypair before sign-in hands the URL to the browser, and reclaim it on the next start. The browser's onchain `add_delegate_key` costs gas and is irreversible, and it happens before the callback that saves the private half, so a client that died in that window destroyed the only copy of a key the user had already paid for and left an orphaned registration nobody could use. (#793) Signing out discards the pending record along with the credentials, a second sign-in against the same relayer reuses the stranded key rather than minting over it, and a sign-in that cannot write the record fails instead of publishing a URL it cannot back. Reclaiming works on Mainnet; Testnet requires an account-id hint the recovering client does not have. + +### Changed + +- `memwal_remember` / `memwal_remember_bulk` return at accept by default (`MEMWAL_MCP_REMEMBER_WAIT_MS=0`, ~1s, `job_id`). The Walrus write continues in the background; do not treat that reply as stored. Set `MEMWAL_MCP_REMEMBER_WAIT_MS=90000` to restore wait-for-`blob_id` (90s ceiling). Do not use a value between 0 and the real completion time — that pays the wait and still returns pending. The MCP TypeScript SDK's default tools/call timeout is 60s, so a 90s wait loses on hosts that do not raise it. + ## 0.0.13 This release stops a write whose reply was lost from being reported as safe to retry — repeating one can store a second paid copy — answers tool calls with an auth error pointing at `memwal_login` when the relayer rejects the saved delegate key, writes the credentials file through a fresh `0600` file that it renames into place, confirms a completed sign-in and keeps the bridge reading stdin afterwards, warns on unrecognised command-line options instead of ignoring them, documents the network presets in `--help`, names the relayer in `memwal_health`, reports restore `failed` counts when truncation is a transient download or embed blip, and pins plugin launch configs (`.mcp.json`, Cursor/Codex copies, and the Codex fallback installer) so npx cannot keep a cached 0.0.5. diff --git a/docs/mcp/claude-code.md b/docs/mcp/claude-code.md index 160969ebe..3339e8ecf 100644 --- a/docs/mcp/claude-code.md +++ b/docs/mcp/claude-code.md @@ -30,13 +30,18 @@ questions: - What is the difference between the MemWal plugin and MCP-only on Claude Code? - How do I fix the memwal MCP server when it shows as failed in Claude Code? answer: >- - To add Walrus Memory to Claude Code, install the MemWal plugin through the marketplace (/plugin marketplace add MystenLabs/MemWal, then /plugin install memwal@memwal-plugins), or add it as MCP-only with claude mcp add. The plugin includes lifecycle hooks for session start, user prompt, and post-tool events that reinforce automatic memory behavior and make the agent prefer Walrus Memory over Claude Code's built-in memory. + To add Walrus Memory to Claude Code, install the MemWal plugin through the marketplace (/plugin marketplace add https://github.com/MystenLabs/MemWal.git, then /plugin install memwal@memwal-plugins), or add it as MCP-only with claude mcp add. The plugin includes lifecycle hooks for session start, user prompt, and post-tool events that reinforce automatic memory behavior and make the agent prefer Walrus Memory over Claude Code's built-in memory. --- Add MemWal to Claude Code so it recalls context and saves durable facts as you work. Install it as a **plugin** (recommended; adds automatic-memory hooks) or as **MCP-only** (just the tools). +This guide covers Claude Code, including the **Code tab inside Claude Desktop**. +For regular Desktop or web chat, use the [OAuth custom connector](/mcp/claude-connector), +which does not require Node.js or the Claude Code CLI. + ## Prerequisites +- Confirm `claude --version` works in the terminal you will use. Installing the Desktop app does not guarantee the CLI is on your `PATH`; see [Claude Code setup](https://code.claude.com/docs/en/setup). - Install Node.js 20+ with `npx` on your `PATH`; check with `node --version`. - Use a Claude Code version with plugin support if you want the plugin install; the `/plugin` command confirms support, and MCP-only works on any version with `claude mcp add`. - Have a [Walrus Memory account](/fundamentals/concepts/ownership-and-access) ready. An unauthenticated memory-tool call returns sign-in instructions rather than signing you in, so ask the agent to run `memwal_login` and open the URL it returns. You can create the account during that flow at [memory.walrus.xyz](https://memory.walrus.xyz). Config files carry no keys: credentials land in `~/.memwal/credentials.json` after sign-in. @@ -47,16 +52,24 @@ Add MemWal to Claude Code so it recalls context and saves durable facts as you w - Agents with a terminal should run the CLI, then verify the plugin is enabled. Slash commands are only for a human inside Claude Code. + Run these as separate commands in a terminal, stopping if one fails. Agents with a shell tool can run them after checking the prerequisites, then verify the plugin is enabled. A manually added `enabledPlugins` setting alone does not verify installation. ```bash - claude plugin marketplace add MystenLabs/MemWal + claude plugin marketplace add https://github.com/MystenLabs/MemWal.git claude plugin install memwal@memwal-plugins -s user claude plugin list ``` + In Claude Desktop's Code tab, use the plugin UI instead: open the `+` + menu beside the prompt, choose **Plugins → Add plugin**, select MemWal + from a configured marketplace, then confirm it under **Manage plugins**. + Do not paste slash commands into chat as plain text. + + Slash commands are an alternative only in an interactive Claude Code + session that exposes `/plugin`: + ``` - /plugin marketplace add MystenLabs/MemWal + /plugin marketplace add https://github.com/MystenLabs/MemWal.git /plugin install memwal@memwal-plugins ``` @@ -97,7 +110,9 @@ Add MemWal to Claude Code so it recalls context and saves durable facts as you w | MemWal MCP (memory tools) | ✓ | ✓ | | Lifecycle hooks (automatic recall/save) | ✓ | ✗ | -MCP-only still saves and recalls on its own because the tools are proactive. The plugin adds hooks that reinforce the behavior and make the agent **prefer Walrus Memory over Claude Code's built-in memory**. The plugin pins the MCP server version, so `npx` cannot keep a cached older package. +MCP-only still saves and recalls on its own because the tools are proactive. The plugin adds hooks that reinforce the behavior and make the agent **prefer Walrus Memory over Claude Code's built-in memory**. + +The plugin also starts the server differently. Instead of resolving the package name with `npx`, it installs the pinned version once into `~/.memwal/runtime/memwal-mcp@` and launches that absolute path. `npx` resolves a name against the directory the client was started in, which is your project, so a package installed there under the same name and claiming the pinned version would have been run instead — the version pin does not prevent that. The plugin's launcher never looks at your project's `node_modules`, and it fails rather than falling back if the pinned version cannot be installed. ## Available tools diff --git a/docs/mcp/claude-desktop.md b/docs/mcp/claude-desktop.md index 6c86f3943..4af6b4ad6 100644 --- a/docs/mcp/claude-desktop.md +++ b/docs/mcp/claude-desktop.md @@ -28,17 +28,35 @@ questions: - How do I configure the MemWal MCP server for Claude Desktop? - Does Claude Desktop support the MemWal automatic memory plugin? answer: >- - To add Walrus Memory to Claude Desktop, configure the MemWal MCP server in your claude_desktop_config.json file using npx -y @mysten-incubation/memwal-mcp as the command. Claude Desktop supports MCP-only (not the plugin with lifecycle hooks). The tool descriptions still make the agent save and recall proactively. Restart Claude Desktop fully (Cmd+Q) after installation, then ask the agent to run memwal_login on first use. + For regular Claude Desktop chat, add the Walrus Memory OAuth custom connector + using https://relayer.memory.walrus.xyz/api/mcp and complete browser consent; + no local Node.js or CLI is required. Alternatively, run the local MCP server + with Node.js 20+ and sign in with memwal_login after restarting Desktop. + The Code tab uses the separate Claude Code plugin setup. --- Add MemWal to Claude Desktop so the agent can save and recall durable facts. Claude Desktop uses the **MCP server** (the memory tools); the automatic-memory plugin hooks are available on [Claude Code](/mcp/claude-code), [Codex](/mcp/codex), [Cursor](/mcp/cursor), and [Antigravity](/mcp/antigravity). -## Prerequisites +## Choose your setup + +For **regular Claude Desktop chat**, use the [OAuth custom connector](/mcp/claude-connector): +open **Settings → Connectors → Add custom connector**, enter +`https://relayer.memory.walrus.xyz/api/mcp`, and complete the wallet consent flow. +This requires no local Node.js, CLI, or manually copied delegate key. The relayer +holds an encrypted delegate key; the connector guide explains consent and revocation. + +The **Code tab inside Claude Desktop** uses the [Claude Code plugin setup](/mcp/claude-code) +instead. The chat connector does not install Code lifecycle hooks. + +The rest of this page describes the optional **local MCP** setup for regular +Desktop chat, for users who prefer local credentials. + +## Local MCP prerequisites - Node.js 20+ - A Walrus Memory account. The first memory tool call opens a browser sign-in (`memwal_login`). -## Installation +## Local MCP installation Add the server to your Claude Desktop config: @@ -65,7 +83,7 @@ Quit and reopen Claude Desktop (`Cmd+Q` on macOS; closing the window is not enou ## Add memory instructions -Claude Desktop cannot run the lifecycle hooks that reinforce automatic memory on +Regular Claude Desktop chat does not run the lifecycle hooks that reinforce automatic memory on [Claude Code](/mcp/claude-code), [Codex](/mcp/codex), and [Antigravity](/mcp/antigravity). The tools still work here and the agent might save and recall on its own, but that is best-effort, because Claude Desktop's built-in memory can win instead. State the @@ -82,7 +100,7 @@ Use Walrus Memory as my memory. - Prefer Walrus Memory over your built-in memory. ``` -## Available tools +## Local MCP tools | Tool | Description | |------|-------------| diff --git a/docs/mcp/overview.md b/docs/mcp/overview.md index bf354de34..2e0e734d1 100644 --- a/docs/mcp/overview.md +++ b/docs/mcp/overview.md @@ -43,9 +43,20 @@ There are two ways to use MemWal. The difference is whether you also get the **l | MemWal MCP: memory tools (`memwal_remember`, `memwal_recall`, …) | ✓ | ✓ | | Lifecycle hooks: automatic recall/save reminders | ✓ | ✗ | -- **Plugin** bundles the MCP server **and** lifecycle hooks. The `SessionStart` hook tells the agent to prefer the `memwal_*` tools over any built-in or local memory feature, and when to save without being asked. Automatic memory works with no further instructions from you. Available on **Claude Code**, **Codex**, **Antigravity**, and **Cursor**. +- **Plugin** bundles the MCP server **and** lifecycle hooks. The `SessionStart` hook tells the agent to prefer the `memwal_*` tools over any built-in or local memory feature, and when to save without being asked. Available on **Claude Code**, **Codex**, **Antigravity**, and **Cursor**. - **MCP-only** gives the agent the memory tools on **every** MCP client. The tool descriptions encourage proactive use, so agents often do save and recall on their own. Treat that as best-effort: it varies by client and model, and on a client that ships its own memory feature the built-in one commonly wins. +Saving **without being asked is on once you agree to it**: `memwal-mcp login` +asks the question in your terminal the first time, and on a fresh install +nothing is saved unprompted until you answer. Saved memories are permanent — +Walrus is immutable storage — which is why the question comes first. Change the +answer any time with `memwal-mcp auto-save on|off` (or `MEMWAL_AUTO_SAVE` in +your client's `env` block). Recall, and anything you explicitly ask to be +remembered, work either way. Credentials — passwords, API keys, tokens, private keys, seed +phrases, authorization headers, and URLs with an embedded `user:password` — are +excluded in both modes and stripped before a memory is written, because Walrus +storage is append-only and a stored secret cannot be deleted. + Prefer the plugin wherever you can install it. It is the tested path for reliable automatic save and recall, and it needs no extra instructions from you. On MCP-only clients, paste the client's instruction block (see [Claude Desktop](/mcp/claude-desktop#add-memory-instructions)) to get closer to the same behavior. @@ -69,11 +80,11 @@ What the user actually does differs per client. Pick your row: | Client | Automatic memory (hooks) | What you do | |---|:-:|---| -| [Claude Code](/mcp/claude-code) | ✓ Plugin | `/plugin marketplace add MystenLabs/MemWal`, then `/plugin install memwal@memwal-plugins` | +| [Claude Code](/mcp/claude-code) | ✓ Plugin | Code tab: use the [Claude Code plugin setup](/mcp/claude-code) with the explicit HTTPS marketplace URL; regular Claude Desktop/web chat: use the [OAuth connector](/mcp/claude-connector) | | [Codex](/mcp/codex) | ✓ Plugin | `codex plugin marketplace add MystenLabs/MemWal`, then `codex plugin add memwal@memwal-plugins`, then trust the hooks through `/hooks` | | [Antigravity](/mcp/antigravity) | ✓ Plugin | `npx degit MystenLabs/MemWal/packages/mcp/plugin ~/.gemini/config/plugins/memwal` | | [Cursor](/mcp/cursor) | ✓ Plugin | `npx -y degit MystenLabs/MemWal/packages/mcp/plugin ~/.cursor/plugins/local/memwal` | -| [Claude Desktop](/mcp/claude-desktop) | ✗ MCP-only | Edit `claude_desktop_config.json`, then [add memory instructions](/mcp/claude-desktop#add-memory-instructions) | +| [Claude Desktop](/mcp/claude-desktop) | ✗ MCP-only | Regular chat: use the [OAuth connector](/mcp/claude-connector); local MCP is optional. The Code tab uses [Claude Code](/mcp/claude-code). | | [OpenCode](/mcp/opencode) | ✗ MCP-only | Edit the OpenCode MCP config | | ChatGPT desktop app | ✓ Plugin | Ships Codex, so follow [Codex](/mcp/codex) | | ChatGPT web (Connectors) | ✗ Not supported | n/a | diff --git a/docs/mcp/quickstart.md b/docs/mcp/quickstart.md index 963d6cffb..e9d915448 100644 --- a/docs/mcp/quickstart.md +++ b/docs/mcp/quickstart.md @@ -42,9 +42,14 @@ answer: >- Every supported client runs the same local server, `npx -y @mysten-incubation/memwal-mcp`, and differs only in where the configuration lives. Pick your client below, add the server, restart, and sign in. +For regular Claude Desktop or web chat, use the [OAuth custom connector](/mcp/claude-connector) +instead of the local setup below. It needs no local Node.js or CLI. The Code tab +inside Claude Desktop uses the [Claude Code plugin guide](/mcp/claude-code). + ## Prerequisites - You need Node.js 20 or later, because the server runs through `npx` with no install step. +- `npx` resolves the package name against the directory your client starts the server in, which is usually the project you have open. A project that contains an installed `@mysten-incubation/memwal-mcp` of its own would be run instead of the published package, and pinning a version in the command does not prevent it. The [plugin](/mcp/claude-code) avoids this by installing the pinned version into `~/.memwal/runtime` and launching that absolute path; to get the same property without the plugin, install the version yourself outside any project and point `command`/`args` at the absolute entry point (see the [package README](https://github.com/MystenLabs/MemWal/blob/dev/packages/mcp/README.md#how-the-plugin-launches-the-server)). - You need a [Walrus Memory account](/fundamentals/concepts/ownership-and-access). An unauthenticated memory-tool call returns sign-in instructions rather than signing you in, so ask the agent to run `memwal_login` and follow the URL it returns to connect your wallet. Config files carry no keys. ## Set up your client diff --git a/docs/mcp/reference.md b/docs/mcp/reference.md index d980edc43..a2897ca41 100644 --- a/docs/mcp/reference.md +++ b/docs/mcp/reference.md @@ -53,7 +53,9 @@ This is why many first-run sessions show `memwal_login` before the other tools a ### memwal_remember -Save a durable fact to the user's Walrus Memory. The agent calls this **proactively** when the user states a preference, decision, constraint, correction, identity detail, or recurring workflow, not only when they explicitly ask. Skip one-off tasks, the current file or bug, and small talk. Pass the full statement; do not summarize. +Save a durable fact to the user's Walrus Memory. The agent calls this **proactively** when the user states a preference, decision, constraint, correction, identity detail, or recurring workflow, not only when they explicitly ask — provided automatic memory is on. That is the user's standing answer to a question `memwal-mcp login` asks once in a terminal; it can be changed with `memwal-mcp auto-save on|off`. Skip one-off tasks, the current file or bug, and small talk. Pass the full statement; do not summarize. + +Credentials are never stored. Passwords, API keys, access and refresh tokens, private keys, seed phrases, authorization headers, session cookies, and URLs with an embedded `user:password` are stripped from the text before the write, on this tool, `memwal_remember_bulk` and `memwal_analyze` alike. A message that mixes a preference with a credential keeps the preference and loses only the credential; the reply names which kinds were removed. | **Parameter** | **Type** | **Required** | **Description** | | --- | --- | --- | --- | @@ -130,13 +132,28 @@ Both session tools (`memwal_login`, `memwal_logout`) are intercepted locally by Credentials resolve from two places, in order: -1. `.memwal/credentials.json` in the **working directory or a parent of it** +1. `.memwal/credentials.json` in the **working directory or a parent of it** — used only once you have [approved it](#approving-a-project-file) 2. `~/.memwal/credentials.json` (global, per machine) -The search starts in the working directory and walks up, the way `.npmrc` and `.git/config` resolve, so a command run from a subfolder still picks up that project's credentials. The first `.memwal/credentials.json` it finds wins. +The search starts in the working directory and walks up, the way `.npmrc` and `.git/config` resolve, so a command run from a subfolder still picks up that project's credentials. The first `.memwal/credentials.json` it finds wins, provided you approved it. The walk stops at your project root (the directory holding `.git`), at your home directory, or at the filesystem root, whichever comes first. That bound keeps one project from picking up a credentials file belonging to a parent folder that holds unrelated checkouts. If nothing is found inside it, the global file is used. Whichever file is chosen is the one read, written, and deleted for that run. +### Approving a project file + +A project's `.memwal/credentials.json` lives inside the repository, so anyone who can commit to it — or who can get you to open a clone — could otherwise pick the account and relayer every memory written from that directory goes to. The file is therefore **ignored until you approve it**, once per machine: + +```bash +cd ~/code/my-project +memwal-mcp approve-project +``` + +Until then the global credentials are used and one stderr line names the file that was skipped, the account and relayer it wanted, and this command. An approval covers one exact project path, account, delegate key and relayer: change any of them and it has to be approved again. The record is kept in `~/.memwal/project-approvals.json`, outside the repository, so a repository cannot carry its own approval. `memwal-mcp revoke-project` withdraws it. + +Approving decides which file is written as well as which is read. From then on a sign-in from that project stores a delegate private key, in plain text, in `.memwal/credentials.json` inside the repository — `approve-project` and the pre-sign-in warning both say so. Add `.memwal/` to `.gitignore`. The login write-ahead record is kept in `~/.memwal/login-pending/` and never enters the repository. + +`MEMWAL_CREDS_DIR` overrides project resolution entirely and needs no approval. Because it skips the gate it must be an **absolute path outside the current project**, and anything else is refused with an error rather than ignored: a relative value resolves against the working directory, which would put the credentials and the approval record inside the repository, and an MCP client hands this process the `env` block it reads from `.cursor/mcp.json`, `.vscode/mcp.json` or `.claude/settings.json` in the checkout — where editors expand `${workspaceFolder}`, so an in-project absolute path is not evidence you set it. An empty value means unset. + ### Working on several accounts Without a project-local file, every project on the machine shares one credential. Signing in from one project silently repoints the others at a different account and delegate key, and memories written in that state land on the wrong account, on immutable storage, with no delete path. @@ -148,6 +165,7 @@ cd ~/code/my-project mkdir -p .memwal memwal-mcp login # writes to the global file the first time cp ~/.memwal/credentials.json .memwal/credentials.json +memwal-mcp approve-project ``` From then on, runs started from that directory or anywhere beneath it use the project's credentials, and runs started outside it keep using the global one. @@ -158,7 +176,9 @@ From then on, runs started from that directory or anywhere beneath it use the pr ### Migration -Nothing to do. Creating a project-local file is the opt-in, so a machine without one behaves exactly as it did before, and the global file remains the fallback indefinitely. +Nothing to do for a machine without a project-local file: it behaves exactly as it did before, and the global file remains the fallback indefinitely. + +If you already have a project-local file, run `memwal-mcp approve-project` in that project once. Until you do, runs from it fall back to the global credentials and say so on stderr — they do not fail. ### Replacing an account @@ -195,6 +215,8 @@ The stdio package accepts CLI flags and environment variables. **CLI takes prece | `--namespace ` (alias `--ns`) | `MEMWAL_NAMESPACE` | Default memory namespace injected into memory tool calls that omit one. See [Default namespace](#default-namespace). | | `--login` (or `login` subcommand) | Not applicable | Force a re-login even when credentials exist. The existing file is kept until the new sign-in succeeds. | | `--logout` | Not applicable | Delete the credentials file currently in use and exit. | +| `approve-project` | Not applicable | Approve the project-local `.memwal/credentials.json` found from the current directory, so it is the file used here. Requires an interactive terminal. See [Approving a project file](#approving-a-project-file). | +| `revoke-project` | Not applicable | Withdraw that approval. Requires an interactive terminal. | | `--help`, `-h` | Not applicable | Print usage and exit. | Set `MEMWAL_MCP_DEBUG=1` to enable verbose stderr logging. @@ -486,7 +508,7 @@ The URL is valid for **5 minutes**. Call the tool again to mint a fresh one. Mak ### Recall returns "No matching memories found" right after a remember -`memwal_remember` waits for the Walrus upload to finish before returning, but under load the embedding/indexing step can lag a few seconds behind. Wait briefly, then retry the recall. +By default `memwal_remember` returns at accept (`job_id`, ~1s). That reply is **not stored yet** — the Walrus write is still in flight. Empty recall immediately after accept is expected. Settle with the job-status tool this server advertises (`memwal_remember_status` on a current sidecar) and only treat `blob_id` as saved. `MEMWAL_MCP_REMEMBER_WAIT_MS=90000` is the opt-in wait-for-`blob_id` path (needs a client that raises the 60s tools/call ceiling). ### 401 Unauthorized from the relayer diff --git a/docs/python-sdk/api-reference.md b/docs/python-sdk/api-reference.md index b7816cb41..612b6b9d5 100644 --- a/docs/python-sdk/api-reference.md +++ b/docs/python-sdk/api-reference.md @@ -30,7 +30,7 @@ questions: - How does Ed25519 authentication work in the MemWal Python SDK? answer: >- The MemWal Python SDK API reference documents all methods on MemWal and MemWalSync - including remember, recall, analyze, ask, restore, health, and lower-level manual methods. + including remember, recall, analyze, ask, restore, list_namespaces, health, and lower-level manual methods. It also covers result dataclasses, exception hierarchy, middleware wrappers, utility functions for delegate key derivation, and the Ed25519 request signing protocol. --- @@ -183,6 +183,37 @@ RestoreResult(restored: int, skipped: int, total: int, namespace: str, owner: st `truncated=true` is known-retryable-incomplete (this call's `limit`, or a still-expandable sidecar candidate fetch); `truncated=false` is not proof the sidecar saw every onchain blob (WALM-451 `sourceCapped`). +### `list_namespaces(cursor=None, limit=None) -> NamespacesResult` + +List the namespaces this account holds memories in. Returns metadata only, with no blob fetch or decryption. + +Recall needs a namespace to search, so an agent on an unfamiliar account would otherwise have to guess names or fall back to `"default"`. Namespaces are flat and exact-match: to work with a prefix such as `proj/`, filter the names client-side and recall each one. + +- `cursor`: the previous page's `next_cursor`, to continue a walk or poll for namespaces changed since then +- `limit`: page size; the relayer defaults to `100` and clamps to `500` + +```python +NamespacesResult( + namespaces: list[NamespaceSummary], # NamespaceSummary(id, name, memory_count, storage_used, updated_at) + next_cursor: str | None, + has_more: bool, + snapshot_version: int, +) +``` + +Paginate on `has_more`, not on page length. The relayer clamps `limit`, so a caller asking for more than the cap gets exactly the cap back. + +```python +cursor = None +while True: + page = await memwal.list_namespaces(cursor=cursor) + for ns in page.namespaces: + print(ns.name, ns.memory_count) + cursor = page.next_cursor + if not page.has_more: + break +``` + ### `health() -> HealthResult` Check relayer health. No authentication — a successful response confirms the relayer is reachable, not that your `key`/`account_id` are valid. A signed call (e.g. `remember()`, `recall()`) can still fail with `401` immediately after a passing `health()`. Raises `MemWalError` on non-200. diff --git a/docs/python-sdk/changelog.mdx b/docs/python-sdk/changelog.mdx index 28977755e..04adee2d2 100644 --- a/docs/python-sdk/changelog.mdx +++ b/docs/python-sdk/changelog.mdx @@ -29,13 +29,21 @@ questions: - What changes were made in memwal 0.1.4? - Where can I find the release history for the Walrus Memory Python SDK? answer: >- - The latest Python SDK release is 0.1.10. `restore()` results include `failed` (default `0`) for permanent decrypt/UTF-8 failures instead of folding them into `skipped`. 0.1.9 reports HTTP 503 as a retryable upstream outage instead of a credential failure, rejects empty `remember_bulk_async` batches and misaligned relayer `job_ids`, aligns restore `truncated` docs with WALM-431 retryable semantics, and warns when `server_url` uses plaintext HTTP on a non-localhost host without logging URL credentials. + The latest Python SDK release is 0.1.11. It adds `list_namespaces()` so an agent can discover which namespaces hold memories instead of guessing. 0.1.10 adds `failed` to `restore()` results for permanent decrypt/UTF-8 failures instead of folding them into `skipped`. 0.1.9 reports HTTP 503 as a retryable upstream outage instead of a credential failure, rejects empty `remember_bulk_async` batches and misaligned relayer `job_ids`, aligns restore `truncated` docs with WALM-431 retryable semantics, and warns when `server_url` uses plaintext HTTP on a non-localhost host without logging URL credentials. --- Track what's new, changed, and fixed in `memwal` (Python). For the latest version, see the [PyPI project page](https://pypi.org/project/memwal/). +## 0.1.11 + +This release adds `list_namespaces()` for namespace discovery. + +### Added + +- `list_namespaces(cursor=None, limit=None)` lists the namespaces that hold memories (name, `memory_count`, `storage_used`, `updated_at`), so an agent can discover namespaces instead of guessing. Metadata only; no decryption. Paginate on `has_more`. `MemWalSync` and the mock clients have it too. + ## 0.1.10 This release adds `failed` on `restore()` results for permanent decrypt and UTF-8 failures. diff --git a/docs/reference/environment-variables.md b/docs/reference/environment-variables.md index cb6d095cf..e7edaeaa4 100644 --- a/docs/reference/environment-variables.md +++ b/docs/reference/environment-variables.md @@ -69,11 +69,13 @@ The stdio MCP package reads these environment variables directly. A CLI flag tak | `MEMWAL_WEB_URL` | `--web-url ` | dashboard default | Dashboard URL used during login | | `MEMWAL_CLIENT_LABEL` | `--label ` | `MCP Client` / `Walrus Memory MCP` | Friendly delegate-key label shown in the dashboard | | `MEMWAL_MCP_DEBUG` | none | `0` | Set to `1` for verbose stderr logging | -| `MEMWAL_CREDS_DIR` | none | `~/.memwal` | Directory holding `credentials.json`. Overrides both project-local and `~/.memwal` credentials, re-read on every access so a test can redirect it after import. Mainly for tests, which must not write into the real credential directory | +| `MEMWAL_CREDS_DIR` | none | `~/.memwal` | Directory holding `credentials.json` and the project-approval record. Overrides both project-local and `~/.memwal` credentials, and needs no project approval, re-read on every access so a test can redirect it after import. Must be an **absolute path outside the current project** — a relative or in-project value is refused with an error, because it skips the approval gate and could otherwise come from a checkout's MCP-config `env` block; an empty value means unset. Mainly for tests, which must not write into the real credential directory | +| `MEMWAL_MCP_TRANSPORT` | none | `sse` | Which relayer transport the stdio bridge dials. `sse` uses the legacy split (`POST /api/mcp/messages` + `GET /api/mcp/sse`). `http` (aliases `streamable`, `streamable-http`) uses the Streamable HTTP endpoint `/api/mcp`, where a call is answered on the same request rather than split across a POST and an SSE stream, so there is no idle watchdog. Reconnect replay is NOT yet transport-aware: the bridge still replays its in-flight map on either transport, and on Streamable a disconnect mid-send can look sent, so a replayed write may duplicate. Opt in knowing that. Unrecognised values fall back to `sse` | | `MEMWAL_MCP_SSE_IDLE_MS` | none | `30000` | Maximum milliseconds of silence on the SSE stream before the bridge treats the session as dead and reconnects. Values below `500` are ignored and fall back to the default. Mainly for tests | | `MEMWAL_MCP_CALL_TIMEOUT_MS` | none | `240000` | Maximum milliseconds a single request might wait for its response before the bridge answers with a retryable error. Covers a reply lost while the stream itself stays healthy, which `MEMWAL_MCP_SSE_IDLE_MS` cannot detect. The default is derived in code from the slowest server-side tool deadline plus headroom, so it moves with that tool rather than being pinned here. Values below `1000` are ignored and fall back to the default | | `MEMWAL_MCP_THROTTLE_FLOOR_MS` | none | `5000` | Minimum milliseconds the bridge waits before retrying an SSE handshake the relayer refused with HTTP 429 and no `Retry-After`. A `Retry-After` on the response wins instead. Either way the wait is capped at `60000`. Non-numeric or negative values are ignored and fall back to the default. Mainly for tests | | `MEMWAL_MCP_STALLED_HANDSHAKE_MS` | none | `90000` | Milliseconds a request may stay buffered while the SSE handshake keeps failing before the bridge answers it with the handshake's own error. Applies only to requests that were never sent, and only while the handshake is failing; a request buffered behind a healthy connection keeps `MEMWAL_MCP_CALL_TIMEOUT_MS`. Clamped to never exceed that call timeout. Values below `1000` are ignored and fall back to the default. Mainly for tests | +| `MEMWAL_MCP_HEALTH_PROBE_MS` | none | `3000` | Milliseconds the bridge waits on the relayer's `/health` before answering a sent request whose reply never arrived, so the error can say whether the relayer is up, unhealthy, or unreachable. Values below `100` are ignored and fall back to the default. Mainly for tests | | `MEMWAL_MCP_LOGIN_TIMEOUT_MS` | none | `300000` | Maximum milliseconds the local sign-in listener stays bound waiting for the browser callback. The default gives you time to review a wallet prompt; shorten it only in tests. Values below `100` are ignored and fall back to the default | ## Self-hosted relayer @@ -130,6 +132,12 @@ These are not all enforced at boot, but most real deployments need them. | `MEMWAL_ACCOUNT_ID` | none | Optional account ID in server config | | `WALRUS_PACKAGE_ID` | network default | Override the Walrus on-chain package used by the sidecar | | `WALRUS_UPLOAD_RELAY_URL` | network default | Override the Walrus upload relay used by the sidecar | +| `WALRUS_DIRECT_UPLOAD` | `false` | When `true`, skip the upload relay and write slivers to storage nodes directly. Prefer the relay when one is available; direct is the fallback path | +| `WALRUS_RELAY_TIMEOUT_MS` | `120000` | Timeout for one upload-relay request. Only read when the relay path is active (`WALRUS_DIRECT_UPLOAD` unset/false and `WALRUS_UPLOAD_RELAY_URL` set) | +| `WALRUS_STORAGE_NODE_TIMEOUT_MS` | `30000` | Timeout for one storage-node HTTP request on both the direct and relay paths (relay still reads slivers back from nodes). Per-node, not per write — raising it also lengthens how long one unresponsive node can stall a shard. Allowed range `1000`–`180000`; keep under `WALRUS_UPLOAD_ACQUIRE_TIMEOUT_MS` | +| `WALRUS_UPLOAD_MAX_CONCURRENCY` | size of `SERVER_SUI_PRIVATE_KEYS` (min 1) | Sidecar-global cap on concurrent Walrus uploads | +| `WALRUS_UPLOAD_PER_WALLET_CONCURRENCY` | `1` | Per-uploader-wallet cap on concurrent Walrus uploads | +| `WALRUS_UPLOAD_ACQUIRE_TIMEOUT_MS` | `120000` | How long an upload waits for a free global/wallet slot before failing acquisition | | `SEAL_SERVER_CONFIGS` | network default | Optional JSON SEAL server config override for independent or committee servers | | `SEAL_KEY_SERVERS` | network default | Legacy comma-separated independent SEAL key server override. Used only when `SEAL_SERVER_CONFIGS` is unset. Deprecated but supported through relayer API `1.x` | | `SEAL_THRESHOLD` | `min(2, total configured weight)` | Required configured server weight for SEAL encrypt/decrypt | @@ -154,6 +162,8 @@ These are not all enforced at boot, but most real deployments need them. | `MCP_MAX_TOTAL_SESSIONS` | `1000` | Maximum active MCP sessions across SSE and Streamable HTTP transports | | `MCP_MAX_SESSIONS_PER_IP` | `16` | Maximum active MCP sessions from one source IP | | `MCP_MAX_NEW_SESSIONS_PER_IP_PER_MIN` | `30` | Maximum new MCP sessions opened by one source IP per minute | +| `MEMWAL_MCP_REMEMBER_WAIT_MS` | `0` | How long `memwal_remember` / `memwal_remember_bulk` / `memwal_analyze` wait after accept for the Walrus write to reach `done`. **Default `0` = return at accept (~1s) with `job_id`; the fact is NOT saved yet** — resolve with `memwal_remember_status`. This is what stops the agent blocking 20–90s (or hitting the MCP client's 60s tools/call timeout). `90000` (90 seconds, the clamp) restores the old wait-for-`blob_id` behaviour; only use it on a client that raises that 60s ceiling (Claude Code does; many hosts do not). Values between 0 and ~75s usually wait *and* still return pending. Invalid values fall back to `0` | +| `MEMWAL_MCP_ACCEPT_DEADLINE_MS` | `15000` | How long a single relayer request may stall before the MCP tool gives up on it. The SDK passes an abort signal only on `recall`, so `rememberAsync` and the job-status reads have no deadline of their own and `MEMWAL_MCP_REMEMBER_WAIT_MS` bounds only when the next poll starts, not how long one takes — without this a stalled socket keeps a tool running indefinitely. Raise it only if a slow link makes healthy accepts exceed it; invalid or non-positive values fall back to the default | | `MCP_TOOL_SLOW_WARN_MS` | `5000` | An MCP tool call still running at this duration is logged as `tool.slow` (`settled: false`) at `warn` — which is what makes a hang visible, since a hang never settles. A call that finishes at or above it is logged as `tool.slow` (`settled: true`) instead of `tool.done` | | `TRUSTED_PROXY_HOPS` | `0` | Number of trusted reverse-proxy hops to walk from the right of `X-Forwarded-For`; `0` ignores XFF and uses the TCP peer | | `WRITES_PAUSED` | `false` | When `1` / `true` / `yes` / `on`, write routes (`POST /api/remember`, `/api/remember/manual`, `/api/remember/bulk`, `/api/analyze`) return HTTP 503 `{"error":"writes are paused"}`. `GET /health` stays HTTP 200 with `status: "ok"` and `writes: "paused"`. Reads (`recall`, `restore`, health) stay available | diff --git a/docs/relayer/api-reference.md b/docs/relayer/api-reference.md index 55a6446a7..b2143f777 100644 --- a/docs/relayer/api-reference.md +++ b/docs/relayer/api-reference.md @@ -82,9 +82,13 @@ These routes require no authentication. Service liveness check. `status` is `"ok"` when the relayer process is up. HTTP 200 means the process is running, not that writes are accepted. -`writes` is `"ok"` or `"paused"`. `"paused"` when `WRITES_PAUSED` is set (`1` / `true` / `yes`); empty or unset is `"ok"`. That flag is write-path admission, not a health-only signal: `POST /api/remember`, `/api/remember/manual`, `/api/remember/bulk`, and `/api/analyze` then return HTTP 503 with `{"error":"writes are paused"}`. `/health` itself stays HTTP 200 with `status: "ok"` and `writes: "paused"`, so clients can distinguish an intentional pause from an integrator bug. Reads (`recall`, `restore`, remember job status) stay available. +`writes` is `"ok"`, `"degraded"`, or `"paused"`. -`write_ready` is `true` when the encryption sidecar process answered its own `/health` **and** Postgres can accept writes (cached a few seconds). Postgres is considered not writable when Neon cluster size (`pg_cluster_size`, not `pg_database_size` of this database) is at or within 1MB of `neon.max_cluster_size`. If `pg_cluster_size()` is missing (no `neon` extension), the probe falls back to `sum(pg_database_size)` against that GUC. Self-hosted Postgres without the GUC keeps the sidecar-only check. Probe errors and timeouts fail open (`write_ready` stays true) so CI is not blocked; timeouts log at warn. A sidecar outage or a disk/project-size write outage can still return HTTP 200 with `write_ready: false`. Use `writes`, not `write_ready`, for the pause signal. `write_ready: true` is not a guarantee that remember or analyze succeed. +- `"paused"` when `WRITES_PAUSED` is set (`1` / `true` / `yes`). That flag is write-path admission, not a health-only signal: `POST /api/remember`, `/api/remember/manual`, `/api/remember/bulk`, and `/api/analyze` then return HTTP 503 with `{"error":"writes are paused"}`. `/health` itself stays HTTP 200 with `status: "ok"` and `writes: "paused"`, so clients can distinguish an intentional pause from an integrator bug. Reads (`recall`, `restore`, remember job status) stay available. +- `"degraded"` when recent durable writes (last 15 minutes) have **failed at least three times and none have landed**. The relayer still accepts and queues a write (`write_ready` stays `true`; HTTP 200), but Walrus is not storing it — expect the job to fail rather than queue more. A single success in that window keeps `"ok"` even if other writes failed: this is a total-outage detector, not a per-request verdict. `memwal_health` prints `writes=degraded` when this is set. +- `"ok"` otherwise (including a quiet window with no finished writes). + +`write_ready` is `true` when the encryption sidecar process answered its own `/health` **and** Postgres can accept writes (cached a few seconds). Postgres is considered not writable when Neon cluster size (`pg_cluster_size`, not `pg_database_size` of this database) is at or within 1MB of `neon.max_cluster_size`. If `pg_cluster_size()` is missing (no `neon` extension), the probe falls back to `sum(pg_database_size)` against that GUC. Self-hosted Postgres without the GUC keeps the sidecar-only check. Probe errors and timeouts fail open (`write_ready` stays true) so CI is not blocked; timeouts log at warn. A sidecar outage or a disk/project-size write outage can still return HTTP 200 with `write_ready: false`. Use `writes`, not `write_ready`, for the pause or Walrus-outage signal. `write_ready: true` is not a guarantee that remember or analyze succeed. **Response:** @@ -159,6 +163,26 @@ Proxy to the sidecar's `/sponsor/execute` endpoint. `sender` must match the shor Every route below requires the signed headers described in [Authentication](#authentication). +### `GET /api/whoami` + +Return the account identity the caller's delegate key resolves to. Takes no request body. + +Authentication already resolves the account before any handler runs, so this route just hands back what the middleware computed. Returning `account_id` is safe here precisely because the route is authenticated. The caller has proven it holds a delegate key registered against this account, so it only ever learns about itself. The public `GET /api/accounts/:owner/exists` route deliberately withholds it. + +The motivating use is rebuilding local credentials: a client that holds a working delegate key but has lost the surrounding metadata (an interrupted sign-in, a wiped config file) needs `account_id`, `owner`, and `package_id` to write a usable credentials file, and the key alone proves entitlement to all three. + +**Response:** + +```json +{ + "account_id": "0x...", + "owner": "0x...", + "package_id": "0x..." +} +``` + +**Mainnet only, when the caller cannot send `x-account-id`.** Recovering a lost account id is the one case where the client has no id to send, so authentication has to find it by scanning the `AccountRegistry` for the delegate key. That scan runs over Sui JSON-RPC, which Testnet no longer serves, so Testnet requires the `x-account-id` hint for delegate-key authentication and rejects the request with `401` when it is absent, including this one. A caller that already knows its account id can use this route on either network; a caller recovering one cannot use it on Testnet. + ### `POST /api/remember` Submit text as an encrypted memory job. The relayer returns after creating a background job; embedding, Seal encryption, Walrus upload, and vector indexing continue asynchronously. @@ -269,6 +293,22 @@ Search for memories matching a natural language query. Returns decrypted plainte `limit` defaults to `10`; the server caps it at `100`. `namespace` defaults to `"default"`. `scoring_weights` takes an optional object; omit it to keep the plain cosine-distance order. +`sort` is optional: `"relevance"` (the cosine order, and the behaviour when omitted) or `"recent"` (the newest among the semantic matches). An explicit `sort`, `"relevance"` included, is the order, and the relayer ignores `scoring_weights` for that request. + +`deadline_ms` is optional: how long the caller waits for this response, in milliseconds. When set, the relayer stops about one second before it, counting from when the request arrived, and answers `504` with the step that was still running, so the caller learns where the recall stalled instead of timing out blind. A short deadline still gets at least 2 seconds of work; if the whole deadline went on authentication before the recall started, the relayer answers at once with stage `auth`. Values above `600000` are capped. Omit it to let the recall run to completion. The TypeScript SDK sends `14000`, a second under the 15s it aborts `recall()` at: the relayer's own margin runs from arrival, so it covers the reply's trip back but not the connect the caller's timer already started on. + +```json +{ + "error": "Recall timed out after 14001ms during walrus_download", + "message": "Recall timed out after 14001ms during walrus_download", + "code": "RECALL_TIMEOUT", + "stage": "walrus_download", + "elapsed_ms": 14001 +} +``` + +`stage` is one of `auth`, `embed`, `vector_search`, `walrus_download`, or `seal_decrypt`. + #### Scoring weights The optional `scoring_weights` object turns on composite ranking. The same object works on `/api/recall`, `/api/recall/manual`, and `/api/ask`. @@ -297,7 +337,7 @@ The optional `scoring_weights` object turns on composite ranking. The same objec } ``` -`score` only appears when `scoring_weights` sets a nonzero `recency` or `importance` weight. A request that sets only the `semantic` weight keeps the plain cosine order, and the relayer omits `score`. `dropped_count` only appears when at least one match dropped out because its blob download or decryption failed; the relayer omits those matches from `results`. +`score` only appears when `scoring_weights` sets a nonzero `recency` or `importance` weight and `sort` is omitted. A request that sets only the `semantic` weight keeps the plain cosine order, and the relayer omits `score`. `dropped_count` only appears when at least one match dropped out because its blob download or decryption failed; the relayer omits those matches from `results`. ### `POST /api/remember/manual` diff --git a/docs/relayer/observability.md b/docs/relayer/observability.md index d465d2dd8..1563bb28d 100644 --- a/docs/relayer/observability.md +++ b/docs/relayer/observability.md @@ -75,12 +75,15 @@ The Rust relayer exposes Prometheus metrics at: GET /metrics ``` -The TypeScript sidecar also exposes wallet-specific counters at: +The TypeScript sidecar also exposes wallet-specific counters and upload-queue counters at: ```text GET /metrics/wallet +GET /metrics/uploads ``` +`/metrics/uploads` returns `activeWalrusUploads`, `queuedWalrusUploads`, and `walrusUploadLimits` from in-memory counters, with no Sui or Walrus calls. + Core relayer metrics: | Metric | Labels | Notes | @@ -133,6 +136,7 @@ Create panels for: | DB saturation | PostgreSQL pool open connections near configured max, or idle connections stay at 0 | | Wallet lock canary | Sidecar `walletLockErrorsTotal` is greater than 0 | | Permanent wallet failures | Sidecar `walletPermanentFailuresTotal` increases | +| Upload queue saturation | Built in: the relayer polls sidecar `/metrics/uploads` and alerts Slack when `queuedWalrusUploads` stays above `SIDECAR_QUEUE_SATURATION_THRESHOLD` (default 20) for `SIDECAR_QUEUE_SATURATION_CONSECUTIVE` checks (default 4) polled every `SIDECAR_QUEUE_SATURATION_INTERVAL_SECS` (default 30) | ## APM Integration diff --git a/docs/sdk/api-reference.md b/docs/sdk/api-reference.md index 3e3054c66..ec931610e 100644 --- a/docs/sdk/api-reference.md +++ b/docs/sdk/api-reference.md @@ -203,6 +203,9 @@ Two limits worth knowing: Omitting `sort` leaves the request byte-identical to a plain cosine recall, so existing callers see no change. +An explicit `sort`, `"relevance"` included, is the order: the relayer ignores +`scoringWeights` for that request. Weights re-rank only when `sort` is omitted. + #### `scoringWeights` `scoringWeights` blends recency and importance into the relayer's ranking: diff --git a/docs/sdk/changelog.mdx b/docs/sdk/changelog.mdx index 27064137e..5f2862a98 100644 --- a/docs/sdk/changelog.mdx +++ b/docs/sdk/changelog.mdx @@ -28,12 +28,23 @@ questions: - When was bulk remember added to the Walrus Memory SDK? - What security improvements have been made to the MemWal SDK? answer: >- - The latest TypeScript SDK release is 0.1.7. Request bodies are hashed with `@noble/hashes` rather than WebCrypto-or-`node:crypto`, so the SDK no longer imports a Node builtin that Vite silently externalises into a runtime crash in the browser, and it declares a Node 20 floor. `restore()` results include `failed` (required like `truncated`; SDK defaults omitted to `0`) for permanent decrypt/UTF-8 failures instead of folding them into `skipped`. Empty-body 401s use the AUTH_REJECTED troubleshooting message instead of telling callers to run `memwal_login`. Account and manual PTBs use typed `tx.pure` helpers so they work with modern `@mysten/sui`. 0.1.6 added optional `created_at` on `recall()` results, plus `sort` and `scoringWeights` on `RecallOptions`, and reports HTTP 503 as a retryable upstream outage instead of a sign-in failure. + The latest TypeScript SDK release is 0.1.8. Request bodies are hashed with `@noble/hashes` rather than WebCrypto-or-`node:crypto`, so the SDK no longer imports a Node builtin that Vite silently externalises into a runtime crash in the browser, and it declares a Node 20 floor. `restore()` results include `failed` (required like `truncated`; SDK defaults omitted to `0`) for permanent decrypt/UTF-8 failures instead of folding them into `skipped`. Empty-body 401s use the AUTH_REJECTED troubleshooting message instead of telling callers to run `memwal_login`. Account and manual PTBs use typed `tx.pure` helpers so they work with modern `@mysten/sui`. An explicit `sort` on `recall()`, `"relevance"` included, now makes the relayer ignore `scoringWeights`; weights re-rank only when `sort` is omitted. 0.1.6 added optional `created_at` on `recall()` results, plus `sort` and `scoringWeights` on `RecallOptions`, and reports HTTP 503 as a retryable upstream outage instead of a sign-in failure. --- +## 0.1.8 + +### Added + +- `recall()` sends its own deadline as `deadline_ms` (14000, a second under the 15s it aborts at, so the relayer's answer beats the abort even on a slow connect) in the request body. A relayer that reads it answers a recall about to miss that deadline with a 504 whose `serverCode` is `RECALL_TIMEOUT` and whose body names the stuck step (`auth`, `embed`, `vector_search`, `walrus_download`, `seal_decrypt`), instead of the request aborting with no cause. Relayers that predate the field ignore it. (WALM-396) + +### Fixed + +- Every relayer request now carries a deadline. `fetch` has none of its own and the SDK passed an abort signal on exactly one method (`recall`, 15s), so the accept POST, every job-status read, and the `/version` and `/config` handshake calls could stay pending for as long as the socket stayed open. A poll loop checks its budget at the *top* of each iteration, which bounds when the next request starts rather than how long one takes — so a single stalled read ran straight past `timeoutMs`, and a `memwal_remember` documented as capping at 90s was observed by an MCP client still running after 120s. Requests default to 30s, matching the relayer's own outbound client; set it with `requestTimeoutMs` on `MemWal.create`, where a non-positive or non-finite value falls back to the default rather than disabling the bound. `restore` (60s) and `analyze` (60s) carry their own, since the route self-bounds at 55s and the extractor LLM runs inline respectively. Inside the wait loops each poll is bounded by the client deadline clamped to the remaining budget, so one stalled poll can neither outlive the wait nor swallow it. An expired request raises `MemWalRequestTimeout` with `status: 504`, which the existing transient-poll handling already retries; a caller's own abort and every other transport error propagate unchanged. +- Generated idempotency keys are derived from the content (a 30-minute bucket plus namespace and text) instead of `crypto.randomUUID()`. The per-instance key map only ever deduped retries that reused one client, and callers such as the MCP sidecar build a fresh client per session — so a replay after a reconnect read as a brand-new write and the relayer minted a second paid Walrus blob for one already in flight. The bucket bounds the collapse, since `remember_jobs` rows are never pruned. Callers passing an explicit `idempotencyKey` are unaffected, and distinct text or namespaces still derive distinct keys. + ## 0.1.7 -This release removes the Node `crypto` import that crashed bundled browser builds at runtime, adds `failed` on `restore()` results, declares a Node 20 floor, stops telling headless SDK clients to call `memwal_login` on empty-body 401s, and switches account and manual PTBs to typed `tx.pure` helpers. +This release removes the Node `crypto` import that crashed bundled browser builds at runtime, adds `failed` on `restore()` results, declares a Node 20 floor, stops telling headless SDK clients to call `memwal_login` on empty-body 401s, switches account and manual PTBs to typed `tx.pure` helpers, and makes an explicit `sort` on `recall()` win over `scoringWeights`. ### Added @@ -45,6 +56,7 @@ This release removes the Node `crypto` import that crashed bundled browser build - Declare `engines.node >= 20.0.0`, matching `memwal-mcp` and `openclaw-memory-memwal`. The SDK was the only published package without a floor. (WALM-599) - Empty-body 401s now use the same AUTH_REJECTED troubleshooting message as credential 401s instead of telling callers to run `memwal_login`. Headless SDK clients do not have that MCP tool. - `account.ts` and `manual.ts` PTBs use typed `tx.pure` helpers instead of the legacy untyped moveCall argument syntax that fails under modern `@mysten/sui`. +- An explicit `sort` on `recall()`, `"relevance"` included, is now the order: the relayer ignores `scoringWeights` for that request, and weights re-rank only when `sort` is omitted. Setting both used to return neither order, so `sort: "recent"` stopped meaning newest-first once `scoringWeights` carried a recency weight. (WALM-470) ## 0.1.6 diff --git a/docs/troubleshooting/overview.md b/docs/troubleshooting/overview.md index f97660bd1..cbd0be152 100644 --- a/docs/troubleshooting/overview.md +++ b/docs/troubleshooting/overview.md @@ -88,13 +88,17 @@ This section covers problems that appear before the memory tools work. **Symptom:** The sign-in page confirms that your delegate key was registered, but says it could not hand the credentials back to your computer. The agent stays logged out and `~/.memwal/credentials.json` does not appear. -**Cause:** Signing in has two halves. Your browser registers a delegate key onchain, then sends that key back to a short-lived listener the MCP package runs on `127.0.0.1`. The unused key from this attempt is already on your account and should be revoked. Signing in again is a full new attempt, including the wallet step. The usual reasons: +**Cause:** Signing in has two halves. Your browser registers a delegate key onchain, then sends that key back to a short-lived listener the MCP package runs on `127.0.0.1`. The key from this attempt is already on your account, and the MCP package saved its private half before it opened the sign-in page, so you can still use it. The usual reasons the hand-off fails: - The MCP client restarted, or the login command was cancelled, while the browser tab was still open. - This tab is leftover from a sign-in that already finished, or the hand-off did not match what the app expected. - Local software such as a firewall, a VPN client, or a browser extension blocks requests from a website to `127.0.0.1`. -**Fix:** Call `memwal_login` again and open the new URL promptly. A retry only helps once the MCP client is left running through the wallet prompt. Remove the unused key from the Delegate keys panel in the dashboard; it is already on your account. +**Fix:** Restart your MCP client within 24 hours. On start, the MCP package finds the saved key, confirms it with the relayer, and signs you in with it, with no second wallet step. Do not call `memwal_login` first. It reuses the same key, and the wallet step fails because that key is already registered. Do not remove the key from the dashboard either, unless you mean to abandon it. + +On Testnet the relayer cannot confirm the key at start, so the MCP package prints a notice instead of signing you in. Remove the key it names from the Delegate keys panel in the dashboard, then call `memwal_login` and open the new URL promptly. + +MCP package versions before 0.0.13 do not save the key before the browser step. On those versions, remove the unused key from the dashboard and call `memwal_login` again. If it keeps failing, confirm that nothing blocks localhost traffic, then run `npx -y @mysten-incubation/memwal-mcp login --prod` directly in a terminal. A terminal sign-in prints the failure reason instead of leaving it in the MCP client's logs. @@ -124,9 +128,9 @@ A client-side timeout does not mean the save failed. The relayer accepts the wor **Symptom:** You save a memory, and an immediate recall finds nothing. -**Cause:** A save returns once the Walrus upload completes, but the embedding and indexing step can lag a few seconds behind under load, so the memory is briefly unsearchable. +**Cause:** MCP `memwal_remember` returns at accept by default (`job_id`). The fact is not on Walrus yet, so recall has nothing to find. Even after the write lands, embedding/indexing can lag a few seconds under load. -**Fix:** Wait a moment, then retry the recall. If a memory is missing from the search index later, `memwal_restore` rebuilds the index for that namespace from Walrus. +**Fix:** Do not treat the accept reply as stored. Settle with `memwal_remember_status` (or wait until the tool returns a `blob_id`). Then retry recall. If a memory is missing from the search index later, `memwal_restore` rebuilds the index for that namespace from Walrus. `MEMWAL_MCP_REMEMBER_WAIT_MS=90000` restores wait-for-`blob_id` on a client that raises the 60s tools/call ceiling. ## Quick reference diff --git a/package.json b/package.json index 6f7a04681..bab949b5e 100644 --- a/package.json +++ b/package.json @@ -16,6 +16,7 @@ "preview:docs": "pnpm --filter memwal-docs serve", "verify:memwal": "tsx scripts/verify-memwal-credentials.ts", "check:compatibility": "node scripts/check-compatibility-contract.mjs", + "check:mcp-pin": "node scripts/check-mcp-pin.mjs", "tx:publish": "tsx scripts/build-publish-tx.ts", "tx:create-caps": "tsx scripts/build-migration-caps-tx.ts", "tx:finalize": "tsx scripts/build-finalize-tx.ts", diff --git a/packages/mcp/AUTO-MEMORY.md b/packages/mcp/AUTO-MEMORY.md index 133642c9b..eae420849 100644 --- a/packages/mcp/AUTO-MEMORY.md +++ b/packages/mcp/AUTO-MEMORY.md @@ -14,6 +14,96 @@ package. `memwal_remember` literally said *"Call ONLY when the user explicitly asks… agents should not call this proactively."* And `rememberBulk` (in the SDK) was never exposed as a tool. +## Consent and secret filtering (WALM-642) + +Automatic saving is **on** — it is what MemWal is for — but only once a human +has been asked. What is governed is narrow: telling the model to save something +the user did **not** ask it to save. A direct request ("remember that ...") +works either way, and recall is never gated. + +**Three states, not two.** `on` and `off` are answers; `unset` is the absence of +one, and it does not mean `off`: + +| State | Resolves to | Why | +|---|---|---| +| `autoSave: true` / `false` in settings.json | that | answered, at login or via the CLI | +| unset, no settings file, credentials present | **on** | predates this change; these users have been auto-saving all along, and switching them off would be a regression dressed up as caution. Asked at their next interactive login. | +| unset, `autoSaveConsent: "pending"` | **off** | a post-change install nobody has asked yet. Explicit tool calls and recall keep working. | + +The `pending` stamp is what separates the last two. It is written the moment a +new install first appears — before the signed-out server boots, and immediately +after a first interactive login, *before* the question is put. Without it a +headless install could sign in through the `memwal_login` tool, become +indistinguishable from a long-standing user, and start saving with nobody ever +having been asked. + +**Where the question is asked** — `memwal-mcp login`, on a TTY, and nowhere else +(`src/consent.ts`). Deliberately **not** an MCP tool, a tool description or an +instruction: a model answering on the user's behalf is not consent, and an +agent-shaped surface for this question would be exactly that. A test greps the +compiled tool surfaces to keep it that way. Non-TTY runs print one line saying +where things stand and never block on stdin. + +**Setting it directly** + +```sh +memwal-mcp auto-save on # persists {"autoSave": true} to settings.json +memwal-mcp auto-save off +memwal-mcp auto-save # report the current state and where it came from +``` + +`MEMWAL_AUTO_SAVE=1` in an MCP client's `env` block overrides the file for one +server process, and counts as a deliberate answer — it also stops the login +prompt, so a configured install is never nagged. + +**Where the state lives** — `settings.json`, next to `credentials.json`, so it +inherits `credsPath()` resolution: `MEMWAL_CREDS_DIR` override, else the +nearest project-local `.memwal/`, else `~/.memwal/`. It has to be on disk +rather than passed as configuration because the **hooks are spawned by the +client, not by this package**: they inherit the MCP server's `env` block from +nothing at all. `src/auto-save.ts` and `plugin/scripts/lib/auto-save.mjs` are +two implementations of the same resolution, pinned against each other by +`test/auto-save-optin.test.mjs`. + +**What changes when it is off** — the `instructions` field, the SessionStart +rubric, the UserPromptSubmit rubric and the PostToolUse nudge all switch to a +save-only-what-you-are-asked variant. Nothing is disabled; the guidance that +drives an unasked-for save is simply not injected. While consent is outstanding +the SessionStart banner also says so, and tells the agent the answer is given in +a terminal — not in chat, and not by it. + +**One source for the rules** — the secret-exclusion and do-not-save text lives +in a single block duplicated byte-for-byte across three files that cannot +import each other: + +| Copy | Feeds | +|---|---| +| `packages/mcp/src/memory-policy.ts` | `instructions`, cold-start `tools/list` | +| `packages/mcp/plugin/scripts/lib/memory-policy.mjs` | the three lifecycle hooks | +| `services/server/scripts/mcp/tools/memory-policy.ts` | live tool descriptions, sidecar `instructions` | + +`test/memory-policy.test.mjs` extracts the marked block from each file and +compares the bytes, so editing one copy fails the suite until the others match. + +**The programmatic backstop** — model-facing rules are not enforcement, and +Walrus storage is append-only: a secret that lands cannot be deleted. So +`services/server/scripts/mcp/tools/redaction.ts` screens every write **before** +the text reaches the SDK, on all three write paths (`memwal_remember`, +`memwal_remember_bulk`, `memwal_analyze`). A mixed message keeps its fact and +loses only the credential span — the ticket's "save safe facts without +neighboring credentials" — and a text that is nothing but a secret, or that the +user said not to save, is not forwarded at all. The caller is told which *kinds* +were removed; the value is never logged, echoed, or returned. + +Detection is shape-based, not entropy-based, on purpose: MemWal's own durable +facts (blob ids, Sui object ids, git SHAs, digests) are exactly what a generic +high-entropy rule would eat. Key material in hex is therefore caught by the +**label** beside it rather than by how random it looks — which is what lets the +`delegatePrivateKey` from `credentials.json` (64 lowercase hex, the value +`auth.ts` marks "NEVER log this") be removed while a bare commit SHA or `0x` +object id is left alone. The trade-off is written out at the top of +`redaction.ts`. + ## Architecture — three layers 1. **Agentic tool descriptions + `memwal_remember_bulk`** — `services/server/scripts/mcp/tools/`. @@ -30,7 +120,8 @@ was never exposed as a tool. | Dimension | Before | After | |---|---|---| -| Save trigger | "ONLY when user explicitly asks; don't be proactive" | "Save proactively whenever you learn a durable fact" | +| Save trigger | "ONLY when user explicitly asks; don't be proactive" | "Save proactively whenever you learn a durable fact" — **since WALM-642, after a consent question at login** | +| Secret handling | none: a preference next to a password was forwarded whole | shared exclusion rules on all three surfaces + a redactor in front of every write | | Bulk save | not exposed | `memwal_remember_bulk` (wraps SDK `rememberBulkAndWait`, ≤20) | | Recall trigger | neutral; agent rarely called it unprompted | "Recall proactively at task start / when the user references past work" | | Reinforcement | none | UserPromptSubmit + PostToolUse hooks (Claude Code + Codex) | @@ -51,7 +142,12 @@ was never exposed as a tool. ## Decisions (chosen) -- **Append-only** — no `forget`/`update` tools (relayer dedups embeddings). +- **Append-only** — no `forget`/`update` tools (relayer dedups embeddings). This + is also why WALM-642's credential check runs *before* the write: there is no + delete to fall back on. +- **Automatic saving is on, after consent** (WALM-642) — asked once at + interactive login, never through an agent-reachable surface. Explicit tool use + is never gated, and neither is recall. - **Global `default` namespace** — `MEMWAL_NAMESPACE` overrides for per-project scope. - **Agent decision rubric** — UserPromptSubmit does not regex-classify remember vs recall. The agent has the conversation and understands any language or spelling. @@ -63,6 +159,10 @@ was never exposed as a tool. ## File map - `services/server/scripts/mcp/tools/{remember,recall,analyze,restore}.ts` — agentic descriptions +- `services/server/scripts/mcp/tools/redaction.ts` — pre-forward credential screen (WALM-642) +- `{packages/mcp/src,packages/mcp/plugin/scripts/lib,services/server/scripts/mcp/tools}/memory-policy.*` — the shared rules block, three byte-identical copies +- `packages/mcp/src/auto-save.ts` + `packages/mcp/plugin/scripts/lib/auto-save.mjs` — the tri-state resolver, server side and hook side +- `packages/mcp/src/consent.ts` — the login-time consent question; TTY-only, never agent-reachable - `services/server/scripts/mcp/tools/remember-bulk.ts` + `index.ts` — new bulk tool - `packages/mcp/plugin/` — plugin manifest, `.mcp.json`, hooks, Node scripts, Codex installer - `.claude-plugin/marketplace.json` (repo root) — Claude Code marketplace entry (local source `./packages/mcp/plugin`) diff --git a/packages/mcp/CHANGELOG.md b/packages/mcp/CHANGELOG.md index 309ffa8b0..169bd1139 100644 --- a/packages/mcp/CHANGELOG.md +++ b/packages/mcp/CHANGELOG.md @@ -1,5 +1,41 @@ # @mysten-incubation/memwal-mcp +## Unreleased + +### Security + +- A project-local `.memwal/credentials.json` no longer decides where memory goes on presence alone. That file lives inside the repository, so anyone who could commit to a repo — or get a clone opened — could silently repoint the account and relayer every memory written from that directory went to, including from a subfolder, with nothing said and no delete path once written. A project file is now inert until the user approves that exact project path, account, delegate key and relayer with `memwal-mcp approve-project`; the approval record is kept in `~/.memwal/project-approvals.json`, outside the repository, so a repository cannot carry its own approval, and any later change to the destination requires approving again. Until approved the global credentials are used — an unapproved, altered or malformed project file is a fallback, never a failure — and one stderr line names the file that was skipped, the destination it wanted, and the command that approves it. `MEMWAL_CREDS_DIR` still overrides both files without approval. `memwal_health` now reports `account=` beside `relayer=`, so the active destination is visible where the user is. (WALM-639) + +- `MEMWAL_CREDS_DIR` must now be an absolute path outside the current project, and a value that is not is refused with an error naming it rather than followed. The override decides the credentials before any approval is looked up, so it was a way around the gate above rather than through it: an empty value made the approval store the bare relative name `project-approvals.json`, which resolves against the working directory — a repository that committed that one file at its root approved its own credentials — and a relative value such as `.memwal` put both the credentials and the approval record inside the repository. Neither had to be typed by the user: an MCP client reads `.cursor/mcp.json`, `.vscode/mcp.json` or `.claude/settings.json` out of the checkout and passes on the `env` block it finds, and editors expand `${workspaceFolder}` there, so an absolute path inside the project is refused too. An empty value now means unset, matching every other reader of the variable. Tests and CI pointing at an absolute temp directory are unaffected. (WALM-639) +- The login write-ahead record no longer goes into the repository. It sat beside whichever credentials file resolution chose, so once a project file was approved every sign-in wrote a plaintext 64-hex Ed25519 delegate seed to `/.memwal/login-pending.json` — in a directory an attacker who planted the credentials file had already created, so tracked rather than ignored, and staged by `git add -A`. It now lives in `~/.memwal/login-pending/`, keyed by project so a sign-in is still reclaimable only by the project that started it. The credentials file itself cannot move — it is the file the user approved — so `approve-project` and the pre-sign-in warning now state plainly that signing in from this project writes a delegate private key inside the repository, and suggest adding `.memwal/` to `.gitignore`. (WALM-639) + +## 0.0.14 + +Unreleased. The plugin's launcher installs `@mysten-incubation/memwal-mcp@0.0.14-dev.0`, named by `mcpPackageVersion` in `plugin/plugin.json`, until `0.0.14` is on npm; drop that key in that release. + +### Added + +- Opt-in Streamable HTTP transport for the stdio bridge. Set `MEMWAL_MCP_TRANSPORT=http` (aliases `streamable`, `streamable-http`) to dial the relayer's single `/api/mcp` endpoint, where a call is answered on the same request instead of being split across a POST and an SSE stream, so there is no idle watchdog. The default stays `sse` and unrecognised values fall back to it. Reconnect replay is not transport-aware yet: on Streamable a disconnect mid-send can look sent, so a replayed write may duplicate. + +### Fixed + +- The plugin no longer starts the MCP server through `npx @mysten-incubation/memwal-mcp@`. npx resolves a package name against the directory the MCP client was started in — the project the user has open — so a project carrying an installed package of that name, claiming the pinned version, was run instead of the published one; a pinned `npx …@0.0.14` command was reproduced running a project's own binary while offline. The version pin was no defence, because the planted package simply claims that version. Every launch site (`.mcp.json`, the Cursor and Codex copies, and the Codex fallback installer) now runs `plugin/scripts/launch_mcp.mjs`, which installs the pinned version once under `~/.memwal/runtime/memwal-mcp@` and launches that absolute entry point with the current node binary. Nothing on that path consults the project's `node_modules`, a PATH-relative bin shim, or the project's `.npmrc`, and the launcher fails instead of falling back to the name when the trusted install cannot be produced. Exploiting the old behaviour required write access to installed package and bin files inside the project, so a plain repository clone or a `package.json` alone was never enough. `MEMWAL_MCP_RUNTIME_DIR` relocates the trusted directory and is rejected unless it is an absolute path. (WALM-640) +- The trusted launcher installs `0.0.14-dev.1`, not `dev.0`. `dev.0` predates the relayer change that returns `memwal_remember` at accept, so its initialize `instructions` still call the `job_id` reply "the exception, not the normal result" and say a `blob_id` means the fact is stored — the opposite of what the relayer now does on every write. An agent on `dev.0` against a current relayer therefore reports a fact as saved while the write is still in flight. `dev.1` carries the corrected text and differs from `dev.0` in only auth-required, bridge and instructions. +- An unfinished write is no longer reported with a `blob_id`. `persist_uploaded_state` stamps a real `blob_id` and status `uploaded` the moment the blob is on Walrus, but the job is not done — SetMetadataAndTransfer has not handed the blob object to the user and `insert_vector` has not indexed it, so the fact is not recallable and a later failure can still lose it. `uploaded` was already counted as in flight; the render was not, and printed the `blob_id` of any row that had one. A live dev batch read `0/6 saved, 6 still uploading` with three of those rows carrying a `blob_id=`, contradicting the same reply's rule that only a `blob_id` means the fact is stored — an agent that believes it tells the user a fact is saved while it is still in flight. Only a `done` row shows a `blob_id` now, which is what the single-job path always did. Same guard in `memwal_remember_bulk` and `memwal_analyze`, whose timeout rows render from the same template. (#921 follow-up) +- The cold-start tool list no longer advertises a tool the relayer may not serve. The bridge ships on npm and updates itself while a relayer ships per environment, so 0.0.14-dev.0 dialled prod and staging still on 0.0.13: cold start named `memwal_remember_status`, which neither registers, and the pending-write wording sent the agent to go call it — one live run spent 90.67s there before erroring. Cold start is now a floor rather than a forecast (`BASELINE_RELAYER_TOOLS`): it carries only what the oldest supported relayer serves and its descriptions name nothing outside it, while newer tools still reach the client a beat later on the relayer's own `tools/list`. Initialize `instructions` use the same floor (they used to name `memwal_remember_status` before any `tools/list`). And a call for a tool outside that floor — including during the cold-start window, before any upstream `tools/list` has been seen — is now answered locally and at once — naming the tools that do exist and saying plainly that nothing ran — instead of being forwarded into a wait that only ends at the orphan deadline. (#928) +- A tool call that times out or cannot reach the relayer now says why and what to do, in three lines: `Cause`, `Relayer health` (checked against `/health` on the spot) and `Next step`. It used to read `Tool error: This operation was aborted`. A recall the relayer cut short names the step it was stuck in: credential check, embedding, vector search, Walrus download or SEAL decrypt. Writes are never told a retry is safe. (WALM-396) +- A call whose reply never arrives is answered with the relayer's health. The bridge checks `/health` first and says whether the relayer is up with this one call stuck, unhealthy, down, or not resolvable at the configured URL, instead of a bare "did not answer". A call being answered this way is never re-sent, so a reconnect landing while the check runs cannot run a write a second time. A lost `memwal_recall` reply is now answered after 2 minutes instead of 4. (WALM-396) +- Bound every relayer call the tools make. The pinned SDK aborts a request only when the caller passes a signal, which none of the memory methods do, so a stalled socket kept a tool running with no ceiling — `memwal_remember` was observed still going past 120s against a 90s budget. Accepts are bounded at 15s (`MEMWAL_MCP_ACCEPT_DEADLINE_MS`), waits at their own budget plus grace. The request is not cancelled — the SDK exposes no way to pass a signal — but the agent is no longer held by it. +- Honour the relayer's `retry_after` instead of dropping the write. Once the per-delegate-key budget (60 weighted requests/minute) is spent the relayer answers 429 with a cooldown, and nothing backed off: the fact was never written and the agent saw only an opaque tool error. A short cooldown is now absorbed; a long one is reported with the wait named, stating plainly that the fact was NOT saved and pointing at the cheaper shape — one `memwal_remember_bulk` rather than N single calls, one `memwal_remember_status(job_ids)` rather than N status calls. Only rejections that provably never reached the handler retry, so `/api/remember/bulk`, which carries no idempotency key, cannot be duplicated by a retry. +- `memwal_remember` sends a content-derived idempotency key, so the retry its own timeout message invites really does attach to the job already in flight instead of storing a second paid copy. The key is computed by the tool rather than relied on from the SDK, whose published build mints a random UUID per client instance. +- `memwal_remember_status` accepts `job_ids` to settle a whole batch in one call, and reports a mixed batch honestly — a still-uploading row no longer renders the poll timeout as `error=`, which read as a failed write. Settling in one request also matters against the rate limit: 20 ids cost one request, not twenty. +- The bridge's cold-start tool list no longer disagrees with the sidecar's. `memwal_remember_status` advertised only `job_id`, required, under `additionalProperties: false`, so the batch call the tools themselves instruct was rejected until `tools/list_changed` arrived; the `waitMs` ceiling advertised 60000 after the sidecar lowered it to 45000, which came back as an MCP validation error; and `memwal_remember_bulk` still carried its pre-queue description. Tests now pin the parts an agent acts on. +- Persist the delegate keypair before sign-in hands the URL to the browser, and reclaim it on the next start. The browser's onchain `add_delegate_key` costs gas and is irreversible, and it happens before the callback that saves the private half, so a client that died in that window destroyed the only copy of a key the user had already paid for and left an orphaned registration nobody could use. (#793) Signing out discards the pending record along with the credentials, a second sign-in against the same relayer reuses the stranded key rather than minting over it, and a sign-in that cannot write the record fails instead of publishing a URL it cannot back. Reclaiming works on Mainnet; Testnet requires an account-id hint the recovering client does not have. + +### Changed + +- `memwal_remember` / `memwal_remember_bulk` return at accept by default (`MEMWAL_MCP_REMEMBER_WAIT_MS=0`, ~1s, `job_id`). The Walrus write continues in the background; do not treat that reply as stored. Set `MEMWAL_MCP_REMEMBER_WAIT_MS=90000` to restore wait-for-`blob_id` (90s ceiling). Do not use a value between 0 and the real completion time — that pays the wait and still returns pending. The MCP TypeScript SDK's default tools/call timeout is 60s, so a 90s wait loses on hosts that do not raise it. + ## 0.0.13 ### Fixed diff --git a/packages/mcp/README.md b/packages/mcp/README.md index 4a21e8584..d31a26026 100644 --- a/packages/mcp/README.md +++ b/packages/mcp/README.md @@ -23,6 +23,86 @@ Add Walrus Memory MCP to your MCP client config: } ``` +## How the plugin launches the server + +The [MemWal plugin](https://memory.walrus.xyz/mcp/claude-code) does **not** use the +`npx` form above. `npx` resolves a package *name* against the directory the MCP +client was started in — your project — so a project that contains an installed +`@mysten-incubation/memwal-mcp` claiming the pinned version would be run instead of +the published one. Pinning the version in the `npx` command does not prevent that: +the planted package simply claims the pinned version. + +Instead, every plugin launch config runs the plugin's launcher: + +```json +{ + "mcpServers": { + "memwal": { + "command": "node", + "args": ["${CLAUDE_PLUGIN_ROOT}/scripts/launch_mcp.mjs"] + } + } +} +``` + +The launcher installs the pinned version once into a directory it owns and then +runs that absolute entry point with the current `node` binary: + +``` +~/.memwal/runtime/memwal-mcp@/node_modules/@mysten-incubation/memwal-mcp/dist/bin/memwal-mcp.js +``` + +It never consults your project's `node_modules` or a `PATH`-relative bin shim, and +it fails rather than falling back to the package name if the pinned version cannot +be installed. Everything after the script path is forwarded to the server +unchanged, so flags such as `--namespace work` or `--relayer ` work exactly as +they do above. + +The runtime directory is trusted because of what it is, not how it is spelled: + +- It must be outside the project the client started in. An absolute path is not + enough on its own — a client expands `${workspaceFolder}` to an absolute path + inside the repository, and a repository can commit a whole fake install there. +- It must be a real directory (not a symlink), owned by you, and not group- or + world-writable. The launcher creates it with mode `0700` and refuses to run code + out of it otherwise. If you see a refusal, `chmod 700 ~/.memwal ~/.memwal/runtime`. +- `MEMWAL_MCP_RUNTIME_DIR` moves it, subject to exactly the same rules. + +The one-off install is run with `--ignore-scripts` (so no `preinstall` or +`postinstall` from the package or its dependencies executes), against an explicitly +pinned public registry, with the `npm_config_*` and `NODE_OPTIONS` environment +scrubbed for that spawn — npm ranks environment variables above every `.npmrc`, so +a client `env` block would otherwise choose the registry. npm itself is run as +`node ` resolved from the running node binary where that layout exists. + +What that does **not** give you is an integrity check of the package contents: it +establishes where the code came from and that nobody else can write it, not that +the registry served the bytes a reviewer read. If you install from a private +registry, pre-populate the directory yourself (below) — the launcher then finds the +install and never runs npm at all. + +If you configure MemWal without the plugin and want the same property, install the +version you intend to run into a directory outside any project and point your +client at its absolute path: + +```sh +mkdir -p ~/.memwal/runtime/memwal-mcp@0.0.14 +chmod 700 ~/.memwal ~/.memwal/runtime +npm install --ignore-scripts --prefix ~/.memwal/runtime/memwal-mcp@0.0.14 \ + @mysten-incubation/memwal-mcp@0.0.14 +``` + +```json +{ + "mcpServers": { + "memwal": { + "command": "node", + "args": ["/absolute/path/to/home/.memwal/runtime/memwal-mcp@0.0.14/node_modules/@mysten-incubation/memwal-mcp/dist/bin/memwal-mcp.js"] + } + } +} +``` + ## Login Run the login flow manually: @@ -38,6 +118,7 @@ The command opens your browser, asks you to connect your Sui wallet, and saves c ```sh memwal-mcp memwal-mcp login +memwal-mcp auto-save on|off memwal-mcp --logout memwal-mcp --help ``` @@ -52,9 +133,111 @@ Use CLI flags or environment variables to override the default Walrus Memory end | `--web-url ` | `MEMWAL_WEB_URL` | Override the web app URL used during login. | | `--label ` | `MEMWAL_CLIENT_LABEL` | Friendly delegate-key label shown in Walrus Memory. | | `--namespace ` (alias `--ns`) | `MEMWAL_NAMESPACE` | Default memory namespace applied when the agent omits one. | +| `auto-save on\|off` | `MEMWAL_AUTO_SAVE` | Whether the agent saves durable facts unprompted. On once you agree at login; see [Automatic Memory](#automatic-memory). | Enable verbose stderr logging with `MEMWAL_MCP_DEBUG=1`. +Set `MEMWAL_MCP_TRANSPORT=http` to dial the relayer's Streamable HTTP endpoint instead of the default SSE pair. Opt-in: reconnect replay is not transport-aware yet, so a write interrupted mid-send may be retried and duplicated. + +## Automatic Memory + +MemWal saves durable facts — preferences, decisions, constraints, recurring +workflows — as you state them, without asking each time. That is what it is for, +so it is on. It asks you once first. + +The question comes up in your terminal the first time you run `login`: + +``` +MemWal can save things about you automatically. + +What that means: when you state a preference, a decision, or a setting in +chat — "I prefer pnpm", "we deploy from dev", "the relayer is at X" — +MemWal writes it to your memory without asking each time, so it is there +in your next session and in every other client you use. + +Before you choose: + + - Saved memories are permanent. They go to Walrus, which is immutable + storage. You can stop saving new ones at any time, but you cannot + delete one that is already saved. + - They are encrypted to your account. Only your delegate key reads them. + - MemWal strips obvious credentials — API keys, tokens, passwords, + private keys — before saving. Treat that as a safety net, not a + guarantee: do not paste secrets into a session with this on. + + [1] Save automatically recommended, this is what MemWal is for + [2] Only save when I ask nothing is saved unless you say "remember this" + +Your choice [1/2]: +``` + +**Saved memories are permanent.** Walrus is immutable storage: you can stop +saving new ones at any time, but you cannot delete one already saved — which is +why you are asked before it starts rather than after. + +Until you answer, nothing is saved unprompted on a new install; "remember this" +and recall work throughout. If you were using MemWal before this question +existed, it keeps saving as it always has and asks you at your next login. + +The question is only ever asked in a terminal. It is never an MCP tool and never +something the assistant can answer for you. Change it any time: + +```sh +npx -y @mysten-incubation/memwal-mcp auto-save on +npx -y @mysten-incubation/memwal-mcp auto-save off +npx -y @mysten-incubation/memwal-mcp auto-save # report the current setting +``` + +The answer is stored as `{"autoSave": true}` in `settings.json` next to your +credentials file, so it follows the same project-local-beats-global resolution. +To pin one MCP client instead, set the environment variable — it overrides the +file and skips the question: + +```json +{ + "mcpServers": { + "memwal": { + "command": "npx", + "args": ["-y", "@mysten-incubation/memwal-mcp"], + "env": { "MEMWAL_AUTO_SAVE": "1" } + } + } +} +``` + +Recall is never gated, and neither is an explicit "remember this" — the setting +only decides whether the agent saves things you did not ask it to save. Saying +no costs nothing and is not asked about again. + +### What is never saved + +Walrus storage is append-only and encrypted: a memory that lands **cannot be +edited or deleted**. So credentials are excluded whichever way you answer, by the same +rules stated in the server instructions, the tool descriptions and the plugin +hooks — and enforced by a check that runs before any text is sent: + +- passwords, API keys, access and refresh tokens, private keys, seed and + recovery phrases, authorization headers, session cookies, and connection + strings or URLs with an embedded `user:password`; +- anything you say not to save ("don't save this", "off the record"); +- pasted third-party content — a fenced block, a quoted passage — which is not + a fact about you. + +When a message mixes a preference with a credential, the **preference is kept** +and only the credential is removed: "I prefer dark mode, db is +`postgres://admin:hunter2@db.internal/app`" is stored with the password gone and +the host intact. The agent is told which kinds were removed; the secret itself +is never stored, logged, or echoed back. + +Detection targets specific credential shapes rather than "looks random", so +identifiers you *do* want remembered — blob ids, Sui object ids, commit SHAs, +digests — pass through untouched. Where a secret is indistinguishable from an +identifier, the **label** decides: pasting your `credentials.json` has its +`delegatePrivateKey` removed, while the same 64 hex characters with nothing +calling them a key are stored as the digest they look like. The trade-off is +that a secret in no recognisable shape, and with no label near it, can still +slip past the check — which is why the model-facing rules exist alongside it. + ## Default Namespace By default the MCP tool schemas expose an optional `namespace` argument and the @@ -146,12 +329,66 @@ You can also pass explicit URLs: ## Credential Storage -Credentials are stored locally in `~/.memwal/credentials.json`. To remove them: +Credentials are stored locally in `~/.memwal/credentials.json`, and the +automatic-memory setting in `settings.json` beside it. To remove the +credentials: ```sh npx -y @mysten-incubation/memwal-mcp --logout ``` +A sign-in that is still in flight also writes a `login-pending.json` +write-ahead record, beside `credentials.json` in `~/.memwal`. It holds the +delegate keypair minted for that sign-in, written before the browser can +register the public half on-chain so an interrupted login can be reclaimed +instead of paid for a second time. Same owner-only mode `0600` as +`credentials.json`. + +It is removed once the sign-in completes, on `--logout`, and on a successful +recovery at the next start. A record older than 24 hours is discarded rather +than reused. + +### Per-project credentials + +A project can keep its own `.memwal/credentials.json` so memory written from it +goes to a separate account. That file lives inside the repository, where anyone +who can commit to it — or who can get you to open a clone — could otherwise +choose the account and relayer your memories go to. So it is **ignored until you +approve it**, once per machine: + +```sh +cd path/to/project +npx -y @mysten-incubation/memwal-mcp approve-project +``` + +Until then the global credentials are used, and a line on stderr names the file +that was skipped and the destination it wanted. An approval covers one exact +project path, account, delegate key and relayer: if any of those change, it has +to be approved again. The record is kept in `~/.memwal/project-approvals.json`, +outside the repository, so a repository cannot carry its own approval. +`revoke-project` withdraws it. + +Approving also picks the file that is **written**: a later sign-in from that +project saves a delegate private key into `.memwal/credentials.json` in plain +text, inside the repository. `approve-project` says so, and so does the sign-in +warning. Add `.memwal/` to your `.gitignore`. The write-ahead record is the one +thing that stays out: a project sign-in keeps it in `~/.memwal/login-pending/`, +one file per approved project, so no key material lands in the checkout and a +sign-in is still reclaimable only by the project that started it. + +`MEMWAL_CREDS_DIR` points the credentials, the approval record and the +write-ahead record at a directory of your choosing and overrides project +resolution entirely, with no approval. Because it skips the gate, it must be an +**absolute path outside the current project**: a relative value would resolve +against the working directory and put the approval record inside the +repository, and an MCP client passes on the `env` block it reads from +`.cursor/mcp.json` / `.vscode/mcp.json` / `.claude/settings.json` in the +checkout — where `${workspaceFolder}` is expanded, so an in-project absolute +path is not proof you chose it. Anything else is refused with an error naming +the value, rather than quietly ignored. An empty value means unset. + +`memwal_health` reports the destination in use as `account=… relayer=…`. + ## License Apache-2.0 diff --git a/packages/mcp/TESTING.md b/packages/mcp/TESTING.md index d4f4e9dda..d05377295 100644 --- a/packages/mcp/TESTING.md +++ b/packages/mcp/TESTING.md @@ -4,9 +4,11 @@ Step-by-step plan to test the auto-memory work (agentic tools + bulk + health + > **Local note:** the standalone configs below point at your **local build** via > `node /Users/uydev/code/MemWal/packages/mcp/dist/bin/memwal-mcp.js --local` -> (tests your local code, not the published npx package). The **shipped plugin -> `.mcp.json` is prod** (`npx -y @mysten-incubation/memwal-mcp`) — no `--local` in -> the committed file. It still reaches your **local relayer** because your saved +> (tests your local code, not the published package). The **shipped plugin +> `.mcp.json` is prod** — it runs `node "${CLAUDE_PLUGIN_ROOT}/scripts/launch_mcp.mjs"`, +> which installs the pinned version into `~/.memwal/runtime` and launches that +> absolute path (never `npx`, never your project's `node_modules`) — no `--local` +> in the committed file. It still reaches your **local relayer** because your saved > creds (`~/.memwal/credentials.json`) point there. To force a local target by hand > (e.g. fresh/prod creds), `export MEMWAL_SERVER_URL=http://127.0.0.1:8000` in the > shell before launching the client — the prod bin reads it (`index.ts:116-117`), @@ -23,6 +25,16 @@ Step-by-step plan to test the auto-memory work (agentic tools + bulk + health + ``` - [ ] Credentials present and on testnet: `~/.memwal/credentials.json` (relayerUrl = `http://127.0.0.1:8000`). - [ ] MCP package built: `ls packages/mcp/dist/bin/memwal-mcp.js`. +- [ ] **Automatic saving answered** — since WALM-642 `login` asks once in the + terminal, and a fresh install saves nothing unprompted until it is + answered. Every "agent saves on its own" step below depends on the answer + being yes: + ```bash + node packages/mcp/dist/bin/memwal-mcp.js auto-save # → reports the state + node packages/mcp/dist/bin/memwal-mcp.js auto-save on # → "Automatic memory is ON" + ``` + Clear `~/.memwal/settings.json` first if you want to see the consent + prompt itself (§1.6). **How to tell MemWal vs the editor's built-in memory** (use everywhere below): - ✅ **MemWal** → the step shows `Called memwal…` with a **`blob_id`** + `namespace`. @@ -77,6 +89,53 @@ Notes: --- +## 1.6 Consent and secret filtering (WALM-642) + +The redactor and the opt-in resolver have automated coverage +(`services/server/scripts/mcp/__tests__/{secret-redaction,write-path-redaction}.test.ts`, +`packages/mcp/test/{memory-policy,auto-save-optin}.test.mjs`). What is NOT +automatable is the half the ticket calls "model behavior": whether a real model, +reading the injected rules, actually declines to save a secret it was never +programmatically stopped from sending. Run these by hand in a real client. + +**Consent at login** — the prompt is TTY-only, so this part is a terminal, not a +client. + +| # | Action | Expect | OK? | +|---|---|---|---| +| 1 | `rm ~/.memwal/settings.json`, then `memwal-mcp login` in a terminal | The consent prompt renders: consequence first, permanence as the first bullet, redaction described as a safety net. Both options readable. | [ ] | +| 2 | Type `banana` at the prompt | Re-asks. Does **not** assume [1]. | [ ] | +| 3 | Press Enter | Takes [1]; `settings.json` shows `"autoSave": true` | [ ] | +| 4 | Run `memwal-mcp login` again | Does **not** re-ask | [ ] | +| 5 | `rm settings.json`, `login`, answer `2` | `"autoSave": false`; says once how to change it | [ ] | +| 6 | `login` again after answering `2` | Does **not** re-ask, does not nag | [ ] | +| 7 | `rm settings.json`, then start the server from an MCP client (no TTY) | No prompt, no hang. One stderr line saying auto-save is waiting on an answer. | [ ] | +| 8 | Ask the assistant in chat to "turn automatic memory on for me" | It points you at the terminal command; there is no tool it can call to answer for you | [ ] | + +**Behaviour either way** + +| # | Action / prompt | Expect | OK? | +|---|---|---|---| +| 9 | With `auto-save off`: `I prefer pnpm and TypeScript strict mode.` | Agent does **not** call `memwal_remember` on its own | [ ] | +| 10 | Same session: `remember that I prefer pnpm` | Agent **does** call `memwal_remember` — an explicit ask is never gated | [ ] | +| 11 | Same session: `what do you remember about my preferences?` | Agent calls `memwal_recall` — recall is never gated | [ ] | +| 12 | `auto-save on`, restart, repeat #9 | Agent calls `memwal_remember` unprompted again | [ ] | + +**Secret filtering (with `auto-save on`)** + +| # | Action / prompt | Expect | OK? | +|---|---|---|---| +| 13 | `I prefer dark mode, and the staging db is postgres://admin:hunter2@db.internal:5432/app` | A memory IS saved, and it contains the preference and `db.internal:5432/app` but **not** `hunter2`. The tool reply names `url-credentials`. Confirm with a recall in a new chat. | [ ] | +| 14 | Same, but phrased so the model saves it as several facts | `memwal_remember_bulk` — same result per entry; a bare-secret entry is reported as `NOT SAVED (1)` with its position | [ ] | +| 15 | Paste a transcript containing a preference and `ghp_…`, ask to analyse it | `memwal_analyze` — the extracted facts never contain the token | [ ] | +| 16 | `My bank PIN is 4821 — don't save this.` | Nothing is saved; the agent says so | [ ] | +| 17 | Paste a fenced log/code block and say "save this" | Not saved as a fact about you; the agent asks you to restate it | [ ] | +| 18 | **Model behavior:** state a secret in a shape the redactor does not match (e.g. `my door code is seven four nine two`) and see whether the model saves it | The rules say not to; a save here is a model failure, not a code failure — record it, it is the residual risk the redactor cannot close | [ ] | +| 19 | Paste a `credentials.json` line containing `"delegatePrivateKey": "<64 hex>"` | Saved without the key; an unlabelled 64-hex digest in the same sentence survives | [ ] | +| 20 | Check `~/.memwal/settings.json` and the relayer logs after #13-#19 | No secret appears in either — the redactor never logs what it removed | [ ] | + +--- + ## 2. Per-editor tests ### 2A. Claude Code — Plugin (hooks + "prefer MemWal" steer) @@ -195,7 +254,7 @@ Notes: ## 4. Before the PR -- [x] `packages/mcp/plugin/.mcp.json` ships the prod default (`npx -y @mysten-incubation/memwal-mcp`) — no `--local` in the committed file +- [x] `packages/mcp/plugin/.mcp.json` ships the prod default (`node "${CLAUDE_PLUGIN_ROOT}/scripts/launch_mcp.mjs"`, which installs the pin under `~/.memwal/runtime` and runs it by absolute path) — no `--local` in the committed file - [ ] Re-register `memwal-local` if you removed it for the plugin test - [ ] Remove the temporary local `memwal` entries from Claude Desktop / Cursor / Codex / OpenCode configs (or keep for ongoing local dev) diff --git a/packages/mcp/package.json b/packages/mcp/package.json index 44738b58d..69ed21ca3 100644 --- a/packages/mcp/package.json +++ b/packages/mcp/package.json @@ -1,7 +1,7 @@ { "name": "@mysten-incubation/memwal-mcp", - "version": "0.0.13", - "description": "Walrus Memory MCP client — single-binary stdio MCP server that bridges Cursor / Claude Desktop / Antigravity / Claude Code to the Walrus Memory relayer. Handles browser-based wallet login on first run.", + "version": "0.0.14", + "description": "Walrus Memory MCP client \u2014 single-binary stdio MCP server that bridges Cursor / Claude Desktop / Antigravity / Claude Code to the Walrus Memory relayer. Handles browser-based wallet login on first run.", "type": "module", "engines": { "node": ">=20.0.0" diff --git a/packages/mcp/plugin/.claude-plugin/plugin.json b/packages/mcp/plugin/.claude-plugin/plugin.json index 66925ea5d..f291f6394 100644 --- a/packages/mcp/plugin/.claude-plugin/plugin.json +++ b/packages/mcp/plugin/.claude-plugin/plugin.json @@ -1,12 +1,18 @@ { - "name": "memwal", - "version": "0.0.13", - "description": "Automatic Walrus Memory for Claude Code — proactive recall and durable-fact saving via the MemWal MCP + lifecycle hooks.", - "author": { - "name": "Mysten Labs" - }, - "homepage": "https://memory.walrus.xyz", - "repository": "https://github.com/MystenLabs/MemWal", - "license": "Apache-2.0", - "keywords": ["memory", "mcp", "walrus", "sui", "semantic-search"] + "name": "memwal", + "version": "0.0.14", + "description": "Automatic Walrus Memory for Claude Code \u2014 proactive recall and durable-fact saving via the MemWal MCP + lifecycle hooks.", + "author": { + "name": "Mysten Labs" + }, + "homepage": "https://memory.walrus.xyz", + "repository": "https://github.com/MystenLabs/MemWal", + "license": "Apache-2.0", + "keywords": [ + "memory", + "mcp", + "walrus", + "sui", + "semantic-search" + ] } diff --git a/packages/mcp/plugin/.codex-mcp.json b/packages/mcp/plugin/.codex-mcp.json index 345b9c487..5056557f1 100644 --- a/packages/mcp/plugin/.codex-mcp.json +++ b/packages/mcp/plugin/.codex-mcp.json @@ -1,8 +1,8 @@ { "mcpServers": { "memwal": { - "command": "npx", - "args": ["-y", "@mysten-incubation/memwal-mcp@0.0.13"] + "command": "node", + "args": ["${PLUGIN_ROOT}/scripts/launch_mcp.mjs"] } } } diff --git a/packages/mcp/plugin/.codex-plugin/plugin.json b/packages/mcp/plugin/.codex-plugin/plugin.json index bdbfa5d91..3b70638ca 100644 --- a/packages/mcp/plugin/.codex-plugin/plugin.json +++ b/packages/mcp/plugin/.codex-plugin/plugin.json @@ -1,29 +1,38 @@ { - "name": "memwal", - "version": "0.0.13", - "description": "Persistent Walrus Memory for Codex. Remembers decisions, preferences, and project context across sessions.", - "author": { - "name": "Mysten Labs", - "url": "https://memory.walrus.xyz" - }, - "homepage": "https://memory.walrus.xyz", - "repository": "https://github.com/MystenLabs/MemWal", - "license": "Apache-2.0", - "keywords": ["memory", "personalization", "mcp", "walrus", "semantic-search"], - "mcpServers": "./.codex-mcp.json", - "hooks": "./hooks/codex-hooks.json", - "interface": { - "displayName": "MemWal", - "shortDescription": "Portable, encrypted memory layer for AI coding workflows", - "longDescription": "MemWal adds long-term, user-owned memory to Codex. Store decisions, preferences, and session context; memories are encrypted with SEAL and stored on Walrus, and retrieved via semantic search so Codex always has the right context.", - "developerName": "Mysten Labs", - "category": "Productivity", - "capabilities": ["Read", "Write"], - "websiteURL": "https://memory.walrus.xyz", - "defaultPrompt": [ - "Search my memories for recent project decisions", - "Remember that I prefer pnpm and TypeScript strict mode", - "What do you know about my coding preferences?" - ] - } + "name": "memwal", + "version": "0.0.14", + "description": "Persistent Walrus Memory for Codex. Remembers decisions, preferences, and project context across sessions.", + "author": { + "name": "Mysten Labs", + "url": "https://memory.walrus.xyz" + }, + "homepage": "https://memory.walrus.xyz", + "repository": "https://github.com/MystenLabs/MemWal", + "license": "Apache-2.0", + "keywords": [ + "memory", + "personalization", + "mcp", + "walrus", + "semantic-search" + ], + "mcpServers": "./.codex-mcp.json", + "hooks": "./hooks/codex-hooks.json", + "interface": { + "displayName": "MemWal", + "shortDescription": "Portable, encrypted memory layer for AI coding workflows", + "longDescription": "MemWal adds long-term, user-owned memory to Codex. Store decisions, preferences, and session context; memories are encrypted with SEAL and stored on Walrus, and retrieved via semantic search so Codex always has the right context.", + "developerName": "Mysten Labs", + "category": "Productivity", + "capabilities": [ + "Read", + "Write" + ], + "websiteURL": "https://memory.walrus.xyz", + "defaultPrompt": [ + "Search my memories for recent project decisions", + "Remember that I prefer pnpm and TypeScript strict mode", + "What do you know about my coding preferences?" + ] + } } diff --git a/packages/mcp/plugin/.cursor-mcp.json b/packages/mcp/plugin/.cursor-mcp.json index 345b9c487..e3b5dc20a 100644 --- a/packages/mcp/plugin/.cursor-mcp.json +++ b/packages/mcp/plugin/.cursor-mcp.json @@ -1,8 +1,8 @@ { "mcpServers": { "memwal": { - "command": "npx", - "args": ["-y", "@mysten-incubation/memwal-mcp@0.0.13"] + "command": "node", + "args": ["${CURSOR_PLUGIN_ROOT}/scripts/launch_mcp.mjs"] } } } diff --git a/packages/mcp/plugin/.cursor-plugin/plugin.json b/packages/mcp/plugin/.cursor-plugin/plugin.json index 68bfe76ff..d7eb97856 100644 --- a/packages/mcp/plugin/.cursor-plugin/plugin.json +++ b/packages/mcp/plugin/.cursor-plugin/plugin.json @@ -1,12 +1,19 @@ { - "name": "memwal", - "version": "0.0.13", - "description": "Automatic Walrus Memory for Cursor — proactive recall and durable-fact saving via the MemWal MCP + lifecycle hooks.", - "author": { "name": "Mysten Labs" }, - "homepage": "https://memory.walrus.xyz", - "repository": "https://github.com/MystenLabs/MemWal", - "license": "Apache-2.0", - "keywords": ["memory", "mcp", "walrus", "semantic-search"], - "hooks": "./hooks/cursor-hooks.json", - "mcpServers": ".cursor-mcp.json" + "name": "memwal", + "version": "0.0.14", + "description": "Automatic Walrus Memory for Cursor \u2014 proactive recall and durable-fact saving via the MemWal MCP + lifecycle hooks.", + "author": { + "name": "Mysten Labs" + }, + "homepage": "https://memory.walrus.xyz", + "repository": "https://github.com/MystenLabs/MemWal", + "license": "Apache-2.0", + "keywords": [ + "memory", + "mcp", + "walrus", + "semantic-search" + ], + "hooks": "./hooks/cursor-hooks.json", + "mcpServers": ".cursor-mcp.json" } diff --git a/packages/mcp/plugin/.mcp.json b/packages/mcp/plugin/.mcp.json index 345b9c487..baed20fd6 100644 --- a/packages/mcp/plugin/.mcp.json +++ b/packages/mcp/plugin/.mcp.json @@ -1,8 +1,8 @@ { "mcpServers": { "memwal": { - "command": "npx", - "args": ["-y", "@mysten-incubation/memwal-mcp@0.0.13"] + "command": "node", + "args": ["${CLAUDE_PLUGIN_ROOT}/scripts/launch_mcp.mjs"] } } } diff --git a/packages/mcp/plugin/plugin.json b/packages/mcp/plugin/plugin.json index 17b1993e0..708945802 100644 --- a/packages/mcp/plugin/plugin.json +++ b/packages/mcp/plugin/plugin.json @@ -1,12 +1,21 @@ { - "id": "memwal", - "name": "memwal", - "version": "0.0.13", - "description": "Automatic Walrus Memory for Antigravity — proactive recall and durable-fact saving via the MemWal MCP + lifecycle hooks.", - "author": { "name": "Mysten Labs" }, - "homepage": "https://memory.walrus.xyz", - "repository": "https://github.com/MystenLabs/MemWal", - "license": "Apache-2.0", - "keywords": ["memory", "mcp", "walrus", "sui", "semantic-search"], - "contextFileName": "AGENTS.md" + "id": "memwal", + "name": "memwal", + "version": "0.0.14", + "mcpPackageVersion": "0.0.14-dev.2", + "description": "Automatic Walrus Memory for Antigravity \u2014 proactive recall and durable-fact saving via the MemWal MCP + lifecycle hooks.", + "author": { + "name": "Mysten Labs" + }, + "homepage": "https://memory.walrus.xyz", + "repository": "https://github.com/MystenLabs/MemWal", + "license": "Apache-2.0", + "keywords": [ + "memory", + "mcp", + "walrus", + "sui", + "semantic-search" + ], + "contextFileName": "AGENTS.md" } diff --git a/packages/mcp/plugin/scripts/install_codex_hooks.mjs b/packages/mcp/plugin/scripts/install_codex_hooks.mjs index f4d2a71fe..0eab99ec0 100644 --- a/packages/mcp/plugin/scripts/install_codex_hooks.mjs +++ b/packages/mcp/plugin/scripts/install_codex_hooks.mjs @@ -12,8 +12,20 @@ * ${PLUGIN_ROOT} placeholder to this plugin's absolute path, and merges the * entries into ~/.codex/hooks.json. * - * Re-running is idempotent: entries this installer owns (identified by our - * hook script filenames) are removed before fresh entries are added. + * The template is parsed as JSON *before* the placeholder is substituted, and + * the path is POSIX-single-quoted on its way into a hook command, so a plugin + * directory containing $(...), backticks, quotes or backslashes cannot break + * out of either the JSON document or the generated shell command. + * + * Re-running is idempotent: the hooks this installer owns are removed before + * fresh ones are added. Ownership is decided per hook, by a marker this + * installer writes (or, for installs predating that marker, by an exact match + * against the commands it generates for this plugin directory) — never by the + * hook script's filename, which another tool may legitimately share. Hooks + * belonging to anyone else, including siblings inside the same group, and the + * group's own settings are left untouched. An existing [mcp_servers.memwal] + * block is migrated to the launcher when it still points at something else + * (WALM-640) instead of being reported as "already present". * * Usage: * node install_codex_hooks.mjs # install or update @@ -28,32 +40,36 @@ import { readFileSync, writeFileSync, existsSync, mkdirSync } from "node:fs"; import { homedir } from "node:os"; import { join, dirname } from "node:path"; import { fileURLToPath } from "node:url"; +import { substituteHookPlaceholder } from "./lib/hook-template.mjs"; + +import { planMcpRegistration } from "./lib/codex-config.mjs"; const SCRIPT_DIR = dirname(fileURLToPath(import.meta.url)); const PLUGIN_ROOT = dirname(SCRIPT_DIR); -function resolveMcpVersion() { - return JSON.parse(readFileSync(join(PLUGIN_ROOT, "plugin.json"), "utf8")).version; -} - const CODEX_DIR = join(homedir(), ".codex"); const HOOKS_FILE = join(CODEX_DIR, "hooks.json"); const CONFIG_FILE = join(CODEX_DIR, "config.toml"); const TEMPLATE_FILE = join(PLUGIN_ROOT, "hooks", "codex-hooks.json"); -// Entries are "ours" when a hook command references one of our scripts. -const OWNER_MARKERS = [ - "on_session_start.mjs", - "on_user_prompt.mjs", - "on_post_tool.mjs", -]; +const PLACEHOLDER = "${PLUGIN_ROOT}"; + +// Every hook this installer writes carries this marker, so a later run can +// recognise its own entries outright instead of guessing from a filename. +// `on_user_prompt.mjs` is a generic name: another tool's hook may well use it, +// and that hook is not ours to touch. +const MARKER_KEY = "_memwal"; +const MARKER_VALUE = "memwal-plugin-hooks"; function loadTemplate() { - const raw = readFileSync(TEMPLATE_FILE, "utf8").replaceAll( - "${PLUGIN_ROOT}", + // Parse first, substitute second. Substituting into the raw text would let + // a path containing a double quote or a backslash rewrite the JSON + // document, and would leave `$(...)` or backticks live in the hook command. + return substituteHookPlaceholder( + JSON.parse(readFileSync(TEMPLATE_FILE, "utf8")), + PLACEHOLDER, PLUGIN_ROOT ); - return JSON.parse(raw); } function loadExisting() { @@ -66,28 +82,87 @@ function loadExisting() { } } -function isOwned(entry) { - for (const hook of entry.hooks || []) { - const cmd = hook.command || ""; - if (OWNER_MARKERS.some((m) => cmd.includes(m))) return true; +/** + * The exact commands this installer writes for the current PLUGIN_ROOT, plus + * the ones it wrote before WALM-641 quoted the path. Installs made by an older + * build carry no marker, so they are still recognised — but only when the + * command matches ours character for character, which a hook belonging to + * another tool never will. + */ +function ownedCommands() { + const commands = new Set(); + if (!existsSync(TEMPLATE_FILE)) return commands; + let template; + try { + template = JSON.parse(readFileSync(TEMPLATE_FILE, "utf8")); + } catch { + return commands; } - return false; + for (const entries of Object.values(template.hooks || {})) { + for (const entry of entries || []) { + for (const hook of entry.hooks || []) { + if (typeof hook.command !== "string") continue; + // What this build writes, and what pre-WALM-641 builds wrote. + commands.add( + substituteHookPlaceholder(hook, PLACEHOLDER, PLUGIN_ROOT).command + ); + commands.add(hook.command.replaceAll(PLACEHOLDER, PLUGIN_ROOT)); + } + } + } + return commands; +} + +const OWNED_COMMANDS = ownedCommands(); + +/** A single hook — not the group around it — that this installer put there. */ +function isOwnedHook(hook) { + if (!hook || typeof hook !== "object") return false; + if (hook[MARKER_KEY] === MARKER_VALUE) return true; + return typeof hook.command === "string" && OWNED_COMMANDS.has(hook.command); } +/** + * Drop our own hooks and nothing else. + * + * A group may hold hooks from several tools. Removing the group because one of + * its hooks is ours takes the siblings with it (WALM-643), so the group is + * rebuilt with its settings intact and only our hooks filtered out. A group is + * dropped only once it has no hooks left, and an event only once it has no + * groups left. + */ function stripOwned(config) { const hooks = config.hooks || {}; for (const event of Object.keys(hooks)) { - hooks[event] = (hooks[event] || []).filter((e) => !isOwned(e)); - if (hooks[event].length === 0) delete hooks[event]; + const kept = []; + for (const entry of hooks[event] || []) { + if (!entry || !Array.isArray(entry.hooks)) { + kept.push(entry); + continue; + } + const keptHooks = entry.hooks.filter((hook) => !isOwnedHook(hook)); + if (keptHooks.length === entry.hooks.length) kept.push(entry); + else if (keptHooks.length > 0) kept.push({ ...entry, hooks: keptHooks }); + } + if (kept.length === 0) delete hooks[event]; + else hooks[event] = kept; } config.hooks = hooks; return config; } +/** Tag each hook so the next run recognises it without matching commands. */ +function markOwned(entries) { + return entries.map((entry) => ({ + ...entry, + hooks: (entry.hooks || []).map((hook) => ({ ...hook, [MARKER_KEY]: MARKER_VALUE })), + })); +} + function mergeTemplate(config, template) { config.hooks = config.hooks || {}; for (const [event, entries] of Object.entries(template.hooks || {})) { - config.hooks[event] = (config.hooks[event] || []).concat(entries); + config.hooks[event] = (config.hooks[event] || []).concat(markOwned(entries)); } return config; } @@ -97,18 +172,59 @@ function writeHooks(config) { writeFileSync(HOOKS_FILE, JSON.stringify(config, null, 2) + "\n"); } -/** Append [mcp_servers.memwal] to config.toml if it isn't registered yet. */ +/** + * Register — or migrate — [mcp_servers.memwal] in config.toml. + * + * Registers the plugin's launcher by absolute path rather than + * `npx @mysten-incubation/memwal-mcp@`. npx resolves that name against the + * directory Codex is started in, so a package installed in the user's project under + * the same name — claiming the pinned version — was run instead of ours (WALM-640). + * The launcher installs the pinned version under ~/.memwal/runtime and runs that + * absolute entry point, so no project directory takes part in the resolution. + * + * An existing block is REWRITTEN, not skipped. Everyone who ran this installer before + * WALM-640 has the `npx` form on disk, and a "already present" message on re-run would + * leave exactly the resolution this fix exists to remove in place for exactly the + * people who already installed. Unrelated keys in the block (`env`, timeouts, …) are + * preserved — see lib/codex-config.mjs. + */ function ensureMcpRegistered() { mkdirSync(CODEX_DIR, { recursive: true }); - let content = existsSync(CONFIG_FILE) ? readFileSync(CONFIG_FILE, "utf8") : ""; - if (content.includes("[mcp_servers.memwal]")) return false; - const spec = `@mysten-incubation/memwal-mcp@${resolveMcpVersion()}`; - const block = - "\n[mcp_servers.memwal]\n" + - 'command = "npx"\n' + - `args = ["-y", "${spec}"]\n`; - writeFileSync(CONFIG_FILE, (content.trimEnd() + "\n" + block).trimStart()); - return true; + const content = existsSync(CONFIG_FILE) ? readFileSync(CONFIG_FILE, "utf8") : ""; + const launcher = join(SCRIPT_DIR, "launch_mcp.mjs"); + const plan = planMcpRegistration(content, launcher); + if (plan.action !== "unchanged") writeFileSync(CONFIG_FILE, plan.content); + return plan; +} + +function reportMcpRegistration(plan) { + if (plan.action === "added") { + console.log(`Registered [mcp_servers.memwal] in ${CONFIG_FILE}`); + return; + } + if (plan.action === "unchanged") { + console.log(`[mcp_servers.memwal] in ${CONFIG_FILE} already runs the launcher`); + return; + } + console.log(`Migrated [mcp_servers.memwal] in ${CONFIG_FILE}:`); + console.log(` was: command = ${plan.previous.command ?? "(absent)"}`); + console.log(` args = ${plan.previous.args ?? "(absent)"}`); + console.log(` now: command = "node"`); + console.log(` args = [${JSON.stringify(launcherPath())}]`); + if (plan.previous.command?.includes("npx")) { + console.log( + " (the old command resolved the package name against the directory Codex " + + "was started in — WALM-640)" + ); + } + if (plan.preserved.length > 0) { + console.log(` kept your other keys: ${plan.preserved.join(", ")}`); + } + console.log(" Restart Codex for the change to take effect."); +} + +function launcherPath() { + return join(SCRIPT_DIR, "launch_mcp.mjs"); } function featureFlagEnabled() { @@ -154,16 +270,12 @@ function main() { config = mergeTemplate(config, template); writeHooks(config); - const mcpAdded = ensureMcpRegistered(); + const mcpPlan = ensureMcpRegistered(); console.log(`Installed MemWal hooks into ${HOOKS_FILE}`); console.log(`Plugin path: ${PLUGIN_ROOT}`); console.log("Events: SessionStart, UserPromptSubmit, PostToolUse"); - console.log( - mcpAdded - ? `Registered [mcp_servers.memwal] in ${CONFIG_FILE}` - : `[mcp_servers.memwal] already present in ${CONFIG_FILE}` - ); + reportMcpRegistration(mcpPlan); if (!featureFlagEnabled()) printFeatureFlagHint(); return 0; diff --git a/packages/mcp/plugin/scripts/launch_mcp.mjs b/packages/mcp/plugin/scripts/launch_mcp.mjs new file mode 100644 index 000000000..db8012dec --- /dev/null +++ b/packages/mcp/plugin/scripts/launch_mcp.mjs @@ -0,0 +1,80 @@ +#!/usr/bin/env node +/** + * The plugin's MCP launcher (WALM-640). + * + * Replaces `npx -y @mysten-incubation/memwal-mcp@` in every plugin manifest. + * npx resolves that name against the directory the MCP client started in — the + * user's project — so a package installed there under the same name, claiming the + * pinned version, was run instead of ours. This launcher never resolves a name: + * it makes sure the pinned version is installed under ~/.memwal/runtime and runs + * that absolute entry point with the current node binary. + * + * The runtime directory is trusted because of what it is, not because of how it is + * spelled: it must sit outside the project tree (an absolute + * `${workspaceFolder}/.memwal-runtime` is still the project), be a real directory + * owned by the current user, and not be group- or world-writable. See + * lib/mcp-launch.mjs for the checks and for what the install step does and does not + * guarantee. + * + * Everything after the script path is forwarded to the server untouched, so the + * manifests keep working with flags such as `--dev`, `--namespace work` or + * `--relayer `. The environment is inherited as-is, and the working directory + * is left alone so project-local credentials (`.memwal/credentials.json` at or + * above the cwd) still resolve the way they do today. + * + * Usage: + * node launch_mcp.mjs [server args...] # start the MCP stdio server + * node launch_mcp.mjs --print-entry # print the resolved path and exit + */ +import { spawn } from "node:child_process"; + +import { ensureTrustedEntry, pinnedVersion } from "./lib/mcp-launch.mjs"; + +const argv = process.argv.slice(2); +const printOnly = argv[0] === "--print-entry"; + +let entry; +try { + entry = ensureTrustedEntry(); +} catch (err) { + process.stderr.write( + `[memwal-mcp] launcher: could not prepare the pinned server ` + + `(${pinnedVersionSafe()}): ${err?.message ?? String(err)}\n` + + `[memwal-mcp] launcher: refusing to fall back to a package resolved from ` + + `the current directory.\n`, + ); + process.exit(1); +} + +if (printOnly) { + process.stdout.write(entry + "\n"); + process.exit(0); +} + +const child = spawn(process.execPath, [entry, ...argv], { + stdio: "inherit", + env: process.env, +}); + +for (const signal of ["SIGINT", "SIGTERM", "SIGHUP"]) { + process.on(signal, () => { + if (!child.killed) child.kill(signal); + }); +} + +child.on("error", (err) => { + process.stderr.write(`[memwal-mcp] launcher: failed to start ${entry}: ${err.message}\n`); + process.exitCode = 1; +}); + +child.on("exit", (code, signal) => { + process.exitCode = signal ? 1 : (code ?? 0); +}); + +function pinnedVersionSafe() { + try { + return pinnedVersion(); + } catch { + return "unknown version"; + } +} diff --git a/packages/mcp/plugin/scripts/lib/auto-save.mjs b/packages/mcp/plugin/scripts/lib/auto-save.mjs new file mode 100644 index 000000000..b9def610a --- /dev/null +++ b/packages/mcp/plugin/scripts/lib/auto-save.mjs @@ -0,0 +1,172 @@ +/** + * Automatic-save opt-in, hook-side (WALM-642). + * + * This file used to be a hand-written MIRROR of `packages/mcp/src/auto-save.ts` + * — the same walk up the directory tree, the same "nearest project-local + * `.memwal/credentials.json` wins", the same settings lookup beside whatever + * that resolved to. Keeping two implementations of one security decision in + * step is not a thing anyone manages to do, and they drifted: the TypeScript + * side grew an approval gate (WALM-639) and this side did not, so a repository + * containing nothing but `.memwal/credentials.json` — contents never parsed, + * only `existsSync` — read as proof of a long-standing install and switched + * automatic memory ON for anyone who opened it. A committed `settings.json` + * next to it could pin `autoSave: true` outright and silence the + * consent-pending warning while it did. + * + * So the mirror is gone. The MCP server resolves the state — approvals, + * project scoping, the consent answer, all of it — and publishes the ANSWER to + * `auto-save-state.json` in the trusted state dir (`MEMWAL_CREDS_DIR`, else + * `~/.memwal`). This file reads that file and nothing else. + * + * What that buys: + * - No resolution here at all, so there is nothing left to drift. + * - Nothing under `process.cwd()` is ever read, so a checkout cannot + * influence hook behaviour — not through a credentials file, not through a + * settings file, not through anything it can add later. + * - Unreadable, missing, malformed, or written by a newer version reads as + * "state unknown", which is automatic memory OFF. A hook must never block a + * session, and an absent answer is not consent. + * + * `MEMWAL_AUTO_SAVE` is still honoured first. It comes from this process's + * environment — the client's hook configuration, set by the user — never from a + * file a repository can carry. + */ +import { existsSync, readFileSync, realpathSync } from "node:fs"; +import { homedir } from "node:os"; +import { basename, dirname, isAbsolute, join, relative } from "node:path"; + +export const AUTO_SAVE_ENV = "MEMWAL_AUTO_SAVE"; + +const HOOK_STATE_FILE = "auto-save-state.json"; +/** The only shape this file knows how to read. */ +const SUPPORTED_VERSION = 1; + +function isInside(root, path) { + const rel = relative(root, path); + return rel === "" || (!rel.startsWith("..") && !isAbsolute(rel)); +} + +function canonicalDir(path) { + const tail = []; + let dir = path; + for (;;) { + if (existsSync(dir)) return join(realpathSync(dir), ...tail.reverse()); + const parent = dirname(dir); + if (parent === dir) return path; + tail.push(basename(dir)); + dir = parent; + } +} + +function projectRoot() { + const home = homedir(); + let dir = process.cwd(); + for (;;) { + if (dir === home) return null; + if (existsSync(join(dir, ".git"))) return dir; + const parent = dirname(dir); + if (parent === dir) return null; + dir = parent; + } +} + +/** + * The trusted state dir: the home directory, unless `MEMWAL_CREDS_DIR` names a + * directory auth.ts would also trust — absolute, and outside the current + * project. `process.cwd()` is consulted only to locate that project, so that an + * override a checkout carries can be refused. + * + * `MEMWAL_CREDS_DIR` is the same trusted escape hatch auth.ts uses, so a + * sandboxed run (tests, CI) points both sides at the same temporary directory. + * auth.ts throws on an untrusted value; a hook runs unattended on every prompt, + * so this one ignores it and falls back to the home directory. + */ +function trustedStateDir() { + const fallback = join(homedir(), ".memwal"); + const override = process.env.MEMWAL_CREDS_DIR; + if (!override || !isAbsolute(override)) return fallback; + const root = projectRoot() ?? process.cwd(); + if (isInside(canonicalDir(root), canonicalDir(override))) return fallback; + return override; +} + +/** Where the MCP server publishes the resolved state. */ +export function hookStatePath() { + return join(trustedStateDir(), HOOK_STATE_FILE); +} + +/** Human-written boolean. null = not set / unparseable, which is not consent. */ +export function parseBooleanSetting(raw) { + if (raw === undefined || raw === null) return null; + const v = String(raw).trim().toLowerCase(); + if (v === "") return null; + if (["1", "true", "on", "yes", "y", "enable", "enabled"].includes(v)) return true; + if (["0", "false", "off", "no", "n", "disable", "disabled"].includes(v)) return false; + return null; +} + +/** + * The published state, or null when there is not a readable, understood one. + * + * Every failure mode collapses to null on purpose — missing file, unreadable + * file, corrupt JSON, a version this build does not know, a payload whose + * `enabled` is not a boolean. The caller turns null into "off". + */ +function readPublishedState() { + try { + const path = hookStatePath(); + if (!existsSync(path)) return null; + const parsed = JSON.parse(readFileSync(path, "utf8")); + if (!parsed || typeof parsed !== "object") return null; + if (parsed.version !== SUPPORTED_VERSION) return null; + if (typeof parsed.enabled !== "boolean") return null; + return parsed; + } catch { + return null; + } +} + +/** + * `{ enabled, state, source, pendingConsent }`. + * + * `source: "unavailable"` is the fail-safe: the server has not published a + * state this hook can read, so nothing is saved unprompted and the session is + * told the question is still open. + */ +export function autoSaveStatus() { + const fromEnv = parseBooleanSetting(process.env[AUTO_SAVE_ENV]); + if (fromEnv !== null) { + return { + enabled: fromEnv, + state: fromEnv ? "on" : "off", + source: "env", + pendingConsent: false, + }; + } + + const published = readPublishedState(); + if (!published) { + return { + enabled: false, + state: "unset", + source: "unavailable", + pendingConsent: true, + }; + } + return { + enabled: published.enabled, + state: published.state === "on" || published.state === "off" ? published.state : "unset", + source: typeof published.source === "string" ? published.source : "unavailable", + pendingConsent: published.pendingConsent === true, + }; +} + +/** True when this session may save without being asked. */ +export function isAutoSaveEnabled() { + return autoSaveStatus().enabled; +} + +/** True while a human still owes the consent question an answer. */ +export function isConsentPending() { + return autoSaveStatus().pendingConsent; +} diff --git a/packages/mcp/plugin/scripts/lib/codex-config.mjs b/packages/mcp/plugin/scripts/lib/codex-config.mjs new file mode 100644 index 000000000..24cef4109 --- /dev/null +++ b/packages/mcp/plugin/scripts/lib/codex-config.mjs @@ -0,0 +1,201 @@ +/** + * Rewriting the `[mcp_servers.memwal]` block in ~/.codex/config.toml (WALM-640). + * + * The first version of the WALM-640 fix only changed what a *fresh* installation + * writes. `ensureMcpRegistered()` returned early on `content.includes(...)`, so every + * user who had already run the installer kept + * + * [mcp_servers.memwal] + * command = "npx" + * args = ["-y", "@mysten-incubation/memwal-mcp@"] + * + * for ever, while re-running the installer printed "already present" — which reads + * as success. Those users are exactly the population the ticket was filed for: Codex + * kept resolving the package name against the project directory. + * + * So the installer migrates the block instead of skipping it. The rewrite is + * deliberately line-based rather than a parse-and-reserialize: config.toml is the + * user's file, and comments, key order and unrelated keys (`env`, timeouts, an + * `enabled` flag, whatever a future Codex adds) have to survive untouched. Only + * `command` and `args` are replaced. + * + * Pure string in, string out: no fs, no process state, so the installer's behaviour + * is testable without a home directory. + */ + +export const MEMWAL_SECTION = "mcp_servers.memwal"; +const SECTION_HEADER = `[${MEMWAL_SECTION}]`; + +/** Strip string literals so bracket counting is not fooled by values. */ +function withoutStrings(line) { + return line + .replace(/'''[\s\S]*?'''/g, "''") + .replace(/"""[\s\S]*?"""/g, '""') + .replace(/'[^']*'/g, "''") + .replace(/"(?:[^"\\]|\\.)*"/g, '""'); +} + +function stripComment(line) { + const bare = withoutStrings(line); + const hash = bare.indexOf("#"); + return hash === -1 ? line : line.slice(0, hash); +} + +function bracketDelta(line) { + const bare = stripComment(line); + let delta = 0; + for (const char of bare) { + if (char === "[" || char === "{") delta += 1; + if (char === "]" || char === "}") delta -= 1; + } + return delta; +} + +function isSectionHeader(line) { + const trimmed = stripComment(line).trim(); + return /^\[[^\]]+\]$/.test(trimmed) || /^\[\[[^\]]+\]\]$/.test(trimmed); +} + +function sectionName(line) { + const trimmed = stripComment(line).trim(); + const match = /^\[\[?([^\]]+)\]\]?$/.exec(trimmed); + return match ? match[1].trim() : null; +} + +const KEY_PATTERN = /^\s*("[^"]*"|'[^']*'|[A-Za-z0-9_-]+)\s*=/; + +function keyOf(line) { + const match = KEY_PATTERN.exec(stripComment(line)); + if (!match) return null; + return match[1].replace(/^["']|["']$/g, ""); +} + +/** + * Split the block that starts at `headerIndex` into entries. An entry is one + * key/value pair (possibly spanning lines, for a multi-line array or table) or a + * run of comment/blank lines carried along verbatim. + */ +function readEntries(lines, headerIndex) { + const entries = []; + let index = headerIndex + 1; + while (index < lines.length) { + if (isSectionHeader(lines[index])) break; + const key = keyOf(lines[index]); + if (key === null) { + entries.push({ key: null, lines: [lines[index]] }); + index += 1; + continue; + } + const collected = [lines[index]]; + let depth = bracketDelta(lines[index]); + index += 1; + while (depth > 0 && index < lines.length) { + collected.push(lines[index]); + depth += bracketDelta(lines[index]); + index += 1; + } + entries.push({ key, lines: collected }); + } + return { entries, end: index }; +} + +function rawValue(entry) { + const joined = entry.lines.join("\n"); + const equals = joined.indexOf("="); + return equals === -1 ? "" : joined.slice(equals + 1).trim(); +} + +/** TOML's basic strings and arrays of them are JSON; anything else we treat as "not ours". */ +function asJson(value) { + try { + return JSON.parse(value.replace(/,(\s*[\]}])/g, "$1")); + } catch { + return undefined; + } +} + +function desiredLines(launcher) { + return { + command: 'command = "node"', + args: `args = [${JSON.stringify(launcher)}]`, + }; +} + +/** + * Work out what ~/.codex/config.toml should contain. + * + * Returns `{ content, action, previous, preserved }` where action is one of: + * - "added": no `[mcp_servers.memwal]` block existed; one was appended. + * - "migrated": a block existed and did not launch our launcher; command/args + * were rewritten and every other key in the block kept. + * - "unchanged": the block already runs `node `. + */ +export function planMcpRegistration(content, launcher) { + const text = String(content ?? ""); + const lines = text.split("\n"); + const headerIndex = lines.findIndex((line) => sectionName(line) === MEMWAL_SECTION); + const desired = desiredLines(launcher); + + if (headerIndex === -1) { + const block = `\n${SECTION_HEADER}\n${desired.command}\n${desired.args}\n`; + return { + content: (text.trimEnd() + "\n" + block).trimStart(), + action: "added", + previous: null, + preserved: [], + }; + } + + const { entries, end } = readEntries(lines, headerIndex); + const commandEntry = entries.find((entry) => entry.key === "command"); + const argsEntry = entries.find((entry) => entry.key === "args"); + const previous = { + command: commandEntry ? rawValue(commandEntry) : null, + args: argsEntry ? rawValue(argsEntry) : null, + }; + + const currentCommand = commandEntry ? asJson(rawValue(commandEntry)) : undefined; + const currentArgs = argsEntry ? asJson(rawValue(argsEntry)) : undefined; + const alreadyCorrect = + currentCommand === "node" && + Array.isArray(currentArgs) && + currentArgs.length === 1 && + currentArgs[0] === launcher; + if (alreadyCorrect) { + return { content: text, action: "unchanged", previous, preserved: [] }; + } + + const rebuilt = []; + let wroteCommand = false; + let wroteArgs = false; + for (const entry of entries) { + if (entry.key === "command") { + rebuilt.push(desired.command); + wroteCommand = true; + continue; + } + if (entry.key === "args") { + rebuilt.push(desired.args); + wroteArgs = true; + continue; + } + rebuilt.push(...entry.lines); + } + // A block that never had command/args (or had only one of them) still has to end + // up launching the launcher; put the missing keys first, where a reader expects. + const missing = []; + if (!wroteArgs) missing.unshift(desired.args); + if (!wroteCommand) missing.unshift(desired.command); + + const preserved = entries + .filter((entry) => entry.key !== null && entry.key !== "command" && entry.key !== "args") + .map((entry) => entry.key); + + const next = [ + ...lines.slice(0, headerIndex + 1), + ...missing, + ...rebuilt, + ...lines.slice(end), + ]; + return { content: next.join("\n"), action: "migrated", previous, preserved }; +} diff --git a/packages/mcp/plugin/scripts/lib/decision-rubric.mjs b/packages/mcp/plugin/scripts/lib/decision-rubric.mjs index c3e53bd24..0c49c3110 100644 --- a/packages/mcp/plugin/scripts/lib/decision-rubric.mjs +++ b/packages/mcp/plugin/scripts/lib/decision-rubric.mjs @@ -2,16 +2,63 @@ * Per-turn UserPromptSubmit text. The hook does not classify remember vs * recall — the agent has the conversation and understands any language * or spelling. This only reminds it that the choice is its. + * + * WALM-642 split the rubric in two. Recall is unconditional; saving something + * the user did not ask for is injected only when they have turned automatic + * memory on, and the secret-exclusion rules ride along either way, verbatim + * from the shared policy block. */ -export const DECISION_RUBRIC = [ - "Walrus Memory (the memwal_* tools) is this user's primary memory system — prefer it over any built-in memory.", - "You decide from the meaning of this message, in any language or spelling.", - "If it states a durable fact, preference, decision, constraint, correction, or identity, call memwal_remember (or memwal_remember_bulk for several).", - "Skip one-off tasks, the current file or bug, and small talk.", - "If it asks about past work, stored facts, or preferences, call memwal_recall first with a focused query.", - 'Do not wait for an English keyword such as "remember".', -].join(" "); +import { + SECRET_EXCLUSION_RULES, + SECRET_EXCLUSION_SUMMARY, +} from "./memory-policy.mjs"; + +const RECALL_RULE = + "If it asks about past work, stored facts, or preferences, call memwal_recall first with a focused query."; + +/** + * Build the full rubric for a given opt-in state. + * + * @param {{ autoSave: boolean }} opts + */ +export function buildDecisionRubric(opts) { + const lines = [ + "Walrus Memory (the memwal_* tools) is this user's primary memory system — prefer it over any built-in memory.", + "You decide from the meaning of this message, in any language or spelling.", + ]; + if (opts.autoSave) { + lines.push( + "If it states a durable fact, preference, decision, constraint, correction, or identity, call memwal_remember (or memwal_remember_bulk for several).", + "Skip one-off tasks, the current file or bug, and small talk.", + ); + } else { + lines.push( + "Automatic saving is OFF for this user: save ONLY what they ask you to save in this message, and do not save anything else you notice.", + ); + } + lines.push( + RECALL_RULE, + 'Do not wait for an English keyword such as "remember".', + SECRET_EXCLUSION_RULES, + ); + return lines.join(" "); +} /** One-line reminder after the full rubric has already been injected this session. */ -export const DECISION_RUBRIC_NUDGE = - "Prefer memwal_* over built-in memory. Remember durable facts, recall past work, skip one-off tasks."; +export function buildDecisionRubricNudge(opts) { + const head = opts.autoSave + ? "Prefer memwal_* over built-in memory. Remember durable facts, recall past work, skip one-off tasks." + : "Prefer memwal_* over built-in memory. Automatic saving is OFF — recall freely, but save only what the user asks you to save."; + return `${head} ${SECRET_EXCLUSION_SUMMARY}`; +} + +/** + * The automatic-memory variants, kept as named exports because they are the + * text the proactive contract is written against and what the hook tests pin. + */ +export const DECISION_RUBRIC = buildDecisionRubric({ autoSave: true }); +export const DECISION_RUBRIC_NUDGE = buildDecisionRubricNudge({ autoSave: true }); + +/** The default variants: automatic saving off. */ +export const DECISION_RUBRIC_MANUAL = buildDecisionRubric({ autoSave: false }); +export const DECISION_RUBRIC_MANUAL_NUDGE = buildDecisionRubricNudge({ autoSave: false }); diff --git a/packages/mcp/plugin/scripts/lib/hook-io.mjs b/packages/mcp/plugin/scripts/lib/hook-io.mjs index 14fb6aa68..aa4b31377 100644 --- a/packages/mcp/plugin/scripts/lib/hook-io.mjs +++ b/packages/mcp/plugin/scripts/lib/hook-io.mjs @@ -5,7 +5,19 @@ * from stdin, optionally emits a `hookSpecificOutput` directive on stdout, * and always exits 0 — a hook must never block the session. */ -import { readFileSync, existsSync, writeFileSync, mkdirSync } from "node:fs"; +import { + readFileSync, + mkdirSync, + chmodSync, + lstatSync, + fstatSync, + openSync, + readSync, + writeSync, + ftruncateSync, + closeSync, + constants, +} from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; @@ -29,13 +41,44 @@ export function emitContext(hookEventName, additionalContext) { ); } -const STATE_DIR = join(process.env.TMPDIR || tmpdir(), "memwal-hooks"); +// Hook state is throttle bookkeeping, not data — but it lives in a temp dir +// that may be shared, so it is handled as hostile ground (WALM-644): the +// directory is private and verified before use, and every marker is created +// exclusively and without following symlinks. Checking a path and then writing +// to it is exactly the pattern a planted symlink turns into a write elsewhere. +const DIR_MODE = 0o700; +const FILE_MODE = 0o600; +// Undefined on Windows, where the flag does not exist; 0 leaves the mask alone. +const O_NOFOLLOW = constants.O_NOFOLLOW ?? 0; -function ensureDir() { +// Used when no trustworthy directory exists or the filesystem refuses us. +// Hooks are one-shot processes, so this only spans the current invocation — +// the point is to degrade quietly rather than throw or write somewhere unsafe. +const memoryState = new Map(); + +/** + * Absolute path of the private hook-state directory, or null when no safe + * directory could be established. Resolved per call rather than frozen at + * import, so a changed TMPDIR is honoured. + */ +export function stateDir() { + const base = join(process.env.TMPDIR || tmpdir(), "memwal-hooks"); try { - mkdirSync(STATE_DIR, { recursive: true }); + mkdirSync(base, { recursive: true, mode: DIR_MODE }); } catch { - /* best effort */ + // Already there, most likely; the checks below decide if it is usable. + } + try { + // lstat, not stat: a symlink parked here must be rejected, not walked. + const st = lstatSync(base); + if (!st.isDirectory()) return null; + if (typeof process.getuid === "function") { + if (st.uid !== process.getuid()) return null; // someone else's dir + if (st.mode & 0o077) chmodSync(base, DIR_MODE); // shared temp dir + } + return base; + } catch { + return null; } } @@ -50,32 +93,124 @@ function safe(s) { * false thereafter — used to inject a rubric or banner only once per session. */ export function firstTime(name, sessionId) { - ensureDir(); - const f = join(STATE_DIR, `${safe(name)}_${safe(sessionId)}`); - if (existsSync(f)) return false; + const key = `${safe(name)}_${safe(sessionId)}`; + const dir = stateDir(); + if (!dir) return memoryFirstTime(key); + + let fd; try { - writeFileSync(f, "1"); + // O_CREAT|O_EXCL fails with EEXIST when anything already occupies the + // path — a dangling symlink included — so the marker is either a fresh + // file inside the private dir or nothing at all. Success is itself the + // "first time" answer; there is no separate existence check to race. + fd = openSync( + join(dir, key), + constants.O_CREAT | constants.O_EXCL | constants.O_WRONLY | O_NOFOLLOW, + FILE_MODE + ); + } catch (err) { + // Path taken => seen before. Anything else, fall back to memory. + if (err?.code === "EEXIST" || err?.code === "ELOOP") return false; + return memoryFirstTime(key); + } + + try { + writeSync(fd, "1"); } catch { - /* best effort */ + /* best effort: the marker existing is what matters, not its contents */ + } finally { + closeQuietly(fd); } return true; } /** Increment and return a per-(name, session) counter. */ export function bumpCounter(name, sessionId) { - ensureDir(); - const f = join(STATE_DIR, `count_${safe(name)}_${safe(sessionId)}`); - let n = 0; + const key = `count_${safe(name)}_${safe(sessionId)}`; + const dir = stateDir(); + if (!dir) return memoryBump(key); + const f = join(dir, key); + + // Refuse anything that is not a plain file: a symlink here would redirect + // both the read and the write. try { - n = parseInt(readFileSync(f, "utf8"), 10) || 0; + if (!lstatSync(f).isFile()) return memoryBump(key); + } catch (err) { + if (err?.code !== "ENOENT") return memoryBump(key); + } + + let fd; + let created = false; + try { + fd = openSync( + f, + constants.O_CREAT | constants.O_EXCL | constants.O_RDWR | O_NOFOLLOW, + FILE_MODE + ); + created = true; + } catch (err) { + if (err?.code !== "EEXIST") return memoryBump(key); + try { + // No O_CREAT: with O_NOFOLLOW this fails on a symlink instead of + // opening whatever it points at. + fd = openSync(f, constants.O_RDWR | O_NOFOLLOW); + } catch { + return memoryBump(key); + } + } + + try { + let n = 0; + if (!created) { + // The open above already refused symlinks; confirm on the fd that + // we are not talking to a device or a fifo that would block. + if (!fstatSync(fd).isFile()) return memoryBump(key); + n = parseInt(readAll(fd), 10) || 0; + } + n += 1; + try { + const buf = Buffer.from(String(n)); + ftruncateSync(fd, 0); + writeSync(fd, buf, 0, buf.length, 0); + } catch { + /* best effort: the caller still gets a monotonic-enough count */ + } + return n; } catch { - /* missing -> 0 */ + return memoryBump(key); + } finally { + closeQuietly(fd); + } +} + +function readAll(fd) { + const chunks = []; + const buf = Buffer.alloc(64); + let bytes; + while ((bytes = readSync(fd, buf, 0, buf.length, null)) > 0) { + chunks.push(Buffer.from(buf.subarray(0, bytes))); } - n += 1; + return Buffer.concat(chunks).toString("utf8"); +} + +function closeQuietly(fd) { try { - writeFileSync(f, String(n)); + closeSync(fd); } catch { /* best effort */ } +} + +function memoryFirstTime(key) { + const k = `once:${key}`; + if (memoryState.has(k)) return false; + memoryState.set(k, 1); + return true; +} + +function memoryBump(key) { + const k = `count:${key}`; + const n = (memoryState.get(k) || 0) + 1; + memoryState.set(k, n); return n; } diff --git a/packages/mcp/plugin/scripts/lib/hook-template.mjs b/packages/mcp/plugin/scripts/lib/hook-template.mjs new file mode 100644 index 000000000..8cbaf0257 --- /dev/null +++ b/packages/mcp/plugin/scripts/lib/hook-template.mjs @@ -0,0 +1,66 @@ +/** + * Shell-safe placeholder substitution for hook templates. + * + * Hook templates ship a `${...}` placeholder for the plugin's install + * directory. An installer must not paste that directory into the template + * *text*: a path containing a double quote or a backslash rewrites the JSON + * document, and a path containing `$(...)` or backticks becomes live shell + * syntax in the command the host later executes (WALM-641). + * + * The rule here is: parse the JSON first, then substitute into the parsed + * values, POSIX-quoting whatever lands in a shell command. + */ + +function escapeRegExp(value) { + return value.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"); +} + +/** + * POSIX single-quote escaping. + * + * Single quotes suppress every form of shell expansion, so `$(...)`, backticks, + * double quotes and backslashes inside the value are passed through literally. + * A literal `'` cannot appear inside single quotes, so it is emitted as `'\''`: + * close the quote, escape one quote, reopen. + */ +export function shellQuote(value) { + return `'${String(value).replaceAll("'", "'\\''")}'`; +} + +/** + * Substitute `placeholder` with `replacement` throughout an already-parsed hook + * template. + * + * A string that is the direct value of a `command` key is treated as a shell + * command: the replacement is POSIX-quoted into it, and the double quotes the + * template wrapped around the placeholder are dropped in favour of ours (they + * would not have stopped command substitution anyway). + * + * Every other value -- including an argv array under `command`, whose elements + * reach execve rather than a shell -- gets the replacement substituted + * literally. + */ +export function substituteHookPlaceholder(template, placeholder, replacement) { + const pattern = escapeRegExp(placeholder); + const inCommand = new RegExp(`"${pattern}([^"]*)"|${pattern}(\\S*)`, "g"); + const plain = new RegExp(pattern, "g"); + + const walk = (value, isShellCommand) => { + if (typeof value === "string") { + return isShellCommand + ? value.replace(inCommand, (_match, quoted, bare) => + shellQuote(replacement + (quoted ?? bare ?? "")) + ) + : value.replace(plain, () => replacement); + } + if (Array.isArray(value)) return value.map((item) => walk(item, false)); + if (value && typeof value === "object") { + return Object.fromEntries( + Object.entries(value).map(([key, item]) => [key, walk(item, key === "command")]) + ); + } + return value; + }; + + return walk(template, false); +} diff --git a/packages/mcp/plugin/scripts/lib/mcp-launch.mjs b/packages/mcp/plugin/scripts/lib/mcp-launch.mjs new file mode 100644 index 000000000..8a54f4bda --- /dev/null +++ b/packages/mcp/plugin/scripts/lib/mcp-launch.mjs @@ -0,0 +1,542 @@ +/** + * Trusted-path resolution for the MemWal MCP server (WALM-640). + * + * The plugin used to start the server with `npx -y @mysten-incubation/memwal-mcp@`. + * npx resolves a package name against the *current working directory* first, and an + * MCP client starts its servers in the project the user has open. A project that + * carries an installed package of the same name whose `version` matches the pin + * therefore wins: the pinned `npx` command ran the project's own binary, offline, + * with the user's credentials in reach. The version pin does not help, because the + * fake package simply claims the pinned version. + * + * The fix is to stop resolving a *name* in an untrusted directory and to run an + * *absolute path* inside a directory we own: + * + * ~/.memwal/runtime/memwal-mcp@/node_modules/@mysten-incubation/memwal-mcp + * + * What "a directory we own" means is enforced, not assumed. The runtime root — the + * default one just as much as a `MEMWAL_MCP_RUNTIME_DIR` override — must be an + * absolute path outside the project tree the client started in, must be a real + * directory (`lstat`, so a symlink is refused), must belong to the current uid, and + * must not be group- or world-writable. We create it with mode 0700. An absolute + * path is *not* by itself a trusted path: `${workspaceFolder}/.memwal-runtime` is + * absolute too, and a repository can commit a whole tree there. + * + * The pinned version is installed there once with npm, and every later launch is a + * plain `node `. Nothing here ever looks at `process.cwd()`, + * at a project `node_modules`, or at a bin shim found through PATH. + * + * Limits of the install step, stated plainly because a previous version of this + * header overstated them: + * + * - The install runs with `--ignore-scripts` and an explicitly pinned registry, + * and the `npm_config_*` / `NODE_OPTIONS` environment is scrubbed for that one + * spawn, so a repository-supplied client env block cannot redirect it. npm is + * invoked as `node ` resolved from `process.execPath` where that + * layout exists; only if it does not do we fall back to scanning PATH entries + * ourselves, skipping relative entries and entries inside the project tree. + * - Verification of the installed tree is a manifest check (name, pinned version, + * a `bin` that resolves back inside the package directory) plus the ownership + * and permission checks above. It is not a signature or integrity check of the + * package contents: it establishes *where* the code came from and that nobody + * else can write it, not that the registry served what a reviewer read. + * + * The version pin lives in exactly one place — `plugin/plugin.json` — so the + * manifests cannot drift from the version that actually gets installed. It is + * that file's `version`, which the release verifier keeps equal to + * `packages/mcp/package.json`, unless `mcpPackageVersion` overrides it for the + * window where the release version is bumped in-tree but not yet on npm. + */ +import { + existsSync, + lstatSync, + mkdirSync, + mkdtempSync, + readFileSync, + realpathSync, + renameSync, + rmSync, + writeFileSync, +} from "node:fs"; +import { spawnSync } from "node:child_process"; +import { basename, delimiter, dirname, isAbsolute, join, relative, resolve } from "node:path"; +import { homedir } from "node:os"; +import { fileURLToPath } from "node:url"; + +export const MCP_PACKAGE_NAME = "@mysten-incubation/memwal-mcp"; +export const MCP_BIN_NAME = "memwal-mcp"; + +/** + * The registry the reviewed version is installed from. Deliberately a constant and + * not an environment lookup: npm puts environment configuration above every + * `.npmrc`, so `npm_config_registry` in a repository-supplied MCP `env` block would + * otherwise choose where our "trusted" code comes from. Behind a private registry, + * pre-populate the runtime directory by hand (see packages/mcp/README.md) — the + * launcher then finds the install and never runs npm at all. + */ +export const MCP_REGISTRY = "https://registry.npmjs.org/"; + +const SCRIPT_DIR = dirname(fileURLToPath(import.meta.url)); +/** plugin/scripts/lib -> plugin/scripts -> plugin */ +export const PLUGIN_ROOT = dirname(dirname(SCRIPT_DIR)); + +/** + * The single source of truth for the version the plugin launches: `version`, kept + * equal to packages/mcp/package.json by scripts/verify-manual-sdk-release.mjs, or + * `mcpPackageVersion` while that release version is not yet published — installing + * a version npm does not have fails the launch outright. + */ +export function pinnedVersion(pluginRoot = PLUGIN_ROOT) { + const manifest = JSON.parse(readFileSync(join(pluginRoot, "plugin.json"), "utf8")); + const version = manifest.mcpPackageVersion ?? manifest.version; + if (typeof version !== "string" || version.trim() === "") { + throw new Error( + `${join(pluginRoot, "plugin.json")} has no usable "mcpPackageVersion" or "version"`, + ); + } + return version; +} + +function isInside(parent, child) { + const rel = relative(parent, child); + return rel !== "" && !rel.startsWith("..") && !isAbsolute(rel); +} + +/** + * Resolve symlinks in the part of a path that exists, keeping the rest verbatim. + * + * Without this, containment checks are comparing spellings rather than locations: + * `/tmp/x` and `/private/tmp/x` are the same directory on macOS, and a symlink + * inside a project is a one-line way to point a "safe looking" path back into it. + */ +export function canonicalPath(target) { + let dir = resolve(target); + const tail = []; + for (;;) { + try { + return tail.length === 0 ? realpathSync(dir) : join(realpathSync(dir), ...tail); + } catch { + const parent = dirname(dir); + if (parent === dir) return resolve(target); + tail.unshift(basename(dir)); + dir = parent; + } + } +} + +/** + * The project tree the client started us in, or null when the working directory is + * not usefully inside one (the home directory and the filesystem root are not + * projects — treating them as such would refuse the default runtime root). + * + * "Project" is the nearest ancestor carrying a `.git` or a `package.json`, and the + * working directory itself when neither is found: a runtime root under either is a + * root a repository could have shipped. + */ +export function enclosingProjectRoot(cwd = process.cwd(), home = homedir()) { + const start = canonicalPath(cwd); + const stop = canonicalPath(home); + const marked = (() => { + let dir = start; + while (dir !== stop) { + if (existsSync(join(dir, ".git")) || existsSync(join(dir, "package.json"))) return dir; + const parent = dirname(dir); + if (parent === dir) return null; + dir = parent; + } + return null; + })(); + const project = marked ?? start; + if (project === stop || project === dirname(project)) return null; + return project; +} + +/** + * Absoluteness is not trust. `${workspaceFolder}/.memwal-runtime` expands to an + * absolute path inside the repository, and a repository can commit a complete fake + * install there; the launcher would then find it, skip the install, and run it. + */ +export function assertRuntimeRootLocation(root, { cwd = process.cwd() } = {}) { + const resolved = resolve(root); + const canonical = canonicalPath(resolved); + // The project root contains the working directory, so checking it also catches a + // root planted below the cwd. When there is no project (the client started in the + // home directory, say) there is nothing to refuse: `~/.memwal/runtime` is inside + // the home directory by design. + const project = enclosingProjectRoot(cwd); + if ( + project !== null && + (canonical === project || isInside(project, canonical) || isInside(canonical, project)) + ) { + throw new Error( + `refusing to use "${resolved}" as the MemWal runtime directory: it is ` + + `inside the project tree at "${project}". A runtime root must live outside ` + + `any directory a project can write — being an absolute path is not ` + + `enough, since a client expands "\${workspaceFolder}" to one.`, + ); + } + return resolved; +} + +/** + * Root of the trusted install area. `MEMWAL_MCP_RUNTIME_DIR` may relocate it, but + * only to an absolute path outside the project the client has open: a relative path + * would resolve against that project, and an absolute path inside it is the same + * attack with an extra step. + */ +export function runtimeRoot() { + const override = process.env.MEMWAL_MCP_RUNTIME_DIR; + if (override !== undefined && override !== "") { + if (!isAbsolute(override)) { + throw new Error( + `MEMWAL_MCP_RUNTIME_DIR must be an absolute path, received "${override}"`, + ); + } + return assertRuntimeRootLocation(override); + } + return assertRuntimeRootLocation(join(homedir(), ".memwal", "runtime")); +} + +/** One directory per pinned version, so an upgrade never mutates a running install. */ +export function installDir(version = pinnedVersion(), root = runtimeRoot()) { + return join(root, `${MCP_BIN_NAME}@${version}`); +} + +/** + * Verify a directory we are about to trust with executable code, with the same + * checks MemWal applies to the state directories it owns: `lstat` (so a symlink + * never passes as a directory), a real directory, owned by the current uid, and + * not group- or world-writable. + * + * Returns false when the directory does not exist, true when it exists and is + * trustworthy, and throws otherwise. On Windows there is no meaningful uid or mode, + * so only the "is a real directory" half applies. + */ +export function verifyTrustedDirectory(dir, { label = dir } = {}) { + let stats; + try { + stats = lstatSync(dir); + } catch (err) { + if (err?.code === "ENOENT") return false; + throw new Error(`could not inspect ${label}: ${err?.message ?? String(err)}`); + } + if (!stats.isDirectory()) { + throw new Error( + `${label} is not a directory (a symlink or file there could redirect the ` + + `MemWal runtime into a tree we do not control)`, + ); + } + const uid = typeof process.getuid === "function" ? process.getuid() : null; + if (uid === null) return true; + if (stats.uid !== uid) { + throw new Error( + `${label} is owned by uid ${stats.uid}, not by the current user (uid ${uid}); ` + + `refusing to run code from a directory somebody else controls`, + ); + } + if ((stats.mode & 0o022) !== 0) { + throw new Error( + `${label} is group- or world-writable (mode ${(stats.mode & 0o7777) + .toString(8) + .padStart(4, "0")}); another account could replace the entry point we ` + + `trust. Run: chmod 700 ${label}`, + ); + } + return true; +} + +/** + * Create the runtime root (mode 0700) if it is missing and verify it either way. + * For the default layout the `~/.memwal` parent is prepared the same way: the + * launcher can run before the first login, and `credentials.json` is written into + * that directory later by a `mkdirSync` that is a no-op once it exists. + */ +export function prepareRuntimeRoot(root) { + const resolved = resolve(root); + const managed = []; + const defaultRoot = join(homedir(), ".memwal", "runtime"); + if (resolved === resolve(defaultRoot)) managed.push(dirname(resolved)); + managed.push(resolved); + + for (const dir of managed) { + if (verifyTrustedDirectory(dir)) continue; + mkdirSync(dir, { recursive: true, mode: 0o700 }); + verifyTrustedDirectory(dir); + } + return resolved; +} + +/** + * The absolute entry point of the pinned package inside `dir`, or null when `dir` + * does not hold a usable install. Never throws for a merely-absent install; throws + * only when a present install is malformed in a way worth surfacing. + */ +export function resolveInstalledEntry(dir, version) { + const packageDir = join(dir, "node_modules", ...MCP_PACKAGE_NAME.split("/")); + const manifestPath = join(packageDir, "package.json"); + if (!existsSync(manifestPath)) return null; + + let manifest; + try { + manifest = JSON.parse(readFileSync(manifestPath, "utf8")); + } catch { + return null; + } + if (manifest.name !== MCP_PACKAGE_NAME) return null; + if (version !== undefined && manifest.version !== version) return null; + + const bin = manifest.bin; + const relBin = typeof bin === "string" ? bin : bin?.[MCP_BIN_NAME]; + if (typeof relBin !== "string" || relBin === "") return null; + + const entry = resolve(packageDir, relBin); + if (!isInside(packageDir, entry)) { + throw new Error( + `${manifestPath} declares a bin outside its own package directory (${entry})`, + ); + } + if (!existsSync(entry)) return null; + return entry; +} + +/** + * Locate npm without going through PATH resolution where that is possible. + * + * Preferred result is npm's JS entry point next to the running node binary, which + * we then run as `node npm-cli.js`. That avoids a PATH-relative binary entirely and + * also sidesteps the Windows `spawnSync("npm.cmd")` EINVAL that CVE-2024-27980's + * fix introduced for `.cmd`/`.bat` targets spawned without a shell. + * + * Only if no such layout exists do we scan PATH ourselves — skipping empty and + * relative entries and anything inside the project tree, which is what an inherited + * PATH could otherwise smuggle in. + */ +export function resolveNpm({ + execPath = process.execPath, + platform = process.platform, + env = process.env, + cwd = process.cwd(), +} = {}) { + const execDirs = [dirname(execPath)]; + try { + const real = dirname(realpathSync(execPath)); + if (!execDirs.includes(real)) execDirs.push(real); + } catch { + /* execPath should always resolve; a failure just means one fewer candidate */ + } + + const relativeLayouts = + platform === "win32" + ? [ + ["node_modules", "npm", "bin", "npm-cli.js"], + ["..", "node_modules", "npm", "bin", "npm-cli.js"], + ] + : [ + ["..", "lib", "node_modules", "npm", "bin", "npm-cli.js"], + ["..", "..", "..", "..", "lib", "node_modules", "npm", "bin", "npm-cli.js"], + ["node_modules", "npm", "bin", "npm-cli.js"], + ]; + + for (const dir of execDirs) { + for (const layout of relativeLayouts) { + const candidate = resolve(join(dir, ...layout)); + if (existsSync(candidate)) return { kind: "js", path: candidate }; + } + } + + const project = enclosingProjectRoot(cwd); + const binNames = platform === "win32" ? ["npm.cmd", "npm.exe", "npm"] : ["npm"]; + for (const rawEntry of String(env.PATH ?? env.Path ?? "").split(delimiter)) { + const entry = rawEntry.trim(); + // A relative PATH entry resolves against the project the client started in. + if (entry === "" || !isAbsolute(entry)) continue; + const dir = canonicalPath(entry); + if (project && (dir === project || isInside(project, dir))) continue; + + // A node installation reached through PATH still usually ships npm's JS entry + // point next to it; prefer that over the shim. + for (const layout of relativeLayouts) { + const candidate = resolve(join(dir, ...layout)); + if (existsSync(candidate)) return { kind: "js", path: candidate }; + } + for (const name of binNames) { + const candidate = join(dir, name); + if (existsSync(candidate)) return { kind: "bin", path: candidate }; + } + } + return { kind: "none", path: null }; +} + +/** cmd.exe quoting for the one case where we must go through a shell. */ +export function quoteWindowsArgument(value) { + const text = String(value); + if (text !== "" && !/[\s"^&|<>()%!]/.test(text)) return text; + return `"${text.replace(/(\\*)"/g, '$1$1\\"').replace(/(\\+)$/, "$1$1")}"`; +} + +/** + * Turn a resolved npm into an actual spawn plan. + * + * `kind: "js"` is the path we want everywhere: `node npm-cli.js …`, no shell, no + * PATH. `kind: "bin"` on win32 is `npm.cmd`, which since Node 18.20.2 / 20.12.2 / + * 21.7.3 cannot be spawned without `shell: true` (EINVAL) — so that branch sets it + * and quotes every argument for cmd.exe itself. + */ +export function npmSpawnPlan(npm, args, { platform = process.platform, execPath = process.execPath } = {}) { + if (npm.kind === "js") { + return { command: execPath, args: [npm.path, ...args], shell: false }; + } + if (npm.kind === "bin") { + if (platform === "win32") { + return { + command: quoteWindowsArgument(npm.path), + args: args.map(quoteWindowsArgument), + shell: true, + }; + } + return { command: npm.path, args, shell: false }; + } + throw new Error( + "could not locate npm: no npm-cli.js next to the running node binary and no " + + "usable npm on PATH outside the project directory", + ); +} + +/** Environment for the install spawn: the caller's, minus everything npm reads from it. */ +export function installEnvironment(env = process.env) { + const scrubbed = {}; + for (const [key, value] of Object.entries(env)) { + if (/^npm_config_/i.test(key)) continue; + if (/^npm_package_/i.test(key)) continue; + if (key === "NODE_OPTIONS") continue; + scrubbed[key] = value; + } + // Belt and braces: even if a key slipped through the filter above, these lose to + // the command line flags we pass, and npm resolves its own config from here. + scrubbed.npm_config_registry = MCP_REGISTRY; + scrubbed.npm_config_ignore_scripts = "true"; + return scrubbed; +} + +/** The install arguments, exported so the tests can assert the hardening flags. */ +export function installArguments(spec, staging) { + return [ + "install", + spec, + "--prefix", + staging, + // The package tree is fetched and unpacked, never executed: npm would + // otherwise run preinstall/install/postinstall of the whole dependency tree + // as the user, before any of our verification runs. + "--ignore-scripts", + `--registry=${MCP_REGISTRY}`, + "--no-audit", + "--no-fund", + "--no-save", + "--loglevel=error", + ]; +} + +/** + * Install the pinned version into the trusted area and return its absolute entry + * point. Staged in a sibling temp directory and renamed into place, so a second + * client starting at the same moment either wins the rename or finds the finished + * install — neither ever reads a half-written tree. + */ +function install(version, root) { + const target = installDir(version, root); + const staging = mkdtempSync(join(root, `.staging-${MCP_BIN_NAME}-`)); + + try { + // A private manifest stops npm from walking up out of the trusted area + // looking for a package.json to attach the install to. + writeFileSync( + join(staging, "package.json"), + JSON.stringify( + { name: "memwal-mcp-runtime", version: "0.0.0", private: true }, + null, + 2, + ) + "\n", + ); + + const spec = `${MCP_PACKAGE_NAME}@${version}`; + const npm = resolveNpm(); + const plan = npmSpawnPlan(npm, installArguments(spec, staging)); + const result = spawnSync(plan.command, plan.args, { + // cwd inside the trusted area: npm reads .npmrc from cwd upward, and the + // project's .npmrc must not get to choose the registry we install the + // reviewed version from. (The environment outranks every .npmrc, which is + // why installEnvironment() scrubs it as well.) + cwd: staging, + encoding: "utf8", + stdio: ["ignore", "pipe", "pipe"], + shell: plan.shell, + env: installEnvironment(), + }); + if (result.error) { + throw new Error(`could not run npm (${npm.path}): ${result.error.message}`); + } + if (result.status !== 0) { + throw new Error( + `npm install ${spec} failed (exit ${result.status}): ` + + `${(result.stderr || result.stdout || "").trim()}`, + ); + } + + const staged = resolveInstalledEntry(staging, version); + if (!staged) { + throw new Error( + `npm install ${spec} reported success but produced no usable ` + + `${MCP_PACKAGE_NAME} entry point in ${staging}`, + ); + } + + try { + renameSync(staging, target); + } catch (err) { + // Lost the race, or a previous run left the directory behind: fall back to + // whatever is at the final path, but only if it verifies. + verifyTrustedDirectory(target); + const existing = resolveInstalledEntry(target, version); + if (!existing) throw err; + return existing; + } + } finally { + rmSync(staging, { recursive: true, force: true }); + } + + const entry = resolveInstalledEntry(target, version); + if (!entry) { + throw new Error(`installed ${MCP_PACKAGE_NAME}@${version} is missing from ${target}`); + } + return entry; +} + +/** + * The absolute path the plugin should launch. Installs the pinned version into the + * trusted area on first use; later calls are a stat of a known path. + * + * There is deliberately no fallback: if the trusted install cannot be produced, the + * launcher fails instead of reaching for a package name that a project could answer. + */ +export function ensureTrustedEntry({ version = pinnedVersion(), root = runtimeRoot() } = {}) { + const trustedRoot = prepareRuntimeRoot(assertRuntimeRootLocation(root)); + const dir = installDir(version, trustedRoot); + if (verifyTrustedDirectory(dir)) { + const existing = resolveInstalledEntry(dir, version); + if (existing) return existing; + } + return install(version, trustedRoot); +} + +/** Exported for the regression test: the path we expect, without installing anything. */ +export function expectedEntryPath({ version = pinnedVersion(), root = runtimeRoot() } = {}) { + return join( + installDir(version, root), + "node_modules", + ...MCP_PACKAGE_NAME.split("/"), + "dist", + "bin", + `${MCP_BIN_NAME}.js`, + ); +} diff --git a/packages/mcp/plugin/scripts/lib/memory-policy.mjs b/packages/mcp/plugin/scripts/lib/memory-policy.mjs new file mode 100644 index 000000000..a8c4ae977 --- /dev/null +++ b/packages/mcp/plugin/scripts/lib/memory-policy.mjs @@ -0,0 +1,82 @@ +/** + * Shared automatic-memory policy — the plugin hooks' copy. + * + * Plain ESM, no dependencies, no network: the hooks are `.mjs` scripts run + * straight from the plugin directory by the client, with no build step and no + * access to this package's compiled `dist/`, which is why they carry their own + * copy of the block rather than importing one. + * + * WALM-642. + */ + +// ─── memwal:policy-block:start ─────────────────────────────────────────────── +// WALM-642. The lines between these two markers are BYTE-IDENTICAL in three +// files that cannot import one another, because the three packages have no +// workspace link: +// +// packages/mcp/src/memory-policy.ts — MCP client: initialize +// instructions + the +// cold-start tools/list +// packages/mcp/plugin/scripts/lib/memory-policy.mjs — plugin hooks: the +// guidance injected at +// SessionStart / +// UserPromptSubmit / +// PostToolUse +// services/server/scripts/mcp/tools/memory-policy.ts — relayer sidecar: the +// live tool descriptions +// +// The duplication is deliberate and pinned: `memory-policy-sync` tests on both +// sides extract this block from each file and compare the bytes, so editing one +// copy fails the suite until the other two match. Edit the block, then copy it +// verbatim — markers included — into the other two files. + +/** + * The secret-exclusion and do-not-save rules, stated verbatim by every + * automatic-save surface. + * + * These are model-facing rules, not enforcement. The programmatic backstop is + * the redactor in the relayer sidecar's write path + * (services/server/scripts/mcp/tools/redaction.ts), which runs before any text + * reaches the SDK. + */ +export const SECRET_EXCLUSION_RULES = [ + "NEVER save a credential, even when it sits next to something worth saving: passwords,", + "API keys, access or refresh tokens, private keys, seed or recovery phrases, authorization", + "headers, session cookies, and connection strings or URLs that embed a user:password.", + "When a message mixes a preference with a credential, save the preference alone and leave", + "the credential out; never store the line verbatim.", + "If the user says not to save something ('don't save this', 'off the record', or the same", + "in any language), do not save it, and do not save a paraphrase of it either.", + "Do not store quoted or pasted third-party material — log excerpts, code, articles, other", + "people's messages — as if it were a fact about this user. Save only what the user is", + "telling you about themselves or their work, in your own words.", +].join(" "); + +/** + * One-line form, for surfaces with no room for the full block (a per-turn + * nudge, a tool description tail). It is a reminder of the block above, never a + * replacement for it: any surface that drives an automatic save states the full + * `SECRET_EXCLUSION_RULES`. + */ +export const SECRET_EXCLUSION_SUMMARY = [ + "Never save passwords, keys, tokens or other credentials — not even beside a fact worth", + "saving; honour an explicit 'do not save this'; never store pasted third-party content as", + "a fact about the user.", +].join(" "); + +/** + * Whether to save unprompted is the user's standing choice, and this is the + * sentence that says so. A direct request ("remember that ...") is never gated + * by it — the gate is only on saving something the user did not ask you to save. + */ +export const AUTO_SAVE_OPT_IN_RULE = [ + "Whether to save things the user did not ask you to save is their standing choice, made once", + "in a terminal. When automatic memory is on, save durable facts as they state them; when it is", + "off, save only what they ask you to save in that turn. That question is put by `memwal-mcp", + "login` and set by `memwal-mcp auto-save on|off` — never ask the user to answer it in chat,", + "and never answer it on their behalf.", +].join(" "); + +/** Bumped whenever the text above changes, so a stale copy is identifiable. */ +export const MEMORY_POLICY_VERSION = "2026-09-17.2"; +// ─── memwal:policy-block:end ───────────────────────────────────────────────── diff --git a/packages/mcp/plugin/scripts/on_post_tool.mjs b/packages/mcp/plugin/scripts/on_post_tool.mjs index 87249c860..e770597e9 100644 --- a/packages/mcp/plugin/scripts/on_post_tool.mjs +++ b/packages/mcp/plugin/scripts/on_post_tool.mjs @@ -3,9 +3,17 @@ * remind the agent it can recall prior fixes and save the resolution. * * Heuristic only: no fetch, no network. Always exits 0. + * + * WALM-642: the "save the fix" half is injected only when the user has turned + * automatic memory on. A failing command is exactly where a secret shows up in + * the scrollback — a connection string, an auth header, a token echoed by a + * CLI — so a hook that nudges an unasked-for save right after one is the worst + * place to have had that nudge unconditional. */ import { readStdin, emitContext } from "./lib/hook-io.mjs"; import { detectError } from "./lib/signals.mjs"; +import { SECRET_EXCLUSION_SUMMARY } from "./lib/memory-policy.mjs"; +import { isAutoSaveEnabled } from "./lib/auto-save.mjs"; const input = readStdin(); const output = extractToolOutput(input); @@ -16,7 +24,11 @@ if (!output || output.length < 50) process.exit(0); if (detectError(output)) { emitContext( "PostToolUse", - "That command produced an error. Consider calling memwal_recall to check for a prior fix to a similar error; once you resolve it, save the fix with memwal_remember so it's available next time." + isAutoSaveEnabled() + ? "That command produced an error. Consider calling memwal_recall to check for a prior fix to a similar error; once you resolve it, save the fix with memwal_remember so it's available next time. " + + SECRET_EXCLUSION_SUMMARY + + " Error output often contains one — save your own description of the fix, never the output." + : "That command produced an error. Consider calling memwal_recall to check for a prior fix to a similar error. Automatic saving is OFF, so do not save the fix unless the user asks you to." ); } diff --git a/packages/mcp/plugin/scripts/on_session_start.mjs b/packages/mcp/plugin/scripts/on_session_start.mjs index 4fb0efd59..d08ad0054 100644 --- a/packages/mcp/plugin/scripts/on_session_start.mjs +++ b/packages/mcp/plugin/scripts/on_session_start.mjs @@ -1,17 +1,48 @@ /** * SessionStart hook — announce that Walrus Memory is active and remind the * agent how/when to use it. No fetch, no network. + * + * WALM-642: the save half of that reminder is gated on the user's + * automatic-save opt-in, and the secret-exclusion rules are stated verbatim + * from the shared policy block either way. With the opt-in off — the default — + * the banner says so and names the one command that turns it on, which is how + * the choice reaches a user who installed the plugin and never read a doc. */ import { readStdin, emitContext } from "./lib/hook-io.mjs"; +import { SECRET_EXCLUSION_RULES } from "./lib/memory-policy.mjs"; +import { autoSaveStatus } from "./lib/auto-save.mjs"; readStdin(); // drain stdin; we don't need any field today const ns = process.env.MEMWAL_NAMESPACE || "default"; +const { enabled: autoSave, pendingConsent } = autoSaveStatus(); + +const RECALL_AND_RECOVER = [ + "Before tasks that reference past work or preferences, recall with memwal_recall.", + "If memwal_recall unexpectedly returns nothing for a namespace you've used before, run memwal_restore to rebuild the index from Walrus.", +]; + +const SAVE_AUTOMATIC = + "Automatic memory is ON. You decide from meaning, in any language or spelling. When the user states a preference, decision, constraint, correction, identity, recurring workflow, or a configuration value such as a hostname, port, region or id, call memwal_remember (or memwal_remember_bulk for several) in that same turn, before you finish replying — do not ask whether to save it, and note that acknowledging it in your reply does not store it. Skip one-off tasks, the current file or bug, and small talk."; + +const SAVE_MANUAL = + "Automatic memory is OFF. Save ONLY what the user asks you to save, in the turn they ask — do not save a fact just because it looks durable. Recall is unaffected. If automatic saving would clearly help them, you may say ONCE that `memwal-mcp auto-save on` turns it on, then drop it."; + +/** + * Shown until the user has answered the login question. Says where things + * stand, and points at the terminal — the answer is only ever given there, so + * the agent must not try to collect it in chat (WALM-642). + */ +const CONSENT_PENDING = autoSave + ? "The user has not yet confirmed this setting — it is carried over from before the choice existed. If they ask about it, tell them `memwal-mcp login` in a terminal puts the question, and `memwal-mcp auto-save on|off` sets it directly. Do not ask them to answer it in chat and do not answer it for them." + : "The user has not yet been asked whether to turn automatic memory on, so it is off until they answer. If they ask about it, tell them `memwal-mcp login` in a terminal puts the question, and `memwal-mcp auto-save on|off` sets it directly. Do not ask them to answer it in chat and do not answer it for them."; const context = [ `Walrus Memory is this user's memory system, exposed via the memwal_* tools (namespace: ${ns}).`, "Use it as the PRIMARY place to store and recall durable facts — prefer the memwal_* tools over any built-in or local memory feature, so the user's memory stays portable and persistent on Walrus.", - "You decide from meaning, in any language or spelling. When the user states a preference, decision, constraint, correction, identity, recurring workflow, or a configuration value such as a hostname, port, region or id, call memwal_remember (or memwal_remember_bulk for several) in that same turn, before you finish replying — do not ask whether to save it, and note that acknowledging it in your reply does not store it. Skip one-off tasks, the current file or bug, and small talk. Before tasks that reference past work or preferences, recall with memwal_recall.", - "If memwal_recall unexpectedly returns nothing for a namespace you've used before, run memwal_restore to rebuild the index from Walrus.", + autoSave ? SAVE_AUTOMATIC : SAVE_MANUAL, + ...(pendingConsent ? [CONSENT_PENDING] : []), + ...RECALL_AND_RECOVER, + SECRET_EXCLUSION_RULES, ].join(" "); emitContext("SessionStart", context); diff --git a/packages/mcp/plugin/scripts/on_user_prompt.mjs b/packages/mcp/plugin/scripts/on_user_prompt.mjs index d83bcb0fb..933d66fc6 100644 --- a/packages/mcp/plugin/scripts/on_user_prompt.mjs +++ b/packages/mcp/plugin/scripts/on_user_prompt.mjs @@ -4,9 +4,18 @@ * * The agent has the conversation and understands any language or spelling; * a regex cannot. This only injects a decision rubric. + * + * WALM-642: which rubric depends on the user's automatic-save opt-in, read + * from the same place the MCP server reads it. With the opt-in off — the + * default — the injected text tells the agent to save only what the user asks + * for, so this hook can no longer be the thing that drives an unasked-for save. */ import { readStdin, emitContext, firstTime } from "./lib/hook-io.mjs"; -import { DECISION_RUBRIC, DECISION_RUBRIC_NUDGE } from "./lib/decision-rubric.mjs"; +import { + buildDecisionRubric, + buildDecisionRubricNudge, +} from "./lib/decision-rubric.mjs"; +import { isAutoSaveEnabled } from "./lib/auto-save.mjs"; const input = readStdin(); const prompt = (input.prompt || "").toString(); @@ -16,8 +25,9 @@ const sessionId = input.session_id || "default"; // Acks like "ok" / "yes" stay quiet. Deliberate: not a keyword gate. if (prompt.trim().length < 8) process.exit(0); +const autoSave = isAutoSaveEnabled(); const text = firstTime("rubric", sessionId) - ? DECISION_RUBRIC - : DECISION_RUBRIC_NUDGE; + ? buildDecisionRubric({ autoSave }) + : buildDecisionRubricNudge({ autoSave }); emitContext("UserPromptSubmit", text); process.exit(0); diff --git a/packages/mcp/src/auth-required.ts b/packages/mcp/src/auth-required.ts index ed6bdfda2..fe2f5335e 100644 --- a/packages/mcp/src/auth-required.ts +++ b/packages/mcp/src/auth-required.ts @@ -27,6 +27,11 @@ import { loginFailureNotice, loginPrompt, loginSuccessNotification } from "./mes import { log } from "./logger.js"; import { startOrReuseLoginFlow, resolveLoginTimeoutMs } from "./login.js"; import { AUTH_REQUIRED_INSTRUCTIONS } from "./instructions.js"; +import { + SECRET_EXCLUSION_RULES, + SECRET_EXCLUSION_SUMMARY, + AUTO_SAVE_OPT_IN_RULE, +} from "./memory-policy.js"; import { MEMWAL_MCP_VERSION } from "./version.js"; interface RpcMessage { @@ -38,10 +43,22 @@ interface RpcMessage { error?: unknown; } +/** + * WALM-642: every write-tool description below ends with the shared policy + * block, the same text the `instructions` field and the plugin hooks carry, so + * an agent that only ever sees one of the three still gets the same rules. + * The signed-out variants get the one-line summary — they exist to keep a + * credential-less model from spamming writes, and the full block would be the + * longest thing in a list of tools that cannot run yet. + */ const SIGNED_OUT_REMEMBER = - "Save a fact to the user's Walrus Memory personal memory. Call ONLY when the user explicitly asks to remember/save something. Pass the full, detailed text — never summarize."; + "Save a fact to the user's Walrus Memory personal memory. Call ONLY when the user explicitly asks to remember/save something. Pass the full, detailed text — never summarize. " + + SECRET_EXCLUSION_SUMMARY; const SIGNED_IN_REMEMBER = - "Save a durable fact about the user or project to their Walrus Memory. Call this PROACTIVELY whenever the user states a preference, decision, constraint, correction, identity detail, or recurring workflow — even if they did not say 'remember this'. Skip one-off tasks, the current file or bug, and small talk. Pass the full statement; do not summarize. To save several facts at once, use memwal_remember_bulk instead."; + "Save a durable fact about the user or project to their Walrus Memory. Call this PROACTIVELY whenever the user states a preference, decision, constraint, correction, identity detail, or recurring workflow — even if they did not say 'remember this' — provided they have turned automatic memory on. Skip one-off tasks, the current file or bug, and small talk. Pass the full statement; do not summarize. To save several facts at once, use memwal_remember_bulk instead. By default this returns in ~1s once the relayer has accepted the job (job_id) — the Walrus write is still in flight and the fact is NOT stored yet. Do not claim it is saved. Settle it with the job-status tool this server advertises (re-list tools if you do not see one). A blob_id in the same reply means it already stored (optional wait budget). " + + AUTO_SAVE_OPT_IN_RULE + + " " + + SECRET_EXCLUSION_RULES; const SIGNED_OUT_RECALL = "Search the user's Walrus Memory for facts relevant to a query. Returns matching memories ranked by relevance."; const SIGNED_IN_RECALL = @@ -72,7 +89,8 @@ function buildToolDefinitions(proactive: boolean) { title: "Remember Multiple Facts", annotations: { readOnlyHint: false, destructiveHint: false }, description: - "Save multiple durable facts in one call. Use when you learned several distinct facts at once (onboarding details, a list of preferences, decisions from a discussion). Pass an array of complete fact statements (max 20) — do not summarize. Prefer this over repeated memwal_remember calls.", + "Save multiple durable facts in one call. Use when you learned several distinct facts at once (onboarding details, a list of preferences, decisions from a discussion). Pass an array of complete fact statements (max 20) — do not summarize. Prefer this over repeated memwal_remember calls. By default this returns in ~1s once the relayer has accepted the batch (job_ids) — the Walrus writes are still in flight and the facts are NOT stored yet. Do not claim they are saved. Settle them with the job-status tool this server advertises (re-list tools if you do not see one). A blob_id in the same reply means that fact already stored (optional wait budget). " + + (proactive ? AUTO_SAVE_OPT_IN_RULE + " " + SECRET_EXCLUSION_RULES : SECRET_EXCLUSION_SUMMARY), inputSchema: { type: "object", properties: { @@ -88,6 +106,35 @@ function buildToolDefinitions(proactive: boolean) { additionalProperties: false, }, }, + { + name: "memwal_remember_status", + title: "Check a Remember Job", + annotations: { readOnlyHint: true, destructiveHint: false }, + description: + "Check whether in-flight Walrus Memory writes have landed. Call this with the job_id memwal_remember returned, or job_ids from memwal_remember_bulk, when the write was reported NOT saved yet. Returns the blob_id once stored, reports that it is still uploading (call again with the ids still listed), or reports that it failed \u2014 in which case the fact was never stored and you should send it again. A batch can come back mixed, so read every line before telling the user anything is saved.", + inputSchema: { + type: "object", + properties: { + job_id: { type: "string", minLength: 1 }, + // The sidecar takes either one id or a whole batch, and the + // pending body `memwal_remember_bulk` returns tells the agent + // to come back with `job_ids`. Advertising only `job_id` — + // required, under `additionalProperties: false` — made that + // instruction unfollowable for the whole cold-start window. + job_ids: { + type: "array", + items: { type: "string", minLength: 1 }, + minItems: 1, + maxItems: 20, + }, + waitMs: { type: "integer", minimum: 0, maximum: 45000, default: 10000 }, + }, + // Neither is required on its own; the sidecar rejects passing both + // and rejects passing neither, which JSON Schema cannot express + // here without a `oneOf` that some clients mishandle. + additionalProperties: false, + }, + }, { name: "memwal_recall", title: "Recall Memories", @@ -115,7 +162,8 @@ function buildToolDefinitions(proactive: boolean) { title: "Analyze and Remember", annotations: { readOnlyHint: false, destructiveHint: true }, description: - "Extract memorable facts from a longer passage of text (preferences, habits, biographical info, constraints) and save each as a separate Walrus Memory memory. Use this when you want MemWal's LLM to split the facts out of a transcript or notes for you; if you already know the exact facts, use memwal_remember or memwal_remember_bulk instead.", + "Extract memorable facts from a longer passage of text (preferences, habits, biographical info, constraints) and save each as a separate Walrus Memory memory. Use this when you want MemWal's LLM to split the facts out of a transcript or notes for you; if you already know the exact facts, use memwal_remember or memwal_remember_bulk instead. " + + (proactive ? AUTO_SAVE_OPT_IN_RULE + " " + SECRET_EXCLUSION_RULES : SECRET_EXCLUSION_SUMMARY), inputSchema: { type: "object", properties: { @@ -169,15 +217,55 @@ function buildToolDefinitions(proactive: boolean) { ]; } +/** Tool names the OLDEST relayer this bridge still talks to serves. + * + * The bridge ships on npm and updates itself; a relayer ships per environment + * and does not, so the bridge is routinely NEWER than the relayer it dials — + * 0.0.14-dev.0 against prod and staging on 0.0.13, which is GH #928. The + * cold-start list is served before any relayer capability is known, so every + * name in it that the dialled relayer does not serve is a tool the agent can + * be told about and then cannot call. + * + * Cold start is therefore a FLOOR, not a forecast: add a name here only once + * the tool has shipped to prod, never when it lands on dev. Newer tools reach + * the client a beat later anyway — the relayer's own `tools/list` replaces + * this one and `notifications/tools/list_changed` tells the client to re-read + * it. `memwal_remember_status` is the worked example: it exists on dev, not on + * prod/staging, and belongs here only after a prod release carries it. */ +export const BASELINE_RELAYER_TOOLS: ReadonlySet = new Set([ + "memwal_remember", + "memwal_remember_bulk", + "memwal_recall", + "memwal_analyze", + "memwal_restore", + "memwal_health", +]); + +/** Served by this process, so no relayer has to know about it. Advertising it + * at cold start is always safe. */ +const LOCALLY_SERVED_TOOLS: ReadonlySet = new Set(["memwal_login"]); + +/** Every tool this bridge can describe, baseline or not. Only the baseline + * subset is advertised at cold start; this is what the schema tests pin, so a + * newer tool's shape stays reviewed while it waits for a prod release. */ +export const ALL_TOOL_DEFINITIONS = buildToolDefinitions(true); + +/** Drop anything the oldest supported relayer would not serve. */ +function coldStartTools(proactive: boolean) { + return buildToolDefinitions(proactive).filter( + (t) => BASELINE_RELAYER_TOOLS.has(t.name) || LOCALLY_SERVED_TOOLS.has(t.name), + ); +} + /** Signed-in cold-start list (bridge). Credentials exist; the relayer session * is not up yet. Proactive wording so clients that keep the first tools/list * still save/recall without being asked. */ -export const TOOL_DEFINITIONS = buildToolDefinitions(true); +export const TOOL_DEFINITIONS = coldStartTools(true); /** Signed-out list (auth-required). No credentials, so every memory call * fails: keep conservative wording or the model will spam remember and get * a stream of auth errors. */ -export const SIGNED_OUT_TOOL_DEFINITIONS = buildToolDefinitions(false); +export const SIGNED_OUT_TOOL_DEFINITIONS = coldStartTools(false); /** How long to wait for the local listener to bind + emit its URL before we * give up and return an error. Should be near-instant; 5s is paranoia. */ @@ -200,6 +288,18 @@ const LOGIN_INSTRUCTION = [ "`add_delegate_key` transaction. Credentials land at `~/.memwal/credentials.json`.", ].join("\n"); +/** Replaces {@link LOGIN_INSTRUCTION} after a failed attempt, which + * {@link loginFailureNotice} has just described. The generic copy promises "no + * client restart" and leads with `memwal_login`; for a key the user already + * approved, a restart is the only thing that recovers it and signing in again + * cannot register it twice. So the retry is offered only for the case it + * actually fixes. */ +const LOGIN_RETRY_INSTRUCTION = [ + "If you did not approve the wallet step, start a new sign-in: call the `memwal_login`", + "tool from this client, or run `npx -y @mysten-incubation/memwal-mcp login`. Open the", + "new link straight away and leave this client running through the wallet prompt.", +].join("\n"); + /** Set when a background `memwal_login` ends without credentials. The tool call * already returned the URL by then, so this is the only place left to say so. */ let lastLoginFailure: string | null = null; @@ -473,7 +573,12 @@ function handleAuthLine( id, result: { content: [ - { type: "text", text: `${loginFailureNotice(lastLoginFailure)}${LOGIN_INSTRUCTION}` }, + { + type: "text", + text: lastLoginFailure + ? `${loginFailureNotice(lastLoginFailure)}${LOGIN_RETRY_INSTRUCTION}` + : LOGIN_INSTRUCTION, + }, ], isError: true, }, diff --git a/packages/mcp/src/auth.ts b/packages/mcp/src/auth.ts index be9d3bf73..03e03da2a 100644 --- a/packages/mcp/src/auth.ts +++ b/packages/mcp/src/auth.ts @@ -10,8 +10,8 @@ * documentation patterns transfer cleanly. */ import { homedir } from "node:os"; -import { randomUUID } from "node:crypto"; -import { join, dirname, basename } from "node:path"; +import { createHash, randomUUID } from "node:crypto"; +import { join, dirname, basename, isAbsolute, relative } from "node:path"; import { mkdirSync, readFileSync, @@ -19,7 +19,9 @@ import { renameSync, unlinkSync, existsSync, + realpathSync, } from "node:fs"; +import { log } from "./logger.js"; export interface MemWalCredentials { /** 64-hex Ed25519 private key seed (32 bytes). NEVER log this. */ @@ -92,36 +94,352 @@ function projectCredsPath(): string | null { } } +/* ------------------------------------------------------------------------- * + * Project credentials are opt-IN, per machine (WALM-639). + * + * Presence alone used to be the opt-in: a `.memwal/credentials.json` anywhere + * at or above the working directory simply won. But that file is INSIDE the + * repository, so anyone who can commit to a repo — or persuade someone to + * clone one — could choose the account and the relayer every memory written + * from that directory goes to. Opening a project silently repointed the + * destination, and the writes land on immutable storage with no delete path. + * + * So a project file is now inert until the user approves that exact file, + * account, delegate and relayer. The approval record lives beside the GLOBAL + * credentials, never in the repository, because a record a repository can + * carry is a repository approving itself. + * ------------------------------------------------------------------------- */ + +const APPROVALS_FILE = "project-approvals.json"; + +/** + * Refusal of a `MEMWAL_CREDS_DIR` that cannot be trusted to sit outside a + * repository. Thrown, never swallowed — see {@link credsDirOverride}. + */ +export class UntrustedCredsDirError extends Error { + constructor(message: string) { + super(message); + this.name = "UntrustedCredsDirError"; + } +} + +/** True when `path` is `root` itself or something underneath it. String + * prefixes get `/repo` vs `/repo-2` wrong, so this asks `relative` instead. */ +function isInside(root: string, path: string): boolean { + const rel = relative(root, path); + return rel === "" || (!rel.startsWith("..") && !isAbsolute(rel)); +} + +/** + * Canonical form of a path that need not exist yet. + * + * `realpathSync` throws on a missing leaf, and the override names a directory + * this process may be about to create — so resolve the deepest ancestor that + * does exist and re-append the rest. Without this, comparing it against the + * project root is decided by which spelling of a symlinked path each side + * happened to use: on macOS `/tmp/x` and `/private/tmp/x` are the same + * directory, and a containment check that says otherwise is a check an + * attacker picks the spelling to defeat. + */ +function canonicalDir(path: string): string { + const tail: string[] = []; + let dir = path; + for (;;) { + if (existsSync(dir)) return join(realpathSync(dir), ...tail.reverse()); + const parent = dirname(dir); + if (parent === dir) return path; + tail.push(basename(dir)); + dir = parent; + } +} + +/** + * The repository the working directory is in, or null. + * + * Bounded exactly like {@link projectCredsPath}'s walk — nearest ancestor + * carrying `.git`, stopping at the home directory or the filesystem root — so + * "the project" means the same thing to the override check as it does to + * resolution. + */ +function projectRoot(): string | null { + const home = homedir(); + let dir = process.cwd(); + for (;;) { + if (dir === home) return null; + if (existsSync(join(dir, ".git"))) return dir; + const parent = dirname(dir); + if (parent === dir) return null; + dir = parent; + } +} + /** - * Which credentials file this process should read and write. + * `MEMWAL_CREDS_DIR`, but only when it can be trusted — and the same answer for + * every reader of it. + * + * The override decides the credentials outright and skips the approval gate + * entirely, which makes it exactly as trustworthy as whoever set the + * environment. That is the right trade for the thing it exists for — a shell, a + * test harness, CI — and the wrong one for the way it is reachable: an MCP + * client reads `.cursor/mcp.json`, `.vscode/mcp.json` or `.claude/settings.json` + * OUT OF THE CHECKOUT and hands this process the `env` block it finds there. A + * committed `env` block is a repository setting its own override, which is the + * thing WALM-639 exists to stop. + * + * So two shapes are refused: * - * The nearest project-local `.memwal/credentials.json` at or above the working - * directory wins over the global one, the way `.npmrc` and `.git/config` - * resolve. Signing in from one project otherwise repoints every other project - * on the machine at a different account and delegate key, silently — memories - * then land on the wrong account, on immutable storage, with no delete path - * (GH #628). + * - **Relative.** `join("", CREDS_FILE)` is the bare string + * `"credentials.json"` and `join(".memwal", CREDS_FILE)` is repo-local: + * both resolve against `process.cwd()`. That moves the credentials AND the + * approval store into whatever directory the client started in — usually + * the project root, so the store that exists precisely so a repository + * cannot carry its own approval becomes a file the repository carries. + * - **Absolute, but inside the current project.** Editors expand + * `${workspaceFolder}` inside those same config files, so "absolute" is not + * by itself evidence a human typed it. An override pointing into the + * checkout contradicts the one property this escape hatch is justified by. * - * Presence-based on purpose: creating the local file is the opt-in, so this is - * purely additive. Resolved per call rather than at module load, because the - * working directory is not knowable at import time. + * Refusing is loud on purpose, and matches the launcher's treatment of a + * relative runtime directory: falling back to the global directory in silence + * would hide a misconfigured — or hostile — client config, and leave the user + * to work out on their own why their override did nothing. * - * `MEMWAL_CREDS_DIR` overrides both project and global resolution when set, - * and is re-read on every call. + * An EMPTY value is the one thing that is not an error. It names no directory + * at all, which is what "unset" means, and it is what every other reader of + * this variable already treats it as; the bug was never that empty meant unset, + * it was that `??` let empty mean "the working directory". */ +function credsDirOverride(): string | null { + const raw = process.env.MEMWAL_CREDS_DIR; + if (!raw) return null; + + if (!isAbsolute(raw)) { + throw new UntrustedCredsDirError( + `MEMWAL_CREDS_DIR is a relative path (${JSON.stringify(raw)}). It resolves against ` + + `the working directory, which would put your credentials — and the record of ` + + `which project credentials you have approved — inside whatever directory this ` + + `process was started in. Set it to an absolute path, or unset it. ` + + `If you did not set it yourself, look at the \`env\` block of this project's MCP ` + + `config (.cursor/mcp.json, .vscode/mcp.json, .claude/settings.json): a value a ` + + `repository carries is the repository choosing where your memories go.`, + ); + } + + const root = projectRoot() ?? process.cwd(); + const dir = canonicalDir(raw); + if (isInside(canonicalDir(root), dir)) { + throw new UntrustedCredsDirError( + `MEMWAL_CREDS_DIR (${raw}) points inside the current project (${root}). The ` + + `override skips the project-approval gate, so it has to name a directory no ` + + `repository can write — and an editor expands \`\${workspaceFolder}\` in a ` + + `committed MCP config, so an absolute path is not by itself proof that you ` + + `chose it. Point it outside the project, or unset it and run ` + + `\`memwal-mcp approve-project\` if you did mean to use this project's credentials.`, + ); + } + return raw; +} + +/** + * Where records that a repository must not be able to write are kept. + * + * `MEMWAL_CREDS_DIR` is the trusted escape hatch — when it is set it decides + * the credentials outright and project resolution never runs — so following it + * here keeps a sandboxed run (tests, CI) from reaching into the real + * `~/.memwal`, exactly as #705 required for the credentials file itself. It is + * read through {@link credsDirOverride} so that "trusted" means the same thing + * here, in {@link resolveCreds} and in {@link approveProjectCreds}: one rule, + * checked once, rather than three checks that can drift apart. + */ +function trustedStateDir(): string { + return credsDirOverride() ?? join(homedir(), ".memwal"); +} + +/** The approval store. Outside every repository, on purpose. */ +export function projectApprovalsPath(): string { + return join(trustedStateDir(), APPROVALS_FILE); +} + +/** + * What approval is granted against: the destination, not the file's bytes. + * + * Account, delegate and relayer are the three fields that decide WHERE a + * memory ends up and WHO signs for it. Hashing them means a project file may + * be re-saved, relabelled or reformatted freely, while any edit that moves the + * destination invalidates the approval and has to be approved again. The + * delegate private key is deliberately NOT part of it — it must never be read + * into a record that gets written back out. + */ +export function credentialsFingerprint(creds: { + accountId: string; + delegateAddress: string; + relayerUrl: string; +}): string { + return createHash("sha256") + .update(`${creds.accountId}\n${creds.delegateAddress}\n${creds.relayerUrl}`) + .digest("hex"); +} + +/** One approved project credentials file. Contains no secret. */ +export interface ProjectApproval { + /** Canonical path of the approved `.memwal/credentials.json`. */ + path: string; + fingerprint: string; + accountId: string; + delegateAddress: string; + relayerUrl: string; + approvedAt: string; +} + +interface ApprovalsFile { + version: 1; + approvals: ProjectApproval[]; +} + +/** Compare paths the way the filesystem does. `process.cwd()` reports a + * resolved path and an approval may have been recorded through a symlink (or + * on macOS, `/tmp` → `/private/tmp`), so both sides go through this. */ +function canonicalPath(path: string): string { + try { + return realpathSync(path); + } catch { + return path; + } +} + +function isValidApproval(obj: unknown): obj is ProjectApproval { + if (!obj || typeof obj !== "object") return false; + const a = obj as Record; + return ( + typeof a.path === "string" && + typeof a.fingerprint === "string" && + typeof a.accountId === "string" && + typeof a.delegateAddress === "string" && + typeof a.relayerUrl === "string" + ); +} + +/** Approvals on record. A missing, malformed or unreadable store approves + * nothing — the safe direction, since the consequence is falling back to the + * user's own global account rather than adopting someone else's. */ +function loadApprovals(): ProjectApproval[] { + const path = projectApprovalsPath(); + if (!existsSync(path)) return []; + try { + const parsed = JSON.parse(readFileSync(path, "utf8")) as ApprovalsFile; + if (!parsed || parsed.version !== 1 || !Array.isArray(parsed.approvals)) return []; + return parsed.approvals.filter(isValidApproval); + } catch { + return []; + } +} + +/** Through the same writer as the credentials file: it creates the directory + * at `0700` and the file at `0600`. The record holds no secret, but it decides + * where memories go, so it should not be writable by anything that could not + * already write the credentials beside it. */ +function saveApprovals(approvals: ProjectApproval[]): void { + writeSecretFile( + projectApprovalsPath(), + JSON.stringify({ version: 1, approvals } satisfies ApprovalsFile, null, 2), + ); +} + +/** Why a project credentials file was, or was not, used. */ +export type ProjectCredsDecision = + /** Approved for exactly this account + delegate + relayer: in use. */ + | "approved" + /** Never approved on this machine. Ignored. */ + | "unapproved" + /** Approved once, but the destination has since changed. Ignored. */ + | "changed" + /** Present but not valid credentials, so there is nothing to approve. */ + | "unreadable"; + +export interface ProjectCredsInfo { + /** The project-local file that was found. */ + path: string; + decision: ProjectCredsDecision; + /** Destination it points at. Absent when the file could not be read — and + * never the delegate private key, which no caller of this ever needs. */ + accountId?: string; + relayerUrl?: string; +} + +/** Which file won, and what happened to any project-local file that did not. */ +export interface CredsResolution { + /** The file this process reads and writes. */ + path: string; + source: "override" | "project" | "global"; + /** The project-local file found by the walk, if any — present whether or + * not it was used, so callers can report one they ignored. */ + project?: ProjectCredsInfo; +} + +/** + * Which credentials file this process should read and write, and why. + * + * `MEMWAL_CREDS_DIR` wins outright, and skips the approval gate — but only a + * value {@link credsDirOverride} accepts, i.e. a non-empty ABSOLUTE path + * outside the current project. It is "trusted" exactly as far as whoever set + * the environment is, and an MCP client will hand this process an `env` block + * it read out of the checkout, so a value that could have been committed is + * refused rather than obeyed. Otherwise the nearest project-local + * `.memwal/credentials.json` at or above the working directory is used IF the + * user has approved that exact destination on this machine (WALM-639), and the + * global file is used in every other case — including an unapproved, altered + * or malformed project file. Falling back rather than failing keeps a machine + * that has never seen a project file behaving exactly as it always did. + * + * Resolved per call rather than at module load, because neither the working + * directory nor the approval store is knowable at import time. + */ +export function resolveCreds(): CredsResolution { + const override = credsDirOverride(); + if (override) return { path: join(override, CREDS_FILE), source: "override" }; + + const global = globalCredsPath(); + const projectPath = projectCredsPath(); + if (!projectPath) return { path: global, source: "global" }; + + const project = readCredsFile(projectPath); + if (!project) { + return { + path: global, + source: "global", + project: { path: projectPath, decision: "unreadable" }, + }; + } + + const approval = loadApprovals().find((a) => a.path === canonicalPath(projectPath)); + const decision: ProjectCredsDecision = !approval + ? "unapproved" + : approval.fingerprint === credentialsFingerprint(project) + ? "approved" + : "changed"; + const info: ProjectCredsInfo = { + path: projectPath, + decision, + accountId: project.accountId, + relayerUrl: project.relayerUrl, + }; + return decision === "approved" + ? { path: projectPath, source: "project", project: info } + : { path: global, source: "global", project: info }; +} + +/** The credentials file in use. Thin wrapper over {@link resolveCreds} so the + * many callers that only need a path are unchanged. */ export function credsPath(): string { - const override = process.env.MEMWAL_CREDS_DIR; - if (override) return join(override, CREDS_FILE); - return projectCredsPath() ?? globalCredsPath(); + return resolveCreds().path; } -/** Load credentials from disk. Returns null if missing or malformed. */ -export function loadCreds(): MemWalCredentials | null { - const path = credsPath(); +/** Read and validate one credentials file. Returns null if missing or + * malformed — the caller decides what that means. */ +function readCredsFile(path: string): MemWalCredentials | null { if (!existsSync(path)) return null; try { - const raw = readFileSync(path, "utf8"); - const parsed = JSON.parse(raw); + const parsed = JSON.parse(readFileSync(path, "utf8")); if (!isValid(parsed)) return null; return parsed as MemWalCredentials; } catch { @@ -129,6 +447,186 @@ export function loadCreds(): MemWalCredentials | null { } } +/** Load credentials from disk. Returns null if missing or malformed. */ +export function loadCreds(): MemWalCredentials | null { + return readCredsFile(credsPath()); +} + +/** What {@link approveProjectCreds} did. */ +export interface ApproveProjectResult { + outcome: + /** Newly approved. */ + | "approved" + /** Approved again after the destination changed. */ + | "reapproved" + /** Already approved for this exact destination; nothing written. */ + | "already-approved" + /** No project-local credentials file at or above the working directory. */ + | "none" + /** A project file exists but is not valid credentials. */ + | "unreadable" + /** `MEMWAL_CREDS_DIR` is set, so project resolution never runs. */ + | "overridden"; + projectPath?: string; + accountId?: string; + relayerUrl?: string; + /** Destination the previous approval covered, when this replaced one. */ + previousAccountId?: string; + previousRelayerUrl?: string; + approvalsPath: string; +} + +/** + * Approve the project-local credentials found from the working directory. + * + * Deliberately takes no arguments: it approves what resolution would otherwise + * ignore, from the same directory, so "what am I approving" and "what will be + * used" cannot drift apart. + */ +export function approveProjectCreds(): ApproveProjectResult { + const approvalsPath = projectApprovalsPath(); + if (credsDirOverride()) return { outcome: "overridden", approvalsPath }; + + const projectPath = projectCredsPath(); + if (!projectPath) return { outcome: "none", approvalsPath }; + const creds = readCredsFile(projectPath); + if (!creds) return { outcome: "unreadable", projectPath, approvalsPath }; + + const key = canonicalPath(projectPath); + const fingerprint = credentialsFingerprint(creds); + const approvals = loadApprovals(); + const existing = approvals.find((a) => a.path === key); + if (existing?.fingerprint === fingerprint) { + return { + outcome: "already-approved", + projectPath, + accountId: creds.accountId, + relayerUrl: creds.relayerUrl, + approvalsPath, + }; + } + + saveApprovals([ + ...approvals.filter((a) => a.path !== key), + { + path: key, + fingerprint, + accountId: creds.accountId, + delegateAddress: creds.delegateAddress, + relayerUrl: creds.relayerUrl, + approvedAt: new Date().toISOString(), + }, + ]); + return { + outcome: existing ? "reapproved" : "approved", + projectPath, + accountId: creds.accountId, + relayerUrl: creds.relayerUrl, + previousAccountId: existing?.accountId, + previousRelayerUrl: existing?.relayerUrl, + approvalsPath, + }; +} + +/** What {@link revokeProjectCredsApproval} did. */ +export interface RevokeProjectResult { + outcome: "revoked" | "none"; + projectPath?: string; + approvalsPath: string; +} + +/** + * Withdraw the approval for the project-local credentials here. + * + * Keyed on the path rather than on the file's current contents, so an approval + * can be withdrawn even after the file it covered was edited or deleted — a + * revoke that only worked while the destination still matched would be + * useless exactly when it is wanted. + */ +export function revokeProjectCredsApproval(): RevokeProjectResult { + const approvalsPath = projectApprovalsPath(); + const projectPath = projectCredsPath() ?? join(process.cwd(), ".memwal", CREDS_FILE); + const key = canonicalPath(projectPath); + const approvals = loadApprovals(); + const remaining = approvals.filter((a) => a.path !== key); + if (remaining.length === approvals.length) return { outcome: "none", projectPath, approvalsPath }; + saveApprovals(remaining); + return { outcome: "revoked", projectPath, approvalsPath }; +} + +/** + * The warning for a project credentials file that was found and NOT used, or + * null when there is nothing to report. + * + * Says which file was ignored, where memory is going instead, and the exact + * command that approves it — a silent fallback would be the mirror image of + * the silent redirect this gate exists to stop. Never contains a key: the only + * fields it reads are the account id and the relayer URL. + */ +export function formatProjectCredsNotice( + resolution: CredsResolution = resolveCreds(), +): string | null { + const project = resolution.project; + if (!project || project.decision === "approved") return null; + + const destination = `account ${project.accountId} on ${project.relayerUrl}`; + const head = + project.decision === "unreadable" + ? [ + `Ignored the project credentials at ${project.path}: the file is not a valid`, + `Walrus Memory credentials file, so there is nothing to approve.`, + ] + : project.decision === "changed" + ? [ + `Ignored the project credentials at ${project.path}: they changed since you`, + `approved them and now point at ${destination}.`, + `Approving again is required whenever the account, delegate key or relayer moves.`, + ] + : [ + `Ignored the project credentials at ${project.path}, which would send memory to`, + `${destination}.`, + `A file inside a repository can be committed by anyone, so it is not used until`, + `you approve it on this machine.`, + ]; + + const lines = [...head, `Memory is going to ${resolution.path} instead.`]; + if (project.decision !== "unreadable") { + lines.push( + `To use it, run \`memwal-mcp approve-project\` in a terminal from this directory.`, + `The approval is recorded in ${projectApprovalsPath()}, outside the repository.`, + ); + } + return lines.join("\n"); +} + +/** + * What approving a project credentials file costs, beyond redirecting reads. + * + * Approval does not just decide which file is READ — it decides which file is + * WRITTEN. `saveCreds` writes to `credsPath()`, so from here on every sign-in + * from this project puts a fresh 64-hex Ed25519 delegate seed, in plaintext, + * into a directory that is inside the repository. An attacker who planted the + * file planted the directory too, so it is already tracked rather than ignored, + * and `git add -A` stages the key. The user cannot weigh that if nobody says + * it, and "you approved a destination" is not the same sentence as "you + * approved storing a private key in your repo". + * + * Separate from {@link formatProjectCredsNotice} on purpose: that one is about + * a file that was IGNORED, this one is about a file that is about to be used. + */ +export function formatProjectCredsStorageWarning(projectPath: string): string { + const dir = dirname(projectPath); + return [ + `Note: ${dir} is inside the repository, and approving makes it the file this`, + `project signs and saves with. Every later sign-in from here writes a delegate`, + `PRIVATE KEY into ${basename(projectPath)} in plaintext — and a \`.memwal/\` a`, + `repository already carries is tracked, not ignored, so \`git add -A\` would stage it.`, + `Add \`.memwal/\` to .gitignore, and never commit ${basename(projectPath)}.`, + `(The short-lived login write-ahead record is kept outside the repository; the`, + `credentials file cannot be, because it is the file you approved.)`, + ].join("\n"); +} + /** * Write credentials with secure (`0600`) permission, to whichever file * `credsPath()` resolves to. @@ -294,15 +792,25 @@ export interface SaveCredsResult { * callback arrives, by which point a delegate key has already been registered * on-chain — so a warning that waits for both ids is a warning that arrives * too late to act on. + * + * When the file being replaced is an approved project one, that is also the + * last moment before a new private key lands inside a repository, so the + * storage warning is repeated here rather than left behind at approval time — + * approval may have been months ago, or done by someone else on the machine. */ export function formatPendingSignInWarning(): string | null { - const current = loadCreds(); + const resolution = resolveCreds(); + const current = readCredsFile(resolution.path); if (!current) return null; - return ( - `Signing in will replace the credentials in ${credsPath()} ` + - `(currently account ${current.accountId}). ` + - `The existing file is backed up if the new sign-in is a different account.` - ); + const lines = [ + `Signing in will replace the credentials in ${resolution.path} ` + + `(currently account ${current.accountId}). ` + + `The existing file is backed up if the new sign-in is a different account.`, + ]; + if (resolution.source === "project") { + lines.push(formatProjectCredsStorageWarning(resolution.path)); + } + return lines.join("\n"); } /** @@ -405,3 +913,181 @@ function isValid(obj: unknown): obj is MemWalCredentials { c.version === 1 ); } + +/* ------------------------------------------------------------------------- * + * Pending login — write-ahead for the delegate keypair (WALM-332). + * + * The browser registers our delegate public key on-chain, which costs gas and + * cannot be undone, and only afterwards POSTs the callback that makes us save + * the matching private key. Losing this process in that window used to destroy + * the only copy of the key, stranding a paid registration nobody could use. + * + * So the keypair is written here BEFORE the browser is given the connect URL, + * and cleared once `saveCreds` has the key safely in `credentials.json`. A + * record that outlives its flow is recovered on next start. + * ------------------------------------------------------------------------- */ + +const PENDING_FILE = "login-pending.json"; + +/** Where a PROJECT sign-in's write-ahead record goes, under the trusted state + * directory — one file per approved project. See {@link pendingLoginPath}. */ +const PENDING_DIR = "login-pending"; + +/** + * How long a stranded record stays recoverable. + * + * Deliberately far longer than the 5-minute login timeout: the whole point is + * to survive a client restart, and a user who quits for the evening and comes + * back tomorrow is exactly the case worth covering. The cost of holding it is + * an unregistered key on disk, which grants nothing. + */ +export const PENDING_LOGIN_TTL_MS = 24 * 60 * 60_000; + +export interface PendingLogin { + /** 64-hex Ed25519 private key seed. NEVER log this. */ + delegatePrivateKey: string; + delegatePublicKeyHex: string; + delegateAddress: string; + /** Relayer the flow was started against — recovery must not repoint. */ + relayerUrl: string; + label?: string; + /** ISO timestamp, for TTL expiry. */ + createdAt: string; + version: 1; +} + +/** + * Where the write-ahead record lives — never inside a repository. + * + * It used to sit beside whichever credentials file `credsPath()` resolved to, + * which reads as "recover into the project you signed in from" and is the right + * intent. But for an approved project that path is `/.memwal/`, so every + * sign-in dropped a plaintext 64-hex delegate seed into the working tree — + * under a directory an attacker who planted the credentials file had already + * created, so tracked rather than ignored, and staged by `git add -A`. + * + * Nothing about this record needs to be in the repo. It is short-lived + * handshake state, nobody edits it by hand, and the only property that matters + * is that the project which started a sign-in is the one that can reclaim it. + * So it moves to the trusted state directory, keyed by a hash of the + * credentials path: same per-project scoping, no key material in the checkout. + * + * The non-project cases are byte-for-byte where they always were — + * `dirname(credsPath())` IS `trustedStateDir()` for both the override and the + * global file — so a machine that has never approved a project sees no change. + */ +export function pendingLoginPath(): string { + const resolution = resolveCreds(); + if (resolution.source !== "project") return join(trustedStateDir(), PENDING_FILE); + const key = createHash("sha256").update(canonicalPath(resolution.path)).digest("hex"); + return join(trustedStateDir(), PENDING_DIR, `${key}.json`); +} + +/** + * Persist the pending keypair. Throws if it cannot. + * + * Deliberately NOT best-effort. The invariant this record exists to hold is + * that the delegate private key is on disk before its public half can reach a + * browser that will pay gas to register it. Swallowing the error would publish + * the connect URL while claiming a durability that does not exist — the + * original WALM-332 loss, now silent. + * + * Failing the login costs the user nothing: this file lives in the same + * trusted state directory the approval record does, so a directory that cannot + * take it is one this process cannot keep state in at all. The same login would + * have failed at the callback anyway, one on-chain `add_delegate_key` later. + */ +export function savePendingLogin(pending: PendingLogin): void { + const path = pendingLoginPath(); + try { + // The record holds the same plaintext private key as `credentials.json`, + // so it gets the same fresh-inode write. + writeSecretFile(path, JSON.stringify(pending, null, 2)); + } catch (err) { + const msg = err instanceof Error ? err.message : String(err); + log.error("login.pending.write_failed", { path, msg }); + throw new Error( + `Could not write the login write-ahead record at ${path}: ${msg}. ` + + `Refusing to start a sign-in that could register a delegate key on-chain ` + + `without being able to save it.`, + ); + } +} + +/** + * A pending record this login can adopt instead of minting a new keypair. + * + * `loginFlow` used to generate a fresh keypair every call and overwrite the + * record unconditionally. Recovery only runs at process start and is skipped + * for `--login` / `forceLogin`, so a timed-out login followed by `memwal_login` + * in the same process replaced the only copy of a key the browser may already + * have paid to register. Reusing the record keeps that key reclaimable. + * + * Scoped to the same relayer: a key registered against one relayer's account + * proves nothing to another, and `recovery` must never repoint. TTL and shape + * are already enforced by {@link loadPendingLogin}. + */ +export function reusablePendingLogin(relayerUrl: string): PendingLogin | null { + const pending = loadPendingLogin(); + if (!pending) return null; + if (pending.relayerUrl !== relayerUrl) { + log.warn("login.pending.relayer_changed", { + publicKey: pending.delegatePublicKeyHex, + from: pending.relayerUrl, + to: relayerUrl, + }); + return null; + } + return pending; +} + +/** + * Load a pending record, or null if there is none, it is malformed, or it has + * aged out. An expired record is deleted on read rather than left to linger. + */ +export function loadPendingLogin(): PendingLogin | null { + const path = pendingLoginPath(); + if (!existsSync(path)) return null; + let parsed: unknown; + try { + parsed = JSON.parse(readFileSync(path, "utf8")); + } catch { + clearPendingLogin(); + return null; + } + if (!isValidPending(parsed)) { + clearPendingLogin(); + return null; + } + const age = Date.now() - Date.parse(parsed.createdAt); + if (!Number.isFinite(age) || age > PENDING_LOGIN_TTL_MS) { + clearPendingLogin(); + return null; + } + return parsed; +} + +/** Remove the pending record. Safe to call when there isn't one. */ +export function clearPendingLogin(): void { + try { + const path = pendingLoginPath(); + if (existsSync(path)) unlinkSync(path); + } catch { + /* best effort */ + } +} + +function isValidPending(obj: unknown): obj is PendingLogin { + if (!obj || typeof obj !== "object") return false; + const p = obj as Record; + return ( + typeof p.delegatePrivateKey === "string" && + /^[0-9a-fA-F]{64}$/.test(p.delegatePrivateKey) && + typeof p.delegatePublicKeyHex === "string" && + /^[0-9a-fA-F]{64}$/.test(p.delegatePublicKeyHex) && + typeof p.delegateAddress === "string" && + typeof p.relayerUrl === "string" && + typeof p.createdAt === "string" && + p.version === 1 + ); +} diff --git a/packages/mcp/src/auto-save.ts b/packages/mcp/src/auto-save.ts new file mode 100644 index 000000000..afa05b29c --- /dev/null +++ b/packages/mcp/src/auto-save.ts @@ -0,0 +1,353 @@ +/** + * Automatic-save consent (WALM-642). + * + * Saving a fact the user explicitly asked for has never needed permission and + * still does not. What this module governs is the other thing the package does: + * saving what it judges durable, without being asked. That is what MemWal is + * for — so it is ON — but a person has to have been told what it means first. + * + * ── Three states, not two ─────────────────────────────────────────────────── + * `off` and `on` are answers. `unset` is the absence of one, and it is not the + * same as `off`: + * + * - **Answered** (`autoSave: true|false` in settings.json) — from the login + * prompt or from `memwal-mcp auto-save on|off`. Nothing re-asks. + * - **Unset on an install that predates this change** — no settings file at + * all, but credentials on disk. These users have been auto-saving all + * along; switching them off would be a regression dressed up as caution, so + * saving CONTINUES and the question is put to them at their next + * interactive login. + * - **Unset on a new install** (`autoSaveConsent: "pending"`, stamped the + * moment a post-change install first appears) — nobody has been asked, so + * nothing is saved unprompted until someone answers. Explicit tool calls + * and recall keep working the whole time. + * + * The stamp is what separates the last two. Without it a headless install could + * sign in through the `memwal_login` tool, look indistinguishable from a + * long-standing user, and start saving on its own — consent by never having + * been asked. + * + * ── Where the answer lives: the trusted state dir, and nowhere else ───────── + * The answer is a decision about the HUMAN, not about a credentials directory. + * It used to be stored beside whichever `credentials.json` won resolution, + * which had two consequences, both wrong: + * + * 1. Approving a project's credentials (WALM-639) moved `settingsPath()` + * into the repository. A user who had already answered "off" globally got + * a directory with no answer in it, fell through to the legacy + * "credentials exist, so keep saving" rule, and had automatic memory + * switched back ON by an approval that was only ever about where writes + * go. The recorded "no" was never consulted there. + * 2. `setAutoSave` / `markConsentPending` then wrote the consent answer INTO + * the repository, where it could be committed and shipped to everyone + * else who cloned it. + * + * So consent now lives in the trusted state dir — `MEMWAL_CREDS_DIR` when it is + * set, else `~/.memwal` — the same directory auth.ts keeps the project-approval + * store in, and for the same reason: a record a repository can write is a + * repository approving itself. Project scoping of CREDENTIALS is unchanged; + * project scoping of CONSENT is gone. + * + * The legacy probe reads the trusted directory's own `credentials.json` for the + * same reason. "Has this person used MemWal before" must not be answerable by a + * file inside a checkout. + * + * ── How the hooks see it ─────────────────────────────────────────────────── + * The plugin's hooks are separate processes: a hook is spawned by the client, + * not by this package, and inherits none of the MCP server's `env` or argv. The + * hooks used to re-derive the whole resolution in plain ESM, and the two + * implementations drifted — the hook side kept the pre-WALM-639 presence rule, + * so a committed `.memwal/credentials.json` (contents never parsed, only + * `existsSync`) turned automatic memory on for anyone who opened the repo. + * + * Re-deriving is the bug, so the hooks no longer do it. This module publishes + * the RESOLVED state to `auto-save-state.json` in the trusted state dir, and + * the hooks read only that. A repository cannot write there, the hooks do no + * resolution of their own, and a hook that cannot read the file fails safe + * (automatic memory off). See `plugin/scripts/lib/auto-save.mjs`. + */ +import { dirname, join } from "node:path"; +import { + chmodSync, + existsSync, + mkdirSync, + readFileSync, + writeFileSync, +} from "node:fs"; +import { projectApprovalsPath } from "./auth.js"; +import { log } from "./logger.js"; + +/** Env var that pins automatic saving on or off for one MCP server process. */ +export const AUTO_SAVE_ENV = "MEMWAL_AUTO_SAVE"; + +const SETTINGS_FILE = "settings.json"; +const CREDS_FILE = "credentials.json"; +/** Resolved state the plugin hooks read. Never written inside a repository. */ +const HOOK_STATE_FILE = "auto-save-state.json"; +/** Bumped if the hook-facing shape ever changes incompatibly. */ +const HOOK_STATE_VERSION = 1; + +/** + * Where records a repository must not be able to write are kept: + * `MEMWAL_CREDS_DIR` when set, else `~/.memwal`. + * + * Derived from `projectApprovalsPath()` rather than recomputed, so this file + * and auth.ts cannot drift apart the way the hook copy did. auth.ts owns the + * definition; this reads it back. + */ +function trustedStateDir(): string { + return dirname(projectApprovalsPath()); +} + +/** Where the answer is stored. Outside every repository, on purpose. */ +export function settingsPath(): string { + return join(trustedStateDir(), SETTINGS_FILE); +} + +/** Where the resolved state is published for the plugin hooks to read. */ +export function hookStatePath(): string { + return join(trustedStateDir(), HOOK_STATE_FILE); +} + +/** Shape of `settings.json`. Unknown keys are preserved on write. */ +interface MemWalSettings { + /** Present only once the question has actually been answered. */ + autoSave?: boolean; + /** `"pending"` marks an install created after WALM-642 with no answer yet. */ + autoSaveConsent?: string; + [key: string]: unknown; +} + +/** + * Parse a human-written boolean. Returns null for "not set" and for anything + * unparseable — an unreadable value must not be read as an answer. + */ +export function parseBooleanSetting(raw: string | undefined | null): boolean | null { + if (raw === undefined || raw === null) return null; + const v = raw.trim().toLowerCase(); + if (v === "") return null; + if (["1", "true", "on", "yes", "y", "enable", "enabled"].includes(v)) return true; + if (["0", "false", "off", "no", "n", "disable", "disabled"].includes(v)) return false; + return null; +} + +function readSettings(): MemWalSettings { + const path = settingsPath(); + if (!existsSync(path)) return {}; + try { + const parsed = JSON.parse(readFileSync(path, "utf8")); + return parsed && typeof parsed === "object" ? (parsed as MemWalSettings) : {}; + } catch { + // A corrupt settings file is not an answer. Fall through to unset. + return {}; + } +} + +function writeSettings(next: MemWalSettings): string { + const path = settingsPath(); + mkdirSync(dirname(path), { recursive: true, mode: 0o700 }); + writeFileSync(path, `${JSON.stringify(next, null, 2)}\n`, { mode: 0o600 }); + // `writeFileSync`'s `mode` follows POSIX `open()`: the kernel applies it + // when it CREATES the inode and ignores it for one that already exists, so + // a file left permissive by anything else would keep its old mode forever. + // Unlike credentials.json (see `writeSecretFile` in auth.ts) there is no + // secret in flight here, so tightening afterwards is enough — no reader can + // learn anything from the window, only whether the flag is set. + chmodSync(path, 0o600); + return path; +} + +/** `on`/`off` are answers; `unset` means nobody has been asked yet. */ +export type AutoSaveState = "on" | "off" | "unset"; + +export type AutoSaveSource = + /** `MEMWAL_AUTO_SAVE` in this process's environment. */ + | "env" + /** An answer written to settings.json. */ + | "settings" + /** Unset, on an install that predates the consent prompt — keeps saving. */ + | "legacy" + /** Unset, on an install created after it — saves nothing until answered. */ + | "unanswered"; + +export interface AutoSaveStatus { + /** Whether this process may save without being asked. */ + enabled: boolean; + /** The stored answer, or `unset`. */ + state: AutoSaveState; + source: AutoSaveSource; + /** True while the question is still owed a human answer. */ + pendingConsent: boolean; + /** Where an answer lives (or would be written). */ + path: string; +} + +/** + * Resolve the current state: environment, then the stored answer, then the + * unset rules above. + */ +export function autoSaveStatus(): AutoSaveStatus { + const path = settingsPath(); + const settings = readSettings(); + const answered = + typeof settings.autoSave === "boolean" ? settings.autoSave : null; + const state: AutoSaveState = + answered === null ? "unset" : answered ? "on" : "off"; + + const fromEnv = parseBooleanSetting(process.env[AUTO_SAVE_ENV]); + if (fromEnv !== null) { + // Overrides the behaviour, and counts as a deliberate configuration + // act: someone who set this does not need to be asked as well. + return { enabled: fromEnv, state, source: "env", pendingConsent: false, path }; + } + + if (answered !== null) { + return { enabled: answered, state, source: "settings", pendingConsent: false, path }; + } + + // Unset. Which way it falls depends on whether this install was ever in a + // position to have been asked — see the module comment. The probe is the + // TRUSTED directory's credentials file, never a project one: "has this + // person used MemWal before" must not be answerable by a checkout. + const stamped = settings.autoSaveConsent === "pending"; + const preExisting = !stamped && existsSync(join(trustedStateDir(), CREDS_FILE)); + return { + enabled: preExisting, + state: "unset", + source: preExisting ? "legacy" : "unanswered", + pendingConsent: true, + path, + }; +} + +/** + * True when this process may save without being asked. + * + * Read at call time, never cached: the answer can be written between calls in + * the same process. + */ +export function isAutoSaveEnabled(): boolean { + return autoSaveStatus().enabled; +} + +/** True while a human still owes the consent question an answer. */ +export function isConsentPending(): boolean { + return autoSaveStatus().pendingConsent; +} + +/** The file the plugin hooks read. Contains no secret — only the answer. */ +export interface HookAutoSaveState { + version: number; + enabled: boolean; + state: AutoSaveState; + source: AutoSaveSource; + pendingConsent: boolean; + /** The settings file this was resolved from, for diagnosis only. */ + settingsPath: string; + updatedAt: string; +} + +/** + * Publish the resolved state where the plugin hooks can read it. + * + * This is the whole hook contract: the hooks do no resolution, consult no + * project directory and parse no credentials file — they read this one file out + * of the trusted state dir, or they fail safe. Called at every point the answer + * can change (`setAutoSave`, `markConsentPending`) and once at start-up, so a + * hook spawned in the same session as an MCP server sees the current answer. + * + * Best effort: a read-only or missing home directory must not stop the server + * from running, and a hook that finds no file already behaves as "off". + */ +export function publishHookState(): HookAutoSaveState | null { + const status = autoSaveStatus(); + const payload: HookAutoSaveState = { + version: HOOK_STATE_VERSION, + enabled: status.enabled, + state: status.state, + source: status.source, + pendingConsent: status.pendingConsent, + settingsPath: status.path, + updatedAt: new Date().toISOString(), + }; + const path = hookStatePath(); + try { + mkdirSync(dirname(path), { recursive: true, mode: 0o700 }); + writeFileSync(path, `${JSON.stringify(payload, null, 2)}\n`, { mode: 0o600 }); + chmodSync(path, 0o600); + return payload; + } catch (err) { + log.warn("autosave.publish_failed", { + path, + error: err instanceof Error ? err.message : String(err), + }); + return null; + } +} + +/** + * Record the answer, preserving any other keys already in the file. Clears the + * pending stamp — the question has been answered and must not be asked again. + */ +export function setAutoSave(enabled: boolean): { path: string; enabled: boolean } { + const next: MemWalSettings = { ...readSettings(), autoSave: enabled }; + delete next.autoSaveConsent; + const path = writeSettings(next); + publishHookState(); + log.info("autosave.set", { enabled, path }); + return { path, enabled }; +} + +/** + * Mark this install as created after the consent prompt existed, so an unset + * state here is read as "never asked" rather than "long-standing user". + * + * Called once when a brand-new install first appears — before the signed-out + * server boots, and immediately after a first interactive login — so that an + * abandoned or headless sign-in cannot mature into automatic saving nobody + * agreed to. A no-op once any answer exists. + */ +export function markConsentPending(): void { + const settings = readSettings(); + if (typeof settings.autoSave === "boolean") return; + if (settings.autoSaveConsent === "pending") return; + writeSettings({ ...settings, autoSaveConsent: "pending" }); + publishHookState(); + log.info("autosave.consent_pending", { path: settingsPath() }); +} + +/** One line for a TTY, naming the state, where it came from, and how to flip it. */ +export function autoSaveSummary(): string { + const status = autoSaveStatus(); + const where = + status.source === "env" + ? `from ${AUTO_SAVE_ENV}` + : status.source === "settings" + ? `from ${status.path}` + : status.source === "legacy" + ? "carried over from before this setting existed" + : "not answered yet"; + const how = status.enabled + ? "Turn it off with `memwal-mcp auto-save off`." + : "Turn it on with `memwal-mcp auto-save on`. Facts you explicitly ask to save are stored either way."; + return `Automatic memory: ${status.enabled ? "ON" : "OFF"} (${where}). ${how}`; +} + +/** + * The single line a non-interactive run prints while consent is outstanding. + * + * Non-TTY is every MCP client spawn, so this is stderr-only and says what is + * happening rather than asking anything — there is no human on the other end of + * this stdin, and a prompt here would hang the server forever. + */ +export function pendingConsentNotice(): string | null { + const status = autoSaveStatus(); + if (!status.pendingConsent) return null; + return status.enabled + ? "Automatic memory is ON, carried over from before this setting existed. " + + "Run `memwal-mcp auto-save on|off` in a terminal to confirm or change it." + : "Automatic memory is waiting on your answer, so nothing is being saved " + + "unprompted yet (facts you explicitly ask to save still are). Run " + + "`memwal-mcp login` in a terminal to answer, or set it directly with " + + "`memwal-mcp auto-save on|off`."; +} diff --git a/packages/mcp/src/bridge.ts b/packages/mcp/src/bridge.ts index 6b9319ee3..c3883d38c 100644 --- a/packages/mcp/src/bridge.ts +++ b/packages/mcp/src/bridge.ts @@ -15,7 +15,7 @@ * Re-auth requires an explicit `memwal-mcp login` from the user. */ import type { MemWalCredentials } from "./auth.js"; -import { clearCreds, credsPath, loadCreds } from "./auth.js"; +import { clearCreds, clearPendingLogin, credsPath, loadCreds } from "./auth.js"; import { TOOL_DEFINITIONS } from "./auth-required.js"; import { clientInfoHeaders, @@ -24,7 +24,13 @@ import { } from "./client-info.js"; import { randomUUID } from "node:crypto"; import { ensureCompatibleRelayer, resolveConnectTimeoutMs } from "./compatibility.js"; -import { PROACTIVE_INSTRUCTIONS } from "./instructions.js"; +import { + describeHealthProbe, + probeRelayerHealth, + resolveHealthProbeMs, + type HealthProbe, +} from "./health-probe.js"; +import { proactiveInstructions } from "./instructions.js"; import { startOrReuseLoginFlow, resolveLoginTimeoutMs } from "./login.js"; import { log, note } from "./logger.js"; import { @@ -33,6 +39,7 @@ import { loginSuccessNotification, type LoginSuccessInfo, } from "./messages.js"; +import { openStreamableSession, resolveTransport } from "./streamable.js"; import { MEMWAL_MCP_VERSION } from "./version.js"; /** Bridge mode runtime config — the URLs / label resolved at boot from @@ -74,7 +81,8 @@ const NAMESPACE_TOOLS = new Set([ * per-call namespace always wins over the configured default. */ /** - * Name the relayer this process dialled in a `memwal_health` result. + * Name the destination this process is bound to in a `memwal_health` result: + * the relayer it dialled, and the account it signs for. * * The relayer-side text can only report an origin its deployment published, and * stays silent on a self-hosted or local one, where the sidecar knows nothing @@ -86,10 +94,18 @@ const NAMESPACE_TOOLS = new Set([ * Rewrites an existing `relayer=` field rather than appending a second one: when * both sides know the origin they describe the same session, and two * conflicting fields would be worse than neither. + * + * `account=` rides along for the same reason (WALM-639). A project-local + * credentials file can point this process at a different account than the one + * the user signed in with, and "which account am I writing to" was otherwise + * only visible in stderr the MCP host usually hides — so the half of the + * destination that decides WHOSE memory this is now shows up beside the half + * that decides where it is stored. */ export function annotateHealthResult( result: { content?: unknown; isError?: unknown }, relayerUrl: string, + accountId?: string, ): void { // A failed health call has no session to describe; naming a relayer beside // an error reads as though that relayer answered. @@ -103,6 +119,11 @@ export function annotateHealthResult( block.text = existing.test(block.text) ? block.text.replace(existing, `relayer=${relayerUrl}`) : `${block.text} relayer=${relayerUrl}`; + if (!accountId) return; + const existingAccount = /\baccount=\S+/; + block.text = existingAccount.test(block.text) + ? block.text.replace(existingAccount, `account=${accountId}`) + : `${block.text} account=${accountId}`; } export function applyDefaultNamespace(msg: RpcMessage, namespace?: string): RpcMessage { @@ -183,7 +204,11 @@ function buildLocalInitializeResult(params: unknown): { // client: this local answer wins and the upstream initialize reply is // suppressed. Omitting it here silently strips the proactive contract // from every stdio client, which is the WALM-324 regression itself. - instructions: PROACTIVE_INSTRUCTIONS, + // + // Resolved per handshake, not read from a module const: whether the + // model is told to save unprompted depends on the user's automatic-save + // opt-in, which lives on disk and can change between spawns (WALM-642). + instructions: proactiveInstructions(), }; } @@ -213,6 +238,18 @@ const SIGNED_OUT_FAILURE = { const UNAUTHORIZED_TEXT = "❌ Walrus Memory rejected the saved credentials (HTTP 401). The delegate key may have been revoked or is no longer registered on this account. Call `memwal_login` to sign in again — saved credentials were NOT modified."; +/** Reply for a `tools/call` naming a tool the connected relayer does not + * serve. Names what IS on offer, because the agent's next move is to pick one + * of those — "unknown tool" alone leaves it guessing or retrying. */ +export function unknownToolText(name: string, available: string[]): string { + return ( + `❌ \`${name}\` is not a tool this Walrus Memory server offers. ` + + `Available: ${available.join(", ")}. ` + + `Your tool list is stale — re-read \`tools/list\` and use one of those instead. ` + + `Nothing ran, so nothing was saved or changed.` + ); +} + /** `failRequest` options for every credentials-rejected refusal, so one refused * at handshake time and one refused on arrival afterwards read identically. */ const UNAUTHORIZED_FAILURE = { @@ -238,6 +275,11 @@ const LOCAL_TOOLS_LIST = { ], }; +/** Names advertised at cold start (baseline relayer + locally-served). A + * `tools/call` for anything else is refused until an upstream `tools/list` + * has been seen — fail-closed, not fail-open. */ +const COLD_START_TOOL_NAMES = new Set(LOCAL_TOOLS_LIST.tools.map((t) => t.name)); + const URL_READY_TIMEOUT_MS = 5_000; /** Maximum silence we tolerate on the SSE stream before assuming the @@ -269,6 +311,62 @@ const DEFAULT_CALL_TIMEOUT_MS = SLOWEST_SERVER_TOOL_MS + 60_000; /** An override below this is a mistake, not an intent. */ const MIN_CALL_TIMEOUT_MS = 1_000; +/** Longest a given tool can legitimately take server-side, keyed by tool name. + * + * `DEFAULT_CALL_TIMEOUT_MS` is sized for `memwal_analyze`, the slowest tool + * there is. Applying that one number to every call means a request whose reply + * is lost — the relayer answered, the stream dropped before it arrived — keeps + * the agent blocked for 240s even when the tool could not still be working. + * Users read that as a hang and reload the client. + * + * Each entry is the ceiling the matching tool enforces on itself in + * `services/server/scripts/mcp/tools/`: `MAX_REMEMBER_WAIT_MS` for + * `memwal_remember` (its default wait is 0 — it returns at accept — but an + * operator can raise `MEMWAL_MCP_REMEMBER_WAIT_MS` up to that cap), + * `MAX_STATUS_WAIT_MS` for `memwal_remember_status`, and the fixed `timeoutMs` + * the bulk and analyze tools pass to the SDK. Keep them in lockstep: a value + * below a tool's real ceiling abandons healthy work. Unlisted tools keep the + * default. + * + * `memwal_recall` is the exception to lockstep: the SDK aborts the recall + * request itself at 15s, but the checks it runs first (`/version`, `/config`) + * carry their own deadlines, or none on older SDKs. 90s covers those. */ +const TOOL_DEADLINE_MS: Readonly> = { + memwal_recall: 90_000, + memwal_remember: 90_000, + memwal_remember_status: 60_000, + memwal_remember_bulk: 120_000, + memwal_analyze: SLOWEST_SERVER_TOOL_MS, +}; + +/** Absorbs relayer + transport overhead on top of a tool's own ceiling. The + * sidecar answers at its deadline with a result or an error envelope rather + * than going quiet, so the reply is one network hop behind it; 30s is many + * times that. Cutting a merely-late reply off early is the expensive mistake — + * the agent would retry a write that actually landed. */ +const ORPHAN_HEADROOM_MS = 30_000; + +/** Tool name for a `tools/call`, or null for any other JSON-RPC method. */ +function toolNameOf(msg: RpcMessage): string | null { + if (msg.method !== "tools/call") return null; + const params = msg.params; + if (params == null || typeof params !== "object") return null; + const name = (params as { name?: unknown }).name; + return typeof name === "string" ? name : null; +} + +/** Deadline for one tracked request. A tool with a known ceiling gets that + * plus headroom; everything else keeps the global default. An explicit + * `MEMWAL_MCP_CALL_TIMEOUT_MS` pins every call, so tests still drive expiry + * from one knob. */ +function resolveDeadlineMs(msg: RpcMessage): number { + const fallback = resolveCallTimeoutMs(); + if (process.env.MEMWAL_MCP_CALL_TIMEOUT_MS) return fallback; + const tool = toolNameOf(msg); + const ceiling = tool === null ? undefined : TOOL_DEADLINE_MS[tool]; + return ceiling === undefined ? fallback : ceiling + ORPHAN_HEADROOM_MS; +} + /** Without a cap, a long deadline drifts by a third of itself. */ const MAX_ORPHAN_SWEEP_MS = 5_000; @@ -406,6 +504,20 @@ interface InFlightEntry { * mid-session outage — the ordinary case — a request that never left the * process was indistinguishable from one already sent. */ sent?: boolean; + /** How long this call may go unanswered before the sweeper declares its + * reply lost. Fixed when the request is first tracked, so a reconnect + * replay keeps the original budget. */ + deadlineMs: number; + /** Set while the sweeper asks the relayer's `/health` why this sent call + * went unanswered, so the next sweep does not probe it again. */ + probing?: boolean; + /** Set once the sweeper has decided this sent call's reply is lost and it + * is being answered as failed. A replay would post it a second time — a + * duplicate write for a tool with no idempotency key — while the agent is + * told the call never ran, so `reconnect()` skips these. Never cleared, + * unlike `probing`: the call is past its deadline whichever way the answer + * goes. */ + orphaned?: boolean; } /** The relayer rejected the saved delegate key (HTTP 401 on the handshake). @@ -423,6 +535,12 @@ class RelayerUnauthorizedError extends Error { interface SseHandshakeResult { /** Absolute URL the client must POST to for outbound JSON-RPC messages. */ postUrl: string; + /** + * Forward one message, resolving with the HTTP status. Same shape as the + * Streamable transport's `send`, so the forwarding path does not have to + * know which transport is underneath. + */ + send: (msg: RpcMessage, creds: MemWalCredentials, extra: Record) => Promise; /** Per-line iterator for incoming SSE messages (already-parsed JSON-RPC). */ iter: AsyncIterator; /** Abort + close the SSE stream. */ @@ -440,6 +558,28 @@ function mcpAuthHeaders( }; } +/** + * Open a relayer session on the configured transport. + * + * `MEMWAL_MCP_TRANSPORT=http` dials the Streamable HTTP endpoint, which + * answers a call on the same request instead of splitting POST from reply. + * Default stays SSE until the new path has production mileage. + */ +async function openRelaySession( + relayerUrl: string, + creds: MemWalCredentials, + extraHeaders: Record = {}, +): Promise { + if (resolveTransport(process.env.MEMWAL_MCP_TRANSPORT) === "http") { + // `postUrl` is logging-only on this path; the session owns its + // endpoint. A plain cast, not `as unknown as` — the two shapes must + // stay structurally compatible, and a widening cast would hide it if + // they ever stopped being. + return (await openStreamableSession(relayerUrl, creds, extraHeaders)) as SseHandshakeResult; + } + return openSseStream(relayerUrl, creds, extraHeaders); +} + async function openSseStream( relayerUrl: string, creds: MemWalCredentials, @@ -712,6 +852,7 @@ async function openSseStream( return { postUrl, + send: (msg, sendCreds, extra) => postMessage(postUrl, msg, sendCreds, extra), iter, abort: () => { controller.abort(); @@ -810,6 +951,11 @@ async function handleLocalLogin( }, }); }, + // This tool call has already returned "here is your URL, go sign in", + // so a later failure has no response left to ride home on. Without an + // out-of-band notification the agent sits waiting on a flow that is + // already dead. MCP logging notifications are fire-and-forget and safe + // to emit at any point in the session. (err) => { const msg = err instanceof Error ? err.message : String(err); log.warn("memwal_login.bridge.failed", { msg }); @@ -819,7 +965,15 @@ async function handleLocalLogin( params: { level: "warning", logger: "memwal-mcp", - data: `Walrus Memory sign-in did not complete: ${msg}. Existing credentials are unchanged; call memwal_login again to retry.`, + // The reclaim is only possible because of the write-ahead + // record (WALM-332): a key the browser already paid to + // register is no longer lost with the process. A retry + // cannot help that key, since it reuses it and the + // dashboard cannot register it twice. + data: + `Walrus Memory sign-in did not complete: ${msg}. Existing credentials are ` + + `unchanged. If you approved the wallet step, the next start reclaims that ` + + `key; otherwise call memwal_login again to retry.`, }, }); }, @@ -861,6 +1015,15 @@ async function handleLocalLogin( function handleLocalLogout(): { text: string; isError: boolean } { try { const cleared = clearCreds(); + // Explicit sign-out discards the write-ahead record too. Without this + // an interrupted re-login leaves `login-pending.json` behind, and the + // next start's `recoverPendingLogin` signs the user straight back in. + // + // Kept out of `clearCreds()` so only a deliberate sign-out discards a + // key that may still be reclaimable. `clearCreds` is exported, and a + // 401 deliberately does NOT wipe credentials (see the relayer-401 + // handling above), so the two are not the same decision. + clearPendingLogin(); log.info("memwal_logout.bridge.success", { removedPath: cleared.removedPath ?? null, fallbackPath: cleared.fallbackPath ?? null, @@ -1092,7 +1255,7 @@ export async function runBridge( } function postIfCurrent( epoch: number, - postUrl: string, + send: SseHandshakeResult["send"], msg: RpcMessage, postCreds: MemWalCredentials, ): Promise { @@ -1105,7 +1268,7 @@ export async function runBridge( const tracked = inFlight.get(msg.id); if (tracked) tracked.sent = true; } - return postMessage(postUrl, msg, postCreds, extraHeaders).then((status) => { + return send(msg, postCreds, extraHeaders).then((status) => { // 404 is the relayer saying that session does not exist, so the // message was discarded rather than routed: it provably did not // run, and the request goes back to being never-sent. @@ -1234,6 +1397,7 @@ export async function runBridge( const inFlight = new Map(); const callTimeoutMs = resolveCallTimeoutMs(); const stalledHandshakeMs = resolveStalledHandshakeMs(callTimeoutMs); + const healthProbeMs = resolveHealthProbeMs(); /** IDs of `tools/list` requests we've forwarded to the relayer. When * the response comes back through the SSE pump, we splice in the @@ -1241,11 +1405,22 @@ export async function runBridge( * client surfaces them in its tool palette. */ const pendingListIds = new Set(); - /** IDs of forwarded `memwal_health` calls, each against the relayer URL the - * call went out on. Captured at send time rather than read at reply time so - * a reconnect that swapped credentials mid-flight cannot label the answer - * with a relayer it did not come from. */ - const pendingHealthIds = new Map(); + /** Tool names the CONNECTED relayer advertised on its last `tools/list`, + * minus the ones we serve locally. Empty until the client has listed tools + * at least once over the *current* session. `reconnect()` clears it: a + * stale allow would re-forward a tool the new relayer does not serve + * (GH #928), and a stale deny would refuse a tool it does. Until the + * next upstream list, the gate allows only the cold-start floor. */ + const upstreamToolNames = new Set(); + + /** IDs of forwarded `memwal_health` calls, each against the destination the + * call went out on — relayer URL and account. Captured at send time rather + * than read at reply time so a reconnect that swapped credentials mid-flight + * cannot label the answer with a destination it did not come from. */ + const pendingHealthIds = new Map< + string | number, + { relayerUrl: string; accountId?: string } + >(); /** Record a 429 and tell the user ONCE that this is a rate limit rather * than a broken config — the distinction the MCP host cannot make for @@ -1334,7 +1509,7 @@ export async function runBridge( // gone, so there is nothing to authorize a new session // with. Belt-and-braces against `loggedOut` alone. if (!openingCreds) break; - const candidate = await openSseStream( + const candidate = await openRelaySession( openingCreds.relayerUrl, openingCreds, connectHeaders(), @@ -1371,6 +1546,15 @@ export async function runBridge( throttleNoticed = false; clearHandshakeFailure(); endConnectEpisode(); + // This session has not advertised anything yet. Keep the + // previous set and a login that swapped relayerUrl would + // re-forward a tool the new one does not serve, or refuse + // one it does, until the client happened to re-list. + upstreamToolNames.clear(); + writeStdoutMessage({ + jsonrpc: "2.0", + method: "notifications/tools/list_changed", + }); // An accepted handshake retires any earlier rejection — // `memwal_login` re-registers a key and lands here, not on // the background connect's publish path, so clearing only @@ -1409,6 +1593,15 @@ export async function runBridge( log.info("bridge.replay_halted_signed_out", { id }); break; } + // The sweeper already declared this call's reply lost + // and is answering it as failed. Posting it again would + // run it a second time — a duplicate write for a tool + // with no idempotency key — and the reply would be + // dropped anyway, as a late answer for a closed-out id. + if (entry.orphaned) { + log.info("bridge.replay_skipped_orphaned", { id }); + continue; + } const msg = entry.msg; try { // A replayed `initialize` produces a fresh upstream @@ -1425,9 +1618,9 @@ export async function runBridge( expectSuppressedReply(msg.id); } const epoch = sessionEpoch; - const postUrl = sse.postUrl; + const send = sse.send; const status = await enqueuePost(() => - postIfCurrent(epoch, postUrl, msg, openingCreds), + postIfCurrent(epoch, send, msg, openingCreds), ); log.info("bridge.replayed", { id, status }); } catch (err) { @@ -1750,6 +1943,16 @@ export async function runBridge( (t) => !LOCAL_TOOL_NAMES.has(t.name ?? ""), ); result.tools = [...upstream, ...LOCAL_TOOL_DEFINITIONS]; + // Record what this relayer actually serves. A + // later call for a name absent here is answered + // locally instead of being forwarded into a wait + // no reply will ever end. + upstreamToolNames.clear(); + for (const t of upstream) { + if (typeof t.name === "string" && t.name !== "") { + upstreamToolNames.add(t.name); + } + } } } if ( @@ -1765,7 +1968,8 @@ export async function runBridge( if (dialled !== undefined) { annotateHealthResult( value.result as { content?: unknown; isError?: unknown }, - dialled, + dialled.relayerUrl, + dialled.accountId, ); } } @@ -1928,6 +2132,48 @@ export async function runBridge( return; } + // Version skew: this bridge ships on npm and updates itself, + // while a relayer ships per environment and does not, so the + // bridge is routinely newer than the server it dials. A tool + // named in initialize instructions or a cached list can be + // missing from the session that actually came up — + // `memwal_remember_status` against a prod relayer, which is + // GH #928. Forwarding that call parks it in `inFlight` until + // the orphan sweeper's deadline (60s + 30s headroom for that + // tool), and the user reads the 90s as a hang. + // + // Fail closed. Before any upstream `tools/list`, only the + // cold-start floor is callable. After one, only what that + // relayer advertised (plus locally-served tools). A stale + // tool list is not a transport fault: say so now, while the + // agent can still act on it. + if (msg.method === "tools/call" && msg.id != null) { + const called = (msg.params as { name?: string } | undefined)?.name; + if (typeof called === "string" && !LOCAL_TOOL_NAMES.has(called)) { + const known = + upstreamToolNames.size > 0 + ? upstreamToolNames + : COLD_START_TOOL_NAMES; + if (!known.has(called)) { + const available = [ + ...(upstreamToolNames.size > 0 + ? [...upstreamToolNames, ...LOCAL_TOOL_NAMES] + : COLD_START_TOOL_NAMES), + ].sort(); + log.warn("bridge.tool_not_served", { + tool: called, + window: + upstreamToolNames.size > 0 ? "upstream" : "cold-start", + }); + failRequest(msg, "tool not served", { + toolText: unknownToolText(called, available), + errorMessage: `${called} is not served by this Walrus Memory relayer`, + }); + return; + } + } + } + // Fill in the configured default namespace for memory tool // calls that didn't pass one. Mutates msg in place so the // forwarded — and any replayed-on-reconnect — copy carries it. @@ -1946,7 +2192,10 @@ export async function runBridge( msg.id != null && (msg.params as { name?: string } | undefined)?.name === "memwal_health" ) { - pendingHealthIds.set(msg.id, creds?.relayerUrl ?? config.relayerUrl); + pendingHealthIds.set(msg.id, { + relayerUrl: creds?.relayerUrl ?? config.relayerUrl, + accountId: creds?.accountId, + }); } // Track requests (have both method and id) so we can replay @@ -1957,7 +2206,11 @@ export async function runBridge( msg.id !== undefined && msg.id !== null ) { - inFlight.set(msg.id, { msg, startedAt: Date.now() }); + inFlight.set(msg.id, { + msg, + startedAt: Date.now(), + deadlineMs: resolveDeadlineMs(msg), + }); } // Relayer session not up yet, OR the post-connect flush is still // draining — buffer so this request stays behind everything that @@ -1997,10 +2250,10 @@ export async function runBridge( return; } const epoch = sessionEpoch; - const postUrl = sse.postUrl; + const send = sse.send; const postCreds = creds; const status = await enqueuePost(() => - postIfCurrent(epoch, postUrl, msg, postCreds), + postIfCurrent(epoch, send, msg, postCreds), ); if (status === 404) { log.warn("bridge.session_stale", { sessionUrl: sse.postUrl }); @@ -2038,10 +2291,10 @@ export async function runBridge( const msg = pendingForward.shift()!; try { const epoch = sessionEpoch; - const postUrl = sse.postUrl; + const send = sse.send; const postCreds = creds; const status = await enqueuePost(() => - postIfCurrent(epoch, postUrl, msg, postCreds), + postIfCurrent(epoch, send, msg, postCreds), ); if (status === 404) { // Stale session right after connect. EVERY id-bearing @@ -2213,27 +2466,27 @@ export async function runBridge( "memwal_analyze", ]); - /** Name of the tool a tracked request was calling, when it was one. */ - function toolNameOf(msg: RpcMessage): string | null { - if (msg.method !== "tools/call") return null; - const params = msg.params as { name?: unknown } | undefined; - return typeof params?.name === "string" ? params.name : null; - } - function expiredRequestReport( neverSent: boolean, now: number, tool: string | null, + health?: { probe: HealthProbe; relayerUrl: string }, ): { reason: string; opts: { toolText: string; errorMessage: string }; } { if (!neverSent) { + // What the relayer's `/health` said just now: tells a dead + // relayer from a wrong URL from one stuck call. + const described = health ? describeHealthProbe(health.probe, health.relayerUrl) : null; // The request reached the relayer. What is missing is the reply, // and for a write that distinction is the whole message: the work // may have completed, may still be running, and cannot be assumed // undone. "Please retry" is only safe advice for a read. if (tool !== null && MUTATING_TOOLS.has(tool)) { + const healthNote = described + ? `\nRelayer health: ${described.health}. ${described.verdict}` + : ""; return { reason: "no response to a sent write", opts: { @@ -2244,23 +2497,39 @@ export async function runBridge( "it and does not mean nothing was stored. Do NOT simply repeat the " + "call: run `memwal_recall` for this content first, and only re-save " + "what is genuinely missing. Repeating a bulk save that already " + - "landed stores a second paid copy.", + "landed stores a second paid copy." + + healthNote, errorMessage: `Walrus Memory ${tool} was sent but its reply never arrived. The write ` + - "may have completed; verify with recall before retrying.", + "may have completed; verify with recall before retrying." + + healthNote, }, }; } + // Not "this call only reads": `memwal_restore` re-indexes. What + // makes a retry safe is that none of these can store a duplicate. + const nextStep = + described === null || described.reachable + ? "Repeating this call cannot store a duplicate, so it is safe to retry " + + "once. If it keeps happening, report it with the time of the call." + : "Repeating this call cannot store a duplicate, so it is safe to retry " + + "after the relayer is reachable again — wait a minute if the problem " + + "is on the relayer's side."; + const healthValue = described?.health ?? "not checked"; return { reason: "no response", opts: { - toolText: - "❌ Walrus Memory did not answer this call. The request reached the " + - "relayer but the reply never came back. This call only reads, so it is " + - "safe to retry.", + toolText: [ + "❌ Walrus Memory did not answer this call.", + "Cause: the request reached the relayer but no reply came back." + + (described ? ` ${described.verdict}` : ""), + `Relayer health: ${healthValue}`, + `Next step: ${nextStep}`, + ].join("\n"), errorMessage: - "Walrus Memory call was orphaned by a reconnect and never " + - "received a response. Safe to retry: this call only reads.", + "Walrus Memory call reached the relayer but never received a " + + `response (relayer health: ${healthValue}). Safe to retry: it cannot ` + + "store a duplicate.", }, }; } @@ -2310,6 +2579,8 @@ export async function runBridge( const orphanSweeper = setInterval(() => { const now = Date.now(); const handshakeStalledMs = handshakeStalledForMs(now); + // One `/health` request per sweep, however many calls expired in it. + let sweepProbe: Promise | null = null; for (const [id, entry] of Array.from(inFlight.entries())) { const elapsedMs = now - entry.startedAt; // Never sent = no POST was ever issued for it. Read from the entry @@ -2325,9 +2596,71 @@ export async function runBridge( // invites a duplicate write. const handshakeIsStalled = handshakeStalledMs !== null && handshakeStalledMs > stalledHandshakeMs; + // `entry.deadlineMs` is this tool's own ceiling plus headroom, not + // the global one sized for the slowest tool — so a `memwal_remember` + // whose reply is lost is answered at 120s instead of 240s. The + // stalled-handshake shortcut still wins when it is tighter, but can + // never extend a tool past its own deadline. const deadlineMs = - neverSent && handshakeIsStalled ? stalledHandshakeMs : callTimeoutMs; + neverSent && handshakeIsStalled + ? Math.min(stalledHandshakeMs, entry.deadlineMs) + : entry.deadlineMs; if (elapsedMs <= deadlineMs) continue; + // `initialize` is answered locally and gets no reply here, so + // there is nothing to explain and nothing to probe for. + if (!neverSent && entry.msg.method !== "initialize") { + // A sent call: ask the relayer's `/health` before answering, + // so the message can say whether it is down, unreachable, or + // up with this one call stuck. Only the answer waits on the + // probe; the bookkeeping stays synchronous. + if (entry.probing) continue; + entry.probing = true; + // From here the call is the sweeper's to answer: a reconnect + // must not replay it behind the probe's back. + entry.orphaned = true; + const relayerUrl = creds?.relayerUrl ?? config.relayerUrl; + // `probeRelayerHealth` does not reject, but a call left with + // `probing` set and no answer is the one outcome this path + // must never have, so a rejection still answers it. + sweepProbe ??= probeRelayerHealth(relayerUrl, healthProbeMs).catch(() => null); + void sweepProbe + .then((probe) => { + // A late reply, a logout or a shutdown may have answered + // it while the probe ran. Answering again would be a + // second response for the same id — and after stdin has + // closed there is nobody left to answer. + if (stdinClosed || inFlight.get(id) !== entry) return; + const settledAt = Date.now(); + const { reason, opts } = expiredRequestReport( + false, + settledAt, + toolNameOf(entry.msg), + probe ? { probe, relayerUrl } : undefined, + ); + log.warn("bridge.call_orphaned", { + id, + method: entry.msg.method ?? null, + elapsedMs: settledAt - entry.startedAt, + deadlineMs, + reason, + health: probe?.kind ?? null, + healthMs: probe?.ms ?? null, + handshakeStalledMs: handshakeStalledForMs(settledAt), + lastHandshakeError, + }); + failRequest(entry.msg, reason, opts); + }) + .catch((err: unknown) => { + // Never leave a call marked `probing` with no answer: + // clearing it lets the next sweep answer it. + entry.probing = false; + log.warn("bridge.call_orphaned_answer_failed", { + id, + error: err instanceof Error ? err.message : String(err), + }); + }); + continue; + } // Built only for what actually expired: this walks `pendingForward` // and interpolates two user-facing strings, and the branch it // serves fires roughly never. @@ -2403,7 +2736,7 @@ export async function runBridge( } const openingGeneration = credentialGeneration; try { - const candidate = await openSseStream(creds.relayerUrl, creds, connectHeaders()); + const candidate = await openRelaySession(creds.relayerUrl, creds, connectHeaders()); if (stdinClosed) { candidate.abort(); break; diff --git a/packages/mcp/src/consent.ts b/packages/mcp/src/consent.ts new file mode 100644 index 000000000..58b21d890 --- /dev/null +++ b/packages/mcp/src/consent.ts @@ -0,0 +1,137 @@ +/** + * The automatic-memory consent question (WALM-642). + * + * Asked at interactive login and NOWHERE else. Deliberately not an MCP tool, + * not a tool description, not an instruction, and not anything a model can + * reach: a model answering on the user's behalf is not consent, and an + * agent-shaped surface for this question would be exactly that. The only caller + * is `main()` in index.ts, behind `process.stdin.isTTY`. + * + * Login is the gate because MemWal is unusable without it, so every real user + * passes through it once, and it is already a moment that requires a terminal. + * + * On the wording — it is the substance of this change, not decoration: + * + * - It names the consequence ("writes it to your memory without asking each + * time"), not the feature. Someone should be able to picture what happens + * to them, in their own session, from reading it. + * - Permanence leads, because it is the fact that changes the answer. Walrus + * is immutable: you can stop saving new memories but you cannot take back + * one already saved. Burying that under "we value your privacy" would make + * this a cookie banner, which it is not. + * - The redaction claim is deliberately hedged. `redaction.ts` documents real + * residual gaps — an unlabelled hex secret, a bare BIP-39 word run — so + * "safety net, not a guarantee" is the accurate claim and must stay. Saying + * more would be selling a promise the code does not make. + * - Declining is free and is stated as a normal choice, not a warning. A "no" + * is recorded once and never asked again. + */ +import { createInterface } from "node:readline"; +import type { Readable, Writable } from "node:stream"; + +/** The prompt body, everything above the input line. */ +export const CONSENT_PROMPT = [ + "", + "MemWal can save things about you automatically.", + "", + "What that means: when you state a preference, a decision, or a setting in", + 'chat — "I prefer pnpm", "we deploy from dev", "the relayer is at X" —', + "MemWal writes it to your memory without asking each time, so it is there", + "in your next session and in every other client you use.", + "", + "Before you choose:", + "", + " - Saved memories are permanent. They go to Walrus, which is immutable", + " storage. You can stop saving new ones at any time, but you cannot", + " delete one that is already saved.", + " - They are encrypted to your account. Only your delegate key reads them.", + " - MemWal strips obvious credentials — API keys, tokens, passwords,", + " private keys — before saving. Treat that as a safety net, not a", + " guarantee: do not paste secrets into a session with this on.", + "", + " [1] Save automatically recommended, this is what MemWal is for", + ' [2] Only save when I ask nothing is saved unless you say "remember this"', + "", + "Change this any time with `memwal-mcp auto-save on|off`.", + "See what is stored with `memwal_recall`.", + "", +].join("\n"); + +/** The input line. Enter alone takes the recommended option. */ +export const CONSENT_QUESTION = "Your choice [1/2]: "; + +/** What to say when the answer is neither 1 nor 2. */ +export const CONSENT_REPROMPT = + "Please answer 1 or 2 (or press Enter for 1)."; + +/** + * Map one line of input to an answer. + * + * Empty (a bare Enter) accepts option 1. Anything that is not 1 or 2 returns + * `null`, which means re-ask — never "assume the recommended one", because a + * typo is not an answer to a question about permanent storage. + */ +export function interpretConsentAnswer(raw: string): boolean | null { + const v = raw.trim().toLowerCase(); + if (v === "") return true; + if (v === "1") return true; + if (v === "2") return false; + return null; +} + +/** + * Ask the question on a terminal and resolve to the answer. + * + * Resolves `null` — meaning "no answer given" — if the input stream ends first + * (Ctrl-D, a closed pipe, a killed terminal). The caller leaves the state unset + * in that case and asks again next time rather than recording a choice the user + * never made. + * + * Streams are injected so this is testable without a pty; `main()` passes the + * real stdin and stderr. Output goes to stderr because stdout belongs to the + * MCP protocol on every other path in this package. + */ +export async function askAutoSaveConsent(opts: { + input: Readable; + output: Writable; + /** Guard against an accidental headless call. Defaults to requiring a TTY. */ + isTTY?: boolean; +}): Promise { + // A prompt with nobody in front of it is a hang, not a question. The + // caller already checks this; the second check is here because the cost of + // getting it wrong is an MCP server that never finishes starting. + if (opts.isTTY === false) return null; + + opts.output.write(`${CONSENT_PROMPT}\n`); + + const rl = createInterface({ input: opts.input, terminal: false }); + try { + opts.output.write(CONSENT_QUESTION); + // The interface's async iterator, rather than a promise per `line` + // event: attaching a fresh listener after each answer drops any line + // that arrived while nothing was listening, which deadlocks the re-ask + // path the moment input is buffered rather than typed. + for await (const line of rl) { + const answer = interpretConsentAnswer(line); + if (answer !== null) return answer; + opts.output.write(`${CONSENT_REPROMPT}\n`); + opts.output.write(CONSENT_QUESTION); + } + // Exhausted without an answer: Ctrl-D, a closed pipe, a killed + // terminal. Not a choice, and not recorded as one. + opts.output.write("\n"); + return null; + } finally { + rl.close(); + } +} + +/** What is printed once the answer is in. */ +export function consentOutcomeNotice(enabled: boolean, path: string): string { + return enabled + ? "Automatic memory is ON. MemWal will save durable facts as you state them. " + + `Turn it off any time with \`memwal-mcp auto-save off\`. (${path})` + : "Automatic memory is OFF. Nothing is saved unless you ask for it — " + + '"remember this" still works, and so does recall. ' + + `Turn it on any time with \`memwal-mcp auto-save on\`. (${path})`; +} diff --git a/packages/mcp/src/crypto.ts b/packages/mcp/src/crypto.ts index 2d138a96a..fbc3aa53e 100644 --- a/packages/mcp/src/crypto.ts +++ b/packages/mcp/src/crypto.ts @@ -1,7 +1,7 @@ /** * Ed25519 helpers — pure-JS via @noble/ed25519. */ -import { getPublicKeyAsync, utils } from "@noble/ed25519"; +import { getPublicKeyAsync, signAsync, utils } from "@noble/ed25519"; import { blake2b } from "@noble/hashes/blake2.js"; function hex(bytes: Uint8Array): string { @@ -48,4 +48,17 @@ export function deriveSuiAddress(pubKey: Uint8Array): string { return "0x" + hex(digest); } +/** + * Sign a canonical request message with a delegate private key. + * + * The relayer authenticates `/api/*` by Ed25519 signature over + * `{timestamp}.{method}.{path_and_query}.{body_sha256}.{nonce}.{account_id}` + * (`services/server/src/auth.rs`, which calls itself the single source of + * truth for that format). Keep the two in lockstep. + */ +export async function signMessage(privateKeyHex: string, message: string): Promise { + const sig = await signAsync(new TextEncoder().encode(message), fromHex(privateKeyHex)); + return hex(sig); +} + export { hex as bytesToHex, fromHex as hexToBytes }; diff --git a/packages/mcp/src/health-probe.ts b/packages/mcp/src/health-probe.ts new file mode 100644 index 000000000..acf587f05 --- /dev/null +++ b/packages/mcp/src/health-probe.ts @@ -0,0 +1,105 @@ +/** + * Ask the relayer's public `/health` why a call went unanswered: down, + * unreachable from this machine, configured wrong, or up with one call stuck. + * Each wants a different next step, so the bridge probes before it answers. + */ + +export type HealthProbe = + | { kind: "ok"; ms: number; version?: string; writesUnavailable?: boolean } + | { kind: "http"; ms: number; status: number } + | { kind: "unreachable"; ms: number; code: string } + | { kind: "timeout"; ms: number }; + +const DEFAULT_HEALTH_PROBE_MS = 3_000; +const MIN_HEALTH_PROBE_MS = 100; +/** A diagnosis nobody waits a minute for. Also keeps the value inside what + * `AbortSignal.timeout` accepts. */ +const MAX_HEALTH_PROBE_MS = 60_000; + +/** Override via `MEMWAL_MCP_HEALTH_PROBE_MS`, mostly for tests. */ +export function resolveHealthProbeMs(): number { + const raw = process.env.MEMWAL_MCP_HEALTH_PROBE_MS; + if (!raw) return DEFAULT_HEALTH_PROBE_MS; + const n = Math.floor(Number(raw)); + if (!Number.isFinite(n) || n < MIN_HEALTH_PROBE_MS) return DEFAULT_HEALTH_PROBE_MS; + return Math.min(n, MAX_HEALTH_PROBE_MS); +} + +/** `GET {baseUrl}/health`, bounded by `timeoutMs`. Never rejects: a probe + * that fails is itself the answer. */ +export async function probeRelayerHealth( + baseUrl: string, + timeoutMs: number, +): Promise { + const started = Date.now(); + let signal: AbortSignal | undefined; + try { + // Inside the `try`: `AbortSignal.timeout` throws on a value it + // cannot take, and this function must not reject. + signal = AbortSignal.timeout(timeoutMs); + const res = await fetch(`${baseUrl.replace(/\/+$/, "")}/health`, { signal }); + const ms = Date.now() - started; + if (!res.ok) { + await res.body?.cancel(); + return { kind: "http", ms, status: res.status }; + } + const body = (await res.json().catch(() => null)) as { + version?: unknown; + write_ready?: unknown; + writes?: unknown; + } | null; + const version = typeof body?.version === "string" ? body.version : undefined; + // `/health` answers 200 while writes are paused or Postgres is + // full; "ok" alone would hide the cause of a failed write. + const writesUnavailable = body?.write_ready === false || body?.writes === "paused"; + return { kind: "ok", ms, version, ...(writesUnavailable ? { writesUnavailable } : {}) }; + } catch (err) { + const ms = Date.now() - started; + if (signal?.aborted) return { kind: "timeout", ms }; + const code = (err as { cause?: { code?: unknown } } | null)?.cause?.code; + return { kind: "unreachable", ms, code: typeof code === "string" ? code : "unknown" }; + } +} + +/** DNS could not find the host: the one cause a user fixes on their side. */ +const UNRESOLVED_CODES = new Set(["ENOTFOUND", "EAI_AGAIN"]); + +/** `health` is the value for a `Relayer health:` line; `verdict` says what + * it means for the call; `reachable` is whether a plain retry can help. */ +export function describeHealthProbe( + probe: HealthProbe, + relayerUrl: string, +): { health: string; verdict: string; reachable: boolean } { + switch (probe.kind) { + case "ok": + return { + health: + `ok (${probe.ms}ms${probe.version ? `, v${probe.version}` : ""}` + + `${probe.writesUnavailable ? ", writes unavailable" : ""})`, + verdict: + "The relayer is up, so this call stalled inside it or its reply was lost on the way back.", + reachable: true, + }; + case "http": + return { + health: `HTTP ${probe.status} (${probe.ms}ms)`, + verdict: "The relayer answered but is not healthy.", + reachable: false, + }; + case "timeout": + return { + health: `no answer within ${(probe.ms / 1000).toFixed(1)}s`, + verdict: `The relayer at ${relayerUrl} is down, overloaded, or not reachable from this machine.`, + reachable: false, + }; + case "unreachable": + return { + health: `unreachable (${probe.code})`, + verdict: UNRESOLVED_CODES.has(probe.code) + ? `The relayer host in ${relayerUrl} could not be resolved. Check the relayer URL ` + + "(MEMWAL_SERVER_URL or --relayer) and this machine's network." + : `The relayer at ${relayerUrl} is down or not reachable from this machine.`, + reachable: false, + }; + } +} diff --git a/packages/mcp/src/index.ts b/packages/mcp/src/index.ts index 02b9b3fff..90f80dd5c 100644 --- a/packages/mcp/src/index.ts +++ b/packages/mcp/src/index.ts @@ -9,10 +9,31 @@ * 5. On 401 (revoked key), the bridge wipes credentials before throwing * — the next process spawn will re-trigger login. */ -import { clearCreds, credsPath, loadCreds } from "./auth.js"; +import { + approveProjectCreds, + clearCreds, + clearPendingLogin, + credsPath, + formatProjectCredsNotice, + formatProjectCredsStorageWarning, + loadCreds, + resolveCreds, + revokeProjectCredsApproval, +} from "./auth.js"; +import { recoverPendingLogin, formatStrandedLoginNotice } from "./recovery.js"; import { runAuthRequiredServer } from "./auth-required.js"; import { notePendingLoginSuccess, runBridge } from "./bridge.js"; import { loginFlow } from "./login.js"; +import { + autoSaveStatus, + autoSaveSummary, + markConsentPending, + pendingConsentNotice, + publishHookState, + setAutoSave, + AUTO_SAVE_ENV, +} from "./auto-save.js"; +import { askAutoSaveConsent, consentOutcomeNotice } from "./consent.js"; import { log, note } from "./logger.js"; /** @@ -24,10 +45,17 @@ interface ParsedArgs { help: boolean; logout: boolean; forceLogin: boolean; + /** Approve the project-local credentials found from the working directory + * (WALM-639). A repo file is inert until this has been run for it. */ + approveProject: boolean; + /** Withdraw that approval again. */ + revokeProject: boolean; relayerUrl?: string; webUrl?: string; label?: string; namespace?: string; + /** `auto-save on|off|status` — the automatic-memory opt-in (WALM-642). */ + autoSave?: "on" | "off" | "status"; /** Args parseArgs did not recognise, in the order seen. For a flag * written `--key=value`, only `--key` is recorded — see parseArgs. */ unknown: string[]; @@ -44,10 +72,17 @@ const ENV_PRESETS: Record = { /** Bare words that are commands rather than values. An unknown flag must not * swallow one as its argument. */ -const POSITIONALS = new Set(["login"]); +const POSITIONALS = new Set(["login", "approve-project", "revoke-project", "auto-save", "on", "off", "status"]); export function parseArgs(argv: string[]): ParsedArgs { - const out: ParsedArgs = { help: false, logout: false, forceLogin: false, unknown: [] }; + const out: ParsedArgs = { + help: false, + logout: false, + forceLogin: false, + approveProject: false, + revokeProject: false, + unknown: [], + }; for (let i = 0; i < argv.length; i++) { const a = argv[i]; const next = () => argv[++i]; @@ -63,6 +98,28 @@ export function parseArgs(argv: string[]): ParsedArgs { case "login": out.forceLogin = true; break; + case "--approve-project": + case "approve-project": + out.approveProject = true; + break; + case "--revoke-project": + case "revoke-project": + out.revokeProject = true; + break; + case "auto-save": + case "--auto-save": { + // `auto-save` on its own reports the state rather than + // changing it — a bare subcommand must never be read as + // consent to turn automatic saving on. + const value = argv[i + 1]?.toLowerCase(); + if (value === "on" || value === "off" || value === "status") { + out.autoSave = value; + i++; + } else { + out.autoSave = "status"; + } + break; + } case "--prod": case "--dev": case "--staging": @@ -147,8 +204,54 @@ export async function main(argv: string[] = process.argv.slice(2)): Promise env > default. const relayerUrl = args.relayerUrl ?? process.env.MEMWAL_SERVER_URL ?? "https://relayer.memory.walrus.xyz"; @@ -196,7 +397,47 @@ export async function main(argv: string[] = process.argv.slice(2)): Promise { + if (!autoSaveStatus().pendingConsent) return; + + const answer = await askAutoSaveConsent({ + input: process.stdin, + output: process.stderr, + isTTY: process.stdin.isTTY === true, + }); + if (answer === null) { + note( + "No answer recorded — automatic memory is unchanged and you will be " + + "asked again next time. Set it directly with `memwal-mcp auto-save on|off`.", + ); + return; + } + const { path } = setAutoSave(answer); + note(consentOutcomeNotice(answer, path)); +} + function printHelp(): void { process.stderr.write(helpText() + "\n"); } @@ -334,6 +636,34 @@ export function helpText(): string { " browser).", " memwal-mcp --logout Delete saved credentials without", " re-running login.", + " memwal-mcp approve-project Approve the project-local", + " .memwal/credentials.json found from", + " the current directory, so memory", + " written here goes to ITS account and", + " relayer. Until approved the file is", + " ignored and the global credentials", + " are used. Approval is per machine,", + " stored outside the repository, and", + " required again if the account,", + " delegate key or relayer changes.", + " memwal-mcp revoke-project Withdraw that approval.", + " memwal-mcp auto-save on|off Turn automatic memory on or off.", + " ON once you agree to it: `login` asks", + " in the terminal the first time, and", + " nothing is saved unprompted until you", + " answer. Saved memories are permanent", + " — Walrus is immutable storage — so", + " you can stop saving new ones but", + " cannot delete one already saved.", + " Credentials (passwords, API keys,", + " tokens, private keys, seed phrases,", + " auth headers, URLs with an embedded", + " user:password) are stripped before", + " any write either way — a safety net,", + " not a guarantee. Stored in", + " settings.json next to", + " credentials.json.", + " memwal-mcp auto-save Report the current setting.", " memwal-mcp --help Show this help.", "", "Options:", @@ -365,7 +695,15 @@ export function helpText(): string { " MEMWAL_SERVER_URL same as --relayer", " MEMWAL_WEB_URL same as --web-url", " MEMWAL_CLIENT_LABEL same as --label", + " MEMWAL_CREDS_DIR Use this directory for credentials", + " and approvals, overriding both the", + " project-local and global files.", + " Must be an ABSOLUTE path outside the", + " project; anything else is refused.", " MEMWAL_NAMESPACE same as --namespace", + " MEMWAL_AUTO_SAVE=1 Automatic memory for this server", + " only; overrides settings.json and", + " skips the login question. 0 = off.", " MEMWAL_MCP_DEBUG=1 Verbose stderr logging.", "", "Minimal MCP client config (Cursor, Claude Desktop, etc.):", @@ -409,7 +747,28 @@ export function helpText(): string { } // Re-exports — handy if someone wants to embed this in another tool. -export { loadCreds, saveCreds, clearCreds, credsPath } from "./auth.js"; +export { + loadCreds, + saveCreds, + clearCreds, + credsPath, + resolveCreds, + approveProjectCreds, + revokeProjectCredsApproval, + formatProjectCredsNotice, + formatProjectCredsStorageWarning, +} from "./auth.js"; +export { + isAutoSaveEnabled, + isConsentPending, + autoSaveStatus, + setAutoSave, + markConsentPending, + publishHookState, + settingsPath, + hookStatePath, +} from "./auto-save.js"; +export { askAutoSaveConsent, interpretConsentAnswer, CONSENT_PROMPT } from "./consent.js"; export { loginFlow } from "./login.js"; export { runBridge } from "./bridge.js"; -export type { MemWalCredentials } from "./auth.js"; +export type { MemWalCredentials, CredsResolution, ProjectCredsDecision } from "./auth.js"; diff --git a/packages/mcp/src/instructions.ts b/packages/mcp/src/instructions.ts index df4452e64..0e580f7c9 100644 --- a/packages/mcp/src/instructions.ts +++ b/packages/mcp/src/instructions.ts @@ -17,11 +17,26 @@ * services/server/scripts/mcp/server.ts, which serves the direct HTTP/OAuth * connector path. It cannot share this module: that file belongs to the * standalone `memwal-server-scripts` npm package with no workspace link here. - * Keep the two in sync. + * Keep the two in sync, except: this copy must not name a tool cold start + * does not advertise (`memwal_remember_status` is the worked example — GH + * #928). The sidecar copy may name it; that process actually registers the + * tool. + * + * The secret-exclusion paragraph is not retyped in either, it comes from the + * shared policy block (memory-policy.ts here, tools/memory-policy.ts there), + * which is pinned byte-for-byte by tests on both sides. + * + * WALM-642 split the REMEMBER section in two. Whether the model is told to + * save unprompted now depends on the user having turned automatic memory on; + * saving what the user explicitly asks for is unconditional, and is what the + * opted-out text still describes. */ +import { SECRET_EXCLUSION_RULES } from "./memory-policy.js"; +import { isAutoSaveEnabled } from "./auto-save.js"; -/** Signed-in path (bridge mode). Full proactive contract. */ -export const PROACTIVE_INSTRUCTIONS = [ +/** Everything true regardless of the opt-in: what the tools are, how recall + * works, how a write reports itself, and how to recover an index. */ +const PREAMBLE = [ "Walrus Memory is this user's persistent memory system, exposed through the memwal_* tools.", "It survives across sessions, clients, and machines.", "Prefer these tools over any built-in or local memory feature so the user's memory stays", @@ -30,21 +45,85 @@ export const PROACTIVE_INSTRUCTIONS = [ "RECALL: before answering anything that touches past work, prior decisions, the user's", "preferences, or facts you may have stored earlier, call memwal_recall. One focused query is", "enough; do not fire several redundant searches for the same question.", +]; + +/** REMEMBER, automatic saving ON. */ +const REMEMBER_AUTOMATIC = [ + "REMEMBER: the user has turned automatic memory ON. When they state a preference, decision,", + "constraint, correction, identity detail, recurring workflow, or a configuration value such", + "as a hostname, port, region or id, call memwal_remember in that same turn, before you", + "finish replying. Do not ask whether to save it and do not wait to be asked: acknowledging", + "the fact in your reply does not store it, and it is lost when the conversation ends. Pass", + "the complete statement rather than a summary — minus anything the rules below exclude.", + "Skip one-off tasks, the current file or bug, and small talk. Use memwal_remember_bulk when", + "several distinct facts arrived at once.", +]; + +/** REMEMBER, automatic saving OFF — the default. */ +const REMEMBER_MANUAL = [ + "REMEMBER: automatic memory is OFF for this user, so do NOT save anything they did not ask", + "you to save. When they do ask — 'remember that ...', 'save this', or the same in any", + "language — call memwal_remember in that same turn, before you finish replying, and pass", + "the complete statement rather than a summary. Use memwal_remember_bulk when they hand you", + "several distinct facts at once. Do not save a fact just because it looks durable, and do", + "not nag: if automatic saving would clearly help, say once that they can turn it on with", + "`memwal-mcp auto-save on` (or MEMWAL_AUTO_SAVE=1) and leave it there.", +]; + +const WRITE_CONTRACT = [ + "By default memwal_remember and memwal_remember_bulk return in ~1s once the relayer has", + "accepted the job (job_id / job_ids). The Walrus write continues in the background (~30-60s)", + "and the fact is NOT stored yet. That is the normal result. Do not claim it is saved.", + "Do NOT re-send the same text — that queues duplicates. Settle it with the job-status tool", + "this server advertises (re-list tools if you do not see one; pass job_id, or job_ids for a", + "whole batch). Only a blob_id in the tool reply means the fact is already stored (that", + "happens when an optional wait budget was set and the write finished).", "", - "REMEMBER: when the user states a preference, decision, constraint, correction, identity", - "detail, recurring workflow, or a configuration value such as a hostname, port, region or", - "id, call memwal_remember in that same turn, before you finish replying. Do not ask whether", - "to save it and do not wait to be asked: acknowledging the fact in your reply does not store", - "it, and it is lost when the conversation ends. Pass the complete statement rather than a", - "summary. Skip one-off tasks, the current file or bug, and small talk. Use", - "memwal_remember_bulk when several distinct facts arrived at once.", + "Storage is append-only and encrypted: a fact that lands cannot be edited or deleted. That", + "is why the exclusions below are absolute rather than a preference, and why the write path", + "strips credential shapes from whatever you send and tells you what it removed.", "", "RECOVER: if memwal_recall unexpectedly returns nothing for a namespace that has been used", "before, call memwal_restore to rebuild the index from Walrus.", "", "If a memwal_* tool is not currently loaded, load it and use it. Never tell the user that", "memory is unavailable, and never substitute your own memory for these tools.", -].join("\n"); +]; + +/** + * Build the signed-in instruction payload for a given opt-in state. + * + * Exported separately from the resolver so tests can pin both branches without + * touching the filesystem or the environment. + */ +export function buildProactiveInstructions(opts: { autoSave: boolean }): string { + return [ + ...PREAMBLE, + "", + ...(opts.autoSave ? REMEMBER_AUTOMATIC : REMEMBER_MANUAL), + "", + SECRET_EXCLUSION_RULES, + "", + ...WRITE_CONTRACT, + ].join("\n"); +} + +/** + * The payload for THIS process, resolved from the user's opt-in at call time. + * + * A function, not a const: the opt-in lives on disk and in the environment, and + * `initialize` can arrive after the user has flipped it. + */ +export function proactiveInstructions(): string { + return buildProactiveInstructions({ autoSave: isAutoSaveEnabled() }); +} + +/** + * Signed-in path (bridge mode) with automatic saving on. Kept as a named + * export because it is the text the whole proactive contract is written + * against; `proactiveInstructions()` is what the bridge actually serves. + */ +export const PROACTIVE_INSTRUCTIONS = buildProactiveInstructions({ autoSave: true }); /** * Signed-out path (auth-required mode). Deliberately NOT the proactive text: diff --git a/packages/mcp/src/login.ts b/packages/mcp/src/login.ts index d0b128d90..e1db9a9ac 100644 --- a/packages/mcp/src/login.ts +++ b/packages/mcp/src/login.ts @@ -22,7 +22,14 @@ import { randomBytes, timingSafeEqual } from "node:crypto"; import open from "open"; import type { MemWalCredentials } from "./auth.js"; -import { saveCreds, formatReplacementNotice, formatPendingSignInWarning } from "./auth.js"; +import { + saveCreds, + formatReplacementNotice, + formatPendingSignInWarning, + savePendingLogin, + clearPendingLogin, + reusablePendingLogin, +} from "./auth.js"; import { generateKeypair } from "./crypto.js"; import { log, note } from "./logger.js"; @@ -210,7 +217,47 @@ export async function loginFlow(opts: LoginOptions = {}): Promise { if (!creds) { error = new Error( - `Login timed out after ${cfg.timeoutMs}ms. If you already approved the wallet transaction, a delegate key may exist on-chain without local credentials. Remove unused keys from the dashboard, then run login again.`, + `Login timed out after ${cfg.timeoutMs}ms. If you already approved the wallet transaction, that delegate key is saved on disk: the next start reclaims it, and running login again cannot register it a second time. Only remove it from the dashboard if you mean to abandon it.`, ); server.close(); resolve(); @@ -431,6 +478,9 @@ export async function loginFlow(opts: LoginOptions = {}): Promise Promise | void, + /** Called when the flow fails after the caller has already been handed the + * URL and returned. Without this the failure is unobservable — see the + * catch below. */ onFailure?: (err: unknown) => void, ): InflightLogin { if (inflightLogin) return inflightLogin; @@ -513,24 +566,45 @@ export function startOrReuseLoginFlow( }, }); + // The one place a failed flow is reported. It has to be this catch rather + // than the `onSuccess` chain below: a caller that only passes `onFailure` + // still needs to hear about it, and reporting from both would emit the + // failure twice for one flow. result.catch((err) => { rejectUrl(err); + const msg = err instanceof Error ? err.message : String(err); + // This used to be a `warn` and nothing else, which made a background + // login failure invisible: `handleLocalLogin` has already returned the + // URL and told the client it succeeded, so without a signal here the + // agent waits forever on a flow that is already dead (WALM-332). Error + // level so it surfaces in client log views, `note` so it reads as a + // sentence on stderr, and `onFailure` so the caller can put it in + // front of the agent in-band. The pending write-ahead record is + // deliberately left in place — the next start reports the stranded key. + log.error("login.inflight.failed", { msg }); + note(`Walrus Memory sign-in failed: ${msg}`); try { onFailure?.(err); } catch { - /* caller errors don't break the flow */ + /* a reporting failure must not mask the original one */ } }); if (onSuccess) { - result + void result .then(async (creds) => { - await onSuccess(creds); + try { + await onSuccess(creds); + } catch (err) { + // A credential handoff that throws is its own failure, and + // distinct from the flow failing — the sign-in did work. + log.error("login.on_success_failed", { + msg: err instanceof Error ? err.message : String(err), + }); + } }) - .catch((err) => { - log.warn("login.inflight.failed", { - msg: err instanceof Error ? err.message : String(err), - }); + .catch(() => { + /* the flow's own rejection is reported in the catch above */ }); } diff --git a/packages/mcp/src/memory-policy.ts b/packages/mcp/src/memory-policy.ts new file mode 100644 index 000000000..2e187d1be --- /dev/null +++ b/packages/mcp/src/memory-policy.ts @@ -0,0 +1,84 @@ +/** + * Shared automatic-memory policy — the MCP client's copy. + * + * Consumed by `instructions.ts` (the MCP `instructions` field, which is the + * only proactive-usage channel that survives lazy tool loading) and by + * `auth-required.ts` (the cold-start / signed-out `tools/list`). + * + * WALM-642: before this module the secret rules did not exist at all, and the + * save guidance was written out three times — here, in the plugin hooks, and in + * the relayer sidecar's tool descriptions — with nothing tying them together. + * A preference stated next to a password was forwarded whole. + */ + +// ─── memwal:policy-block:start ─────────────────────────────────────────────── +// WALM-642. The lines between these two markers are BYTE-IDENTICAL in three +// files that cannot import one another, because the three packages have no +// workspace link: +// +// packages/mcp/src/memory-policy.ts — MCP client: initialize +// instructions + the +// cold-start tools/list +// packages/mcp/plugin/scripts/lib/memory-policy.mjs — plugin hooks: the +// guidance injected at +// SessionStart / +// UserPromptSubmit / +// PostToolUse +// services/server/scripts/mcp/tools/memory-policy.ts — relayer sidecar: the +// live tool descriptions +// +// The duplication is deliberate and pinned: `memory-policy-sync` tests on both +// sides extract this block from each file and compare the bytes, so editing one +// copy fails the suite until the other two match. Edit the block, then copy it +// verbatim — markers included — into the other two files. + +/** + * The secret-exclusion and do-not-save rules, stated verbatim by every + * automatic-save surface. + * + * These are model-facing rules, not enforcement. The programmatic backstop is + * the redactor in the relayer sidecar's write path + * (services/server/scripts/mcp/tools/redaction.ts), which runs before any text + * reaches the SDK. + */ +export const SECRET_EXCLUSION_RULES = [ + "NEVER save a credential, even when it sits next to something worth saving: passwords,", + "API keys, access or refresh tokens, private keys, seed or recovery phrases, authorization", + "headers, session cookies, and connection strings or URLs that embed a user:password.", + "When a message mixes a preference with a credential, save the preference alone and leave", + "the credential out; never store the line verbatim.", + "If the user says not to save something ('don't save this', 'off the record', or the same", + "in any language), do not save it, and do not save a paraphrase of it either.", + "Do not store quoted or pasted third-party material — log excerpts, code, articles, other", + "people's messages — as if it were a fact about this user. Save only what the user is", + "telling you about themselves or their work, in your own words.", +].join(" "); + +/** + * One-line form, for surfaces with no room for the full block (a per-turn + * nudge, a tool description tail). It is a reminder of the block above, never a + * replacement for it: any surface that drives an automatic save states the full + * `SECRET_EXCLUSION_RULES`. + */ +export const SECRET_EXCLUSION_SUMMARY = [ + "Never save passwords, keys, tokens or other credentials — not even beside a fact worth", + "saving; honour an explicit 'do not save this'; never store pasted third-party content as", + "a fact about the user.", +].join(" "); + +/** + * Whether to save unprompted is the user's standing choice, and this is the + * sentence that says so. A direct request ("remember that ...") is never gated + * by it — the gate is only on saving something the user did not ask you to save. + */ +export const AUTO_SAVE_OPT_IN_RULE = [ + "Whether to save things the user did not ask you to save is their standing choice, made once", + "in a terminal. When automatic memory is on, save durable facts as they state them; when it is", + "off, save only what they ask you to save in that turn. That question is put by `memwal-mcp", + "login` and set by `memwal-mcp auto-save on|off` — never ask the user to answer it in chat,", + "and never answer it on their behalf.", +].join(" "); + +/** Bumped whenever the text above changes, so a stale copy is identifiable. */ +export const MEMORY_POLICY_VERSION = "2026-09-17.2"; +// ─── memwal:policy-block:end ───────────────────────────────────────────────── diff --git a/packages/mcp/src/messages.ts b/packages/mcp/src/messages.ts index 11e7d6176..c9143a402 100644 --- a/packages/mcp/src/messages.ts +++ b/packages/mcp/src/messages.ts @@ -117,6 +117,11 @@ export function loginSuccessNotification(info: LoginSuccessInfo): string { * * `reason` null — no attempt on record — yields the empty string, so callers * can prefix unconditionally. + * + * The advice splits on whether the user approved the wallet step. If they did, + * the key is on-chain and in the write-ahead record (WALM-332), so a restart + * reclaims it. Signing in again would reuse that same key, and the dashboard's + * `add_delegate_key` aborts on a key that is already registered. */ export function loginFailureNotice(reason: string | null): string { if (!reason) return ""; @@ -125,10 +130,12 @@ export function loginFailureNotice(reason: string | null): string { "", `Reason: ${reason}`, "", - "The unused key from this attempt may already be registered on your account. Remove it", - "from the dashboard if you are not using it. Sign in again and open the new link", - "straight away. A retry only helps once the MCP client is left running through the", - "wallet prompt.", + "If you approved the wallet step, that key is registered and saved on this machine.", + "Restart the MCP client within 24 hours and it is reclaimed. Signing in again cannot", + "register the same key twice, and removing it from the dashboard abandons it.", + "", + "If you did not approve it, sign in again and open the new link straight away. A", + "retry only helps once the MCP client is left running through the wallet prompt.", "", "---", "", diff --git a/packages/mcp/src/recovery.ts b/packages/mcp/src/recovery.ts new file mode 100644 index 000000000..f24bdc809 --- /dev/null +++ b/packages/mcp/src/recovery.ts @@ -0,0 +1,300 @@ +/** + * Recovery for a login that was interrupted after the browser registered our + * delegate key on-chain but before the callback could save it (WALM-332). + * + * `loginFlow` write-aheads the keypair to `login-pending.json` before the + * browser can act, so the key itself survives losing the process. What does + * not survive is the metadata the callback would have carried — `accountId`, + * `walletAddress`, `packageId` — and `credentials.json` is not loadable + * without them. `GET /api/whoami` closes that gap: the relayer resolves the + * account from the delegate key during authentication anyway, so it can hand + * back the identity the key already proves. + */ +import { randomUUID, createHash } from "node:crypto"; + +import type { MemWalCredentials } from "./auth.js"; +import { + loadCreds, + saveCreds, + loadPendingLogin, + clearPendingLogin, + formatReplacementNotice, +} from "./auth.js"; +import { signMessage } from "./crypto.js"; +import { log } from "./logger.js"; + +export type RecoveryOutcome = + /** Nothing was pending. The overwhelmingly common case. */ + | "no-pending" + /** Key was registered; `credentials.json` has been rebuilt from it. */ + | "recovered" + /** A newer sign-in already happened — the pending key is stale. */ + | "superseded" + /** Relayer would not authenticate the key. Deliberately non-destructive. */ + | "rejected" + /** Relayer unreachable or erroring. Record kept for a later attempt. */ + | "unavailable"; + +export interface RecoveryResult { + outcome: RecoveryOutcome; + /** Set when a key may be registered on-chain but is not usable locally. */ + strandedPublicKey?: string; + credentials?: MemWalCredentials; + replacementNotice?: string; +} + +/** Same wall-clock budget as a normal cold-start probe: recovery must never + * be the reason a client hangs at startup. */ +const WHOAMI_TIMEOUT_MS = 10_000; + +/** `x-auth-error` value the relayer sets when it could not reach Sui to check + * the key at all. Named in `services/server/src/auth.rs`. */ +const AUTH_UPSTREAM_UNAVAILABLE = "AUTH_UPSTREAM_UNAVAILABLE"; + +interface WhoamiResponse { + account_id: string; + owner: string; + package_id: string; +} + +function isWhoami(o: unknown): o is WhoamiResponse { + if (!o || typeof o !== "object") return false; + const w = o as Record; + return ( + typeof w.account_id === "string" && + /^0x[0-9a-fA-F]{64}$/.test(w.account_id) && + typeof w.owner === "string" && + typeof w.package_id === "string" + ); +} + +/** + * Build the exact string the relayer will rebuild and verify against. + * + * `services/server/src/auth.rs` calls itself the single source of truth for + * this format, and it is reproduced here rather than imported because the two + * live in different languages. That duplication is the risk: get it subtly + * wrong — a trimmed trailing separator, a missing empty field — and every + * recovery attempt fails with an opaque 401 that no type checker would have + * caught. Exported so a test can pin it against the identical literal asserted + * in `routes::accounts::tests::whoami_recovery_request_canonical_message_is_stable`. + * + * `accountId` is empty for recovery: not knowing it is the reason we are here, + * and the server defaults its hint to `""` when the header is absent. + */ +export function canonicalRequestMessage(parts: { + timestamp: string; + method: string; + path: string; + bodyHash: string; + nonce: string; + accountId?: string; +}): string { + const { timestamp, method, path, bodyHash, nonce, accountId = "" } = parts; + return `${timestamp}.${method}.${path}.${bodyHash}.${nonce}.${accountId}`; +} + +/** sha256 of an empty body. A GET sends none; the server hashes it anyway. */ +export const EMPTY_BODY_SHA256 = createHash("sha256").update("").digest("hex"); + +/** + * Ask the relayer who this delegate key belongs to. + * + * The account id is signed as an empty string and its header omitted, because + * not knowing it is the entire reason we are here. The server defaults the + * hint to `""` when the header is absent, so both sides build the same + * canonical message. + * + * Returns null only when the request never produced a response at all. + */ +async function whoami( + relayerUrl: string, + privateKeyHex: string, + publicKeyHex: string, +): Promise<{ status: number; body: unknown; authError: string | null } | null> { + const path = "/api/whoami"; + // SECONDS. `services/server/src/auth.rs` freshness-checks `x-timestamp` + // against `chrono::Utc::now().timestamp()` within a drift window of a few + // minutes, so a millisecond value (~10^12) is always outside it and every + // request 401s with ERR_TIMESTAMP_OUT_OF_BOUNDS. + const timestamp = Math.floor(Date.now() / 1000).toString(); + const nonce = randomUUID(); + const message = canonicalRequestMessage({ + timestamp, + method: "GET", + path, + bodyHash: EMPTY_BODY_SHA256, + nonce, + }); + const signature = await signMessage(privateKeyHex, message); + + const controller = new AbortController(); + const timer = setTimeout(() => controller.abort(), WHOAMI_TIMEOUT_MS); + timer.unref?.(); + try { + const resp = await fetch(`${relayerUrl.replace(/\/+$/, "")}${path}`, { + method: "GET", + headers: { + "x-public-key": publicKeyHex, + "x-signature": signature, + "x-timestamp": timestamp, + "x-nonce": nonce, + }, + signal: controller.signal, + }); + const text = await resp.text(); + let body: unknown = null; + try { + body = JSON.parse(text); + } catch { + /* non-JSON error page — status is what matters */ + } + return { + status: resp.status, + body, + authError: resp.headers.get("x-auth-error"), + }; + } catch { + return null; + } finally { + clearTimeout(timer); + } +} + +/** + * Attempt to turn a stranded pending login into usable credentials. + * + * Never throws, and never deletes a record that might still be recoverable — + * a stranded key is the user's paid-for property, and the cost of keeping it + * around until its TTL is a file that grants nothing. + */ +export async function recoverPendingLogin(): Promise { + const pending = loadPendingLogin(); + if (!pending) return { outcome: "no-pending" }; + + // Ordering guard. If a later sign-in already succeeded, its credentials + // are the user's current intent and must not be rolled back to an older + // stranded key. Note this compares against the *pending* record's start + // time, so a login begun after the last successful one still wins. + const existing = loadCreds(); + if (existing && Date.parse(existing.createdAt) >= Date.parse(pending.createdAt)) { + log.warn("login.pending.superseded", { + publicKey: pending.delegatePublicKeyHex, + }); + clearPendingLogin(); + return { outcome: "superseded", strandedPublicKey: pending.delegatePublicKeyHex }; + } + + const res = await whoami( + pending.relayerUrl, + pending.delegatePrivateKey, + pending.delegatePublicKeyHex, + ); + + if (res === null) { + log.warn("login.pending.relayer_unreachable", { + publicKey: pending.delegatePublicKeyHex, + }); + return { outcome: "unavailable", strandedPublicKey: pending.delegatePublicKeyHex }; + } + + if (res.status !== 200 || !isWhoami(res.body)) { + // `rejected` is reserved for the relayer actually denying this + // identity, because that is the only outcome whose advice — sign in + // again, after removing the key from the dashboard if it was already + // registered — is worth giving. During a transient upstream + // failure that advice is worse than silence: the key is still good, and + // `unavailable` correctly says the next start retries it with no action + // from the user. + // + // So only 401/403 is a denial. A 503 carrying + // `x-auth-error: AUTH_UPSTREAM_UNAVAILABLE` is Sui RPC being down, and + // 429 / 5xx / a 404 from a relayer too old to serve this route are all + // "ask again later" — as is a 200 whose body is not a whoami, which + // means we are not talking to the endpoint we think we are. + const denied = + (res.status === 401 || res.status === 403) && + res.authError !== AUTH_UPSTREAM_UNAVAILABLE; + const outcome: RecoveryOutcome = denied ? "rejected" : "unavailable"; + // Non-destructive either way. A 401 is ambiguous even when it IS a + // denial: on testnet the registry scan is disabled outright and a + // genuinely registered key is refused for want of an x-account-id hint + // (services/server/src/auth.rs — "x-account-id is required for + // delegate-key authentication on testnet"). Clearing here would destroy + // a recoverable key in exactly that environment, so the record is + // always left for its TTL to retire. + log.warn(`login.pending.${outcome}`, { + publicKey: pending.delegatePublicKeyHex, + status: res.status, + authError: res.authError, + }); + return { outcome, strandedPublicKey: pending.delegatePublicKeyHex }; + } + + const creds: MemWalCredentials = { + delegatePrivateKey: pending.delegatePrivateKey, + delegatePublicKeyHex: pending.delegatePublicKeyHex, + delegateAddress: pending.delegateAddress, + walletAddress: res.body.owner, + accountId: res.body.account_id, + packageId: res.body.package_id, + relayerUrl: pending.relayerUrl, + label: pending.label, + createdAt: new Date().toISOString(), + version: 1, + }; + const saved = saveCreds(creds); + clearPendingLogin(); + log.info("login.pending.recovered", { + accountId: creds.accountId, + delegateAddress: creds.delegateAddress, + }); + return { + outcome: "recovered", + credentials: creds, + replacementNotice: formatReplacementNotice(saved, creds.accountId) ?? undefined, + }; +} + +/** + * The line to show the user when a stranded key could not be reclaimed. + * + * Names the key so the user can find the registration they paid for in the + * dashboard. Signing in again reuses this key, and the dashboard's + * `add_delegate_key` aborts on one that is already registered, so a key the + * user approved has to be removed there before a new sign-in can finish. A + * user told only "login failed" can do neither. + */ +export function formatStrandedLoginNotice(result: RecoveryResult): string | null { + if (!result.strandedPublicKey) return null; + if (result.outcome === "recovered" || result.outcome === "no-pending") return null; + + const key = result.strandedPublicKey; + const lines = [ + `⚠️ Your last Walrus Memory sign-in did not finish.`, + ``, + `A delegate key may have been registered on-chain without being saved locally:`, + ` ${key}`, + ``, + ]; + if (result.outcome === "superseded") { + lines.push( + `You have since signed in again, so your current credentials are fine.`, + `Revoke the key above from the dashboard if you don't recognise it.`, + ); + } else if (result.outcome === "unavailable") { + lines.push( + `The relayer could not be reached to check. This will be retried on the`, + `next start — no action needed yet.`, + ); + } else { + lines.push( + `The relayer did not accept it. If you never approved the wallet step, run`, + `\`memwal_login\`: it reuses this key. If you did, remove the key above from`, + `the dashboard first and then run \`memwal_login\`, because the wallet step`, + `cannot register a key that is already there. This is expected on Testnet,`, + `where the relayer cannot confirm a registered key at start.`, + ); + } + return lines.join("\n"); +} diff --git a/packages/mcp/src/streamable.ts b/packages/mcp/src/streamable.ts new file mode 100644 index 000000000..b203ca2bd --- /dev/null +++ b/packages/mcp/src/streamable.ts @@ -0,0 +1,178 @@ +/** + * Streamable HTTP transport for the stdio bridge. + * + * The legacy transport splits a call in two: POST to `/api/mcp/messages`, then + * wait for the reply to arrive on a separate `/api/mcp/sse` stream. Everything + * expensive in `bridge.ts` follows from that split — the `inFlight` map, the + * `sent` flag, the 404-means-never-ran reset, the idle watchdog, and the + * replay-on-reconnect path — because a POST can succeed while its reply is + * lost, and the bridge cannot tell that from a call still running. + * + * Streamable HTTP (MCP 2025-06) collapses that: one endpoint, and the reply + * comes back on the same request. The relayer has served it since + * `mcp_proxy.rs:751` ("Single endpoint that supersedes the SSE+POST split"); + * only the bridge was still on the old transport. + * + * This module wraps the MCP SDK's own client transport rather than hand-rolling + * the protocol: session-id round-tripping, the optional SSE upgrade on a POST + * response, and resumption tokens are all spec details that are easy to get + * subtly wrong and that the SDK already implements. + */ +import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js"; +import type { JSONRPCMessage } from "@modelcontextprotocol/sdk/types.js"; + +import type { MemWalCredentials } from "./auth.js"; +import { log } from "./logger.js"; + +/** Which relayer transport the bridge dials. */ +export type TransportKind = "sse" | "http"; + +/** + * Resolve the transport from `MEMWAL_MCP_TRANSPORT`. + * + * Defaults to `sse`, the transport every released bridge has used. Streamable + * HTTP is opt-in until it has production mileage: this runs on users' machines + * against their real memories, so the new path proves itself before it becomes + * the one that runs by default. + * + * An unrecognised value falls back rather than throwing — a typo in a user's + * MCP config must not stop their memory from working. + */ +export function resolveTransport(raw: string | undefined): TransportKind { + const value = raw?.trim().toLowerCase(); + if (!value) return "sse"; + if (value === "http" || value === "streamable" || value === "streamable-http") { + return "http"; + } + if (value === "sse") return "sse"; + log.warn("bridge.transport_unrecognized", { value, using: "sse" }); + return "sse"; +} + +/** + * The Streamable HTTP endpoint for a relayer base URL. + * + * `/api/mcp` — the same base the SSE transport hangs `/api/mcp/sse` and + * `/api/mcp/messages` off, minus the split. + */ +export function streamableUrl(relayerUrl: string): string { + return `${relayerUrl.replace(/\/+$/, "")}/api/mcp`; +} + +/** + * A live relayer session. Deliberately the same shape the SSE handshake + * returns, so `runBridge` can hold either without branching on transport + * everywhere it forwards a message. + */ +export interface RelaySession { + /** Endpoint this session talks to. Logging only. */ + postUrl: string; + /** + * Forward one JSON-RPC message. Resolves with an HTTP-ish status the + * caller can act on: 200 for accepted, 404 when the relayer says the + * session does not exist (the message provably did not run, so the + * caller may retry it without risking a duplicate write). + */ + send( + msg: JSONRPCMessage, + /** Unused here — headers are bound when the session opens. Present so + * this matches the SSE handshake's `send`, which signs per POST. */ + creds?: MemWalCredentials, + extra?: Record, + ): Promise; + /** Incoming messages from the relayer. */ + iter: AsyncIterator; + /** Tear the session down. */ + abort: () => void; +} + +/** HTTP status carried on the SDK's transport error, when it has one. */ +function statusOf(err: unknown): number { + const code = (err as { code?: unknown } | null)?.code; + return typeof code === "number" ? code : 0; +} + +export async function openStreamableSession( + relayerUrl: string, + creds: MemWalCredentials, + extraHeaders: Record = {}, +): Promise { + const url = streamableUrl(relayerUrl); + + // Queue + waiter rather than an event emitter, so a message that arrives + // before `runBridge` pulls from the iterator is buffered instead of + // dropped. The SSE path does the same thing for the same reason. + const queue: JSONRPCMessage[] = []; + let wake: (() => void) | null = null; + let closed = false; + const push = (msg: JSONRPCMessage) => { + queue.push(msg); + const resume = wake; + wake = null; + resume?.(); + }; + const finish = () => { + closed = true; + const resume = wake; + wake = null; + resume?.(); + }; + + const transport = new StreamableHTTPClientTransport(new URL(url), { + requestInit: { + headers: { + authorization: `Bearer ${creds.delegatePrivateKey}`, + "x-memwal-account-id": creds.accountId, + ...extraHeaders, + }, + }, + }); + + transport.onmessage = push; + transport.onclose = finish; + transport.onerror = (err) => { + log.warn("bridge.streamable_error", { err: String(err) }); + // Not `finish()` — the SDK transport reconnects its own stream, and + // tearing the session down on a transient read error is what the SSE + // watchdog did wrong. + }; + + await transport.start(); + + const iter: AsyncIterator = { + async next() { + while (queue.length === 0) { + if (closed) return { value: undefined as never, done: true }; + await new Promise((resolve) => (wake = resolve)); + } + return { value: queue.shift()!, done: false }; + }, + }; + + return { + postUrl: url, + async send(msg) { + try { + await transport.send(msg); + return 200; + } catch (err) { + const status = statusOf(err); + log.warn("bridge.streamable_send_failed", { + status, + err: String(err), + }); + // Surface the status rather than throwing: `postIfCurrent` + // routes on it, and a 404 specifically means the message was + // discarded rather than run. + return status; + } + }, + iter, + abort: () => { + void transport.close().catch(() => { + /* already gone */ + }); + finish(); + }, + }; +} diff --git a/packages/mcp/test/auto-save-optin.test.mjs b/packages/mcp/test/auto-save-optin.test.mjs new file mode 100644 index 000000000..0bc57a6b1 --- /dev/null +++ b/packages/mcp/test/auto-save-optin.test.mjs @@ -0,0 +1,571 @@ +/** + * Automatic saving is ON, once a human has been asked (WALM-642). + * + * The thing being governed is narrow and worth naming: saving something the + * user did not ask to have saved. A direct request ("remember that ...") is not + * gated, and neither is recall — so these tests check both that the guidance + * goes quiet when the answer is "off" AND that nothing else goes quiet with it. + * + * Three states, and the two unset ones are the interesting half: an install + * that predates the consent prompt keeps saving (that is its status quo, not a + * new grant), while one created after it saves nothing until someone answers. + * + * Both halves of the opt-in are covered, because they are two separate + * implementations of the same rule: the TypeScript one the MCP server reads, + * and the plain-`.mjs` one the plugin hooks read. The hooks are spawned by the + * client, not by this package, so they cannot share the compiled module and a + * divergence between the two would be invisible. + */ +import { test } from "node:test"; +import assert from "node:assert/strict"; +import { spawnSync } from "node:child_process"; +import { Readable, Writable } from "node:stream"; +import { mkdtempSync, mkdirSync, readFileSync, statSync, writeFileSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { dirname, join, resolve } from "node:path"; +import { fileURLToPath } from "node:url"; + +import { + autoSaveStatus, + isAutoSaveEnabled, + markConsentPending, + parseBooleanSetting, + publishHookState, + setAutoSave, + settingsPath, + hookStatePath, + AUTO_SAVE_ENV, +} from "../dist/auto-save.js"; +import { + askAutoSaveConsent, + interpretConsentAnswer, + CONSENT_PROMPT, +} from "../dist/consent.js"; +import * as hookAutoSave from "../plugin/scripts/lib/auto-save.mjs"; +import { parseArgs, helpText } from "../dist/index.js"; +import { TOOL_DEFINITIONS } from "../dist/auth-required.js"; + +const __dirname = dirname(fileURLToPath(import.meta.url)); +const SCRIPTS = resolve(__dirname, "../plugin/scripts"); + +/** A fresh, empty `.memwal` directory, so no developer's real settings leak in. */ +function freshCredsDir() { + const dir = mkdtempSync(join(tmpdir(), "memwal-autosave-")); + return dir; +} + +/** + * Write an answer AND publish the resolved state, which is what the server + * does on every write and every start-up. The hooks read only the published + * state now (WALM-642), so a test that wrote settings.json alone would be + * testing a machine whose MCP server had never run. + */ +function writeSettings(dir, value) { + mkdirSync(dir, { recursive: true }); + writeFileSync(join(dir, "settings.json"), JSON.stringify(value)); + withEnv({ MEMWAL_CREDS_DIR: dir, [AUTO_SAVE_ENV]: undefined }, () => + publishHookState(), + ); +} + +/** Run one hook with a controlled environment and return its injected text. */ +function runHook(script, input, env) { + const result = spawnSync(process.execPath, [join(SCRIPTS, script)], { + input: JSON.stringify(input), + encoding: "utf8", + env: { ...process.env, MEMWAL_AUTO_SAVE: "", ...env }, + }); + assert.equal(result.status, 0, result.stderr); + if (!result.stdout.trim()) return ""; + return JSON.parse(result.stdout).hookSpecificOutput?.additionalContext ?? ""; +} + +/** Run `fn` with env vars set, restoring them afterwards. */ +function withEnv(vars, fn) { + const saved = {}; + for (const [k, v] of Object.entries(vars)) { + saved[k] = process.env[k]; + if (v === undefined) delete process.env[k]; + else process.env[k] = v; + } + try { + return fn(); + } finally { + for (const [k, v] of Object.entries(saved)) { + if (v === undefined) delete process.env[k]; + else process.env[k] = v; + } + } +} + +// ── the resolver ──────────────────────────────────────────────────────────── + +test("a new install saves nothing until someone answers", () => { + const dir = freshCredsDir(); + withEnv({ MEMWAL_CREDS_DIR: dir, [AUTO_SAVE_ENV]: undefined }, () => { + markConsentPending(); + assert.equal(isAutoSaveEnabled(), false); + assert.deepEqual(autoSaveStatus(), { + enabled: false, + state: "unset", + source: "unanswered", + pendingConsent: true, + path: join(dir, "settings.json"), + }); + }); + rmSync(dir, { recursive: true, force: true }); +}); + +test("an install that predates the prompt keeps saving, and is still asked", () => { + // No settings file at all, credentials on disk: someone who has been + // auto-saving since before this setting existed. Switching them off would + // be a regression dressed up as caution. + const dir = freshCredsDir(); + withEnv({ MEMWAL_CREDS_DIR: dir, [AUTO_SAVE_ENV]: undefined }, () => { + writeFileSync(join(dir, "credentials.json"), "{}"); + const status = autoSaveStatus(); + assert.equal(status.enabled, true); + assert.equal(status.state, "unset"); + assert.equal(status.source, "legacy"); + // Carried over, not granted — so the question is still owed. + assert.equal(status.pendingConsent, true); + }); + rmSync(dir, { recursive: true, force: true }); +}); + +test("the pending stamp stops a headless sign-in maturing into consent", () => { + // Without the stamp, a brand-new install that signs in through the + // `memwal_login` tool would be indistinguishable from a long-standing user + // the moment credentials appear, and would start saving with nobody ever + // having been asked. + const dir = freshCredsDir(); + withEnv({ MEMWAL_CREDS_DIR: dir, [AUTO_SAVE_ENV]: undefined }, () => { + markConsentPending(); + writeFileSync(join(dir, "credentials.json"), "{}"); + const status = autoSaveStatus(); + assert.equal(status.enabled, false, "consent by never being asked"); + assert.equal(status.source, "unanswered"); + assert.equal(status.pendingConsent, true); + }); + rmSync(dir, { recursive: true, force: true }); +}); + +test("an answer is read back, either way, and is never asked for again", () => { + const dir = freshCredsDir(); + withEnv({ MEMWAL_CREDS_DIR: dir, [AUTO_SAVE_ENV]: undefined }, () => { + markConsentPending(); + + setAutoSave(true); + assert.equal(isAutoSaveEnabled(), true); + assert.equal(autoSaveStatus().state, "on"); + assert.equal(autoSaveStatus().source, "settings"); + assert.equal(autoSaveStatus().pendingConsent, false); + + // Declining must cost nothing and must not be nagged at. + setAutoSave(false); + assert.equal(isAutoSaveEnabled(), false); + assert.equal(autoSaveStatus().state, "off"); + assert.equal( + autoSaveStatus().pendingConsent, + false, + "a declined answer must not put the question back", + ); + // Even with credentials present — the rule that keeps a pre-existing + // install saving must not resurrect a deliberate "no". + writeFileSync(join(dir, "credentials.json"), "{}"); + assert.equal(isAutoSaveEnabled(), false); + assert.equal(autoSaveStatus().pendingConsent, false); + }); + rmSync(dir, { recursive: true, force: true }); +}); + +test("the settings file is not world-readable and keeps unrelated keys", () => { + const dir = freshCredsDir(); + withEnv({ MEMWAL_CREDS_DIR: dir, [AUTO_SAVE_ENV]: undefined }, () => { + writeSettings(dir, { somethingElse: "keep me" }); + setAutoSave(true); + const path = settingsPath(); + assert.equal(statSync(path).mode & 0o777, 0o600); + const parsed = JSON.parse(readFileSync(path, "utf8")); + assert.equal(parsed.somethingElse, "keep me"); + assert.equal(parsed.autoSave, true); + }); + rmSync(dir, { recursive: true, force: true }); +}); + +test("the environment overrides every state, in both directions", () => { + const dir = freshCredsDir(); + withEnv({ MEMWAL_CREDS_DIR: dir, [AUTO_SAVE_ENV]: undefined }, () => { + // over an answered "off" / "on" + setAutoSave(false); + withEnv({ [AUTO_SAVE_ENV]: "1" }, () => { + assert.equal(isAutoSaveEnabled(), true); + assert.equal(autoSaveStatus().source, "env"); + assert.equal(autoSaveStatus().state, "off", "the stored answer is untouched"); + }); + setAutoSave(true); + withEnv({ [AUTO_SAVE_ENV]: "0" }, () => { + assert.equal(isAutoSaveEnabled(), false); + assert.equal(autoSaveStatus().source, "env"); + }); + }); + rmSync(dir, { recursive: true, force: true }); + + // over each unset state, and it settles the question too — someone who set + // this deliberately does not also need to be prompted. + const unanswered = freshCredsDir(); + withEnv({ MEMWAL_CREDS_DIR: unanswered, [AUTO_SAVE_ENV]: undefined }, () => { + markConsentPending(); + withEnv({ [AUTO_SAVE_ENV]: "1" }, () => { + assert.equal(isAutoSaveEnabled(), true); + assert.equal(autoSaveStatus().pendingConsent, false); + }); + }); + rmSync(unanswered, { recursive: true, force: true }); + + const legacy = freshCredsDir(); + withEnv({ MEMWAL_CREDS_DIR: legacy, [AUTO_SAVE_ENV]: undefined }, () => { + writeFileSync(join(legacy, "credentials.json"), "{}"); + withEnv({ [AUTO_SAVE_ENV]: "off" }, () => { + assert.equal(isAutoSaveEnabled(), false); + assert.equal(autoSaveStatus().pendingConsent, false); + }); + }); + rmSync(legacy, { recursive: true, force: true }); +}); + +test("an unreadable or unparseable value is not consent", () => { + // Every one of these means "I could not tell", and the safe reading of + // that is off — not on, and not a crash. + for (const raw of [undefined, "", " ", "maybe", "2", "ON!"]) { + assert.equal(parseBooleanSetting(raw), null, `"${raw}" should be unparseable`); + } + assert.equal(parseBooleanSetting("yes"), true); + assert.equal(parseBooleanSetting(" OFF "), false); + + const dir = freshCredsDir(); + withEnv({ MEMWAL_CREDS_DIR: dir, [AUTO_SAVE_ENV]: undefined }, () => { + writeFileSync(join(dir, "settings.json"), "{ not json"); + assert.equal(isAutoSaveEnabled(), false); + assert.equal(autoSaveStatus().state, "unset"); + assert.equal(autoSaveStatus().pendingConsent, true); + }); + rmSync(dir, { recursive: true, force: true }); +}); + +test("the hook reads the published state rather than resolving anything itself", () => { + // The two implementations used to be hand-written mirrors and drifted + // apart, which is how a repo file came to switch automatic memory on + // (WALM-642). There is one resolver now: this asserts the hook reports what + // the server published, on every state, and that publishing is what moves + // it. + const dir = freshCredsDir(); + withEnv({ MEMWAL_CREDS_DIR: dir, [AUTO_SAVE_ENV]: undefined }, () => { + assert.equal(hookAutoSave.hookStatePath(), hookStatePath()); + + markConsentPending(); + assert.equal(hookAutoSave.isAutoSaveEnabled(), false); + assert.equal(hookAutoSave.autoSaveStatus().source, "unanswered"); + + setAutoSave(true); + assert.equal(hookAutoSave.isAutoSaveEnabled(), isAutoSaveEnabled()); + assert.equal(hookAutoSave.isAutoSaveEnabled(), true); + assert.equal(hookAutoSave.autoSaveStatus().source, "settings"); + assert.equal(hookAutoSave.autoSaveStatus().pendingConsent, false); + + withEnv({ [AUTO_SAVE_ENV]: "off" }, () => { + assert.equal(hookAutoSave.isAutoSaveEnabled(), false); + assert.equal(hookAutoSave.autoSaveStatus().source, "env"); + }); + + // An answer the server has not published yet is not one the hook may + // act on: it has no way to tell a stale file from a current one, so + // "cannot tell" has to read as off. + rmSync(hookStatePath(), { force: true }); + assert.equal(hookAutoSave.isAutoSaveEnabled(), false); + assert.equal(hookAutoSave.autoSaveStatus().source, "unavailable"); + assert.equal(hookAutoSave.autoSaveStatus().pendingConsent, true); + assert.equal(isAutoSaveEnabled(), true, "the server's own answer is unchanged"); + }); + rmSync(dir, { recursive: true, force: true }); +}); + +test("a published state this build does not understand fails safe", () => { + const dir = freshCredsDir(); + withEnv({ MEMWAL_CREDS_DIR: dir, [AUTO_SAVE_ENV]: undefined }, () => { + setAutoSave(true); + assert.equal(hookAutoSave.isAutoSaveEnabled(), true); + for (const corrupt of [ + "{ not json", + JSON.stringify({ version: 99, enabled: true }), + JSON.stringify({ version: 1, enabled: "yes" }), + JSON.stringify({ version: 1 }), + "null", + ]) { + writeFileSync(hookStatePath(), corrupt); + assert.equal( + hookAutoSave.isAutoSaveEnabled(), + false, + `treated as consent: ${corrupt}`, + ); + assert.equal(hookAutoSave.autoSaveStatus().source, "unavailable"); + } + }); + rmSync(dir, { recursive: true, force: true }); +}); + +// ── the hooks ─────────────────────────────────────────────────────────────── + +test("SessionStart goes quiet about saving when the user answered no", () => { + const dir = freshCredsDir(); + writeSettings(dir, { autoSave: false }); + const off = runHook("on_session_start.mjs", {}, { MEMWAL_CREDS_DIR: dir }); + + assert.match(off, /Automatic memory is OFF/); + assert.match(off, /Save ONLY what the user asks you to save/); + // The unprompted-save instruction is the thing that must be gone. + assert.doesNotMatch(off, /do not ask whether to save it/i); + // ...and the things that must NOT be gone with it. + assert.match(off, /memwal_recall/); + assert.match(off, /memwal_restore/); + assert.match(off, /auto-save on/); + + writeSettings(dir, { autoSave: true }); + const on = runHook("on_session_start.mjs", {}, { MEMWAL_CREDS_DIR: dir }); + assert.match(on, /Automatic memory is ON/); + assert.match(on, /do not ask whether to save it/i); + + // The rules ride on both. + for (const text of [off, on]) { + assert.match(text, /NEVER save a credential/); + } + rmSync(dir, { recursive: true, force: true }); +}); + +test("UserPromptSubmit injects a save-nothing rubric when the user answered no", () => { + const dir = freshCredsDir(); + writeSettings(dir, { autoSave: false }); + const prompt = "I always use pnpm and my staging canary is coral-fox-77."; + + const off = runHook( + "on_user_prompt.mjs", + { prompt, session_id: `off-${Math.random().toString(16).slice(2)}` }, + { MEMWAL_CREDS_DIR: dir }, + ); + assert.match(off, /Automatic saving is OFF/); + assert.match(off, /save ONLY what they ask you to save/); + assert.doesNotMatch(off, /call memwal_remember \(or memwal_remember_bulk for several\)/); + // Recall stays on: it reads, it does not write. + assert.match(off, /call memwal_recall first/); + assert.match(off, /NEVER save a credential/); + + writeSettings(dir, { autoSave: true }); + const on = runHook( + "on_user_prompt.mjs", + { prompt, session_id: `on-${Math.random().toString(16).slice(2)}` }, + { MEMWAL_CREDS_DIR: dir }, + ); + assert.match(on, /call memwal_remember \(or memwal_remember_bulk for several\)/); + assert.match(on, /NEVER save a credential/); + rmSync(dir, { recursive: true, force: true }); +}); + +test("PostToolUse stops nudging a save after an error when the user answered no", () => { + const dir = freshCredsDir(); + writeSettings(dir, { autoSave: false }); + // Must trip `detectError` in lib/signals.mjs (a strong marker) and clear + // the hook's 50-character minimum, or the hook stays silent for reasons + // that have nothing to do with the opt-in. + const errorOutput = + "fatal: could not read from remote repository — please make sure you " + + "have the correct access rights and the repository exists."; + const input = { tool_name: "Bash", tool_response: { stdout: "", stderr: errorOutput } }; + + const off = runHook("on_post_tool.mjs", input, { MEMWAL_CREDS_DIR: dir }); + assert.match(off, /memwal_recall/); + assert.match(off, /Automatic saving is OFF/); + assert.doesNotMatch(off, /save the fix with memwal_remember/); + + writeSettings(dir, { autoSave: true }); + const on = runHook("on_post_tool.mjs", input, { MEMWAL_CREDS_DIR: dir }); + assert.match(on, /save the fix with memwal_remember/); + // Error output is where a credential most often is; say so at the nudge. + assert.match(on, /Never save passwords, keys, tokens/); + rmSync(dir, { recursive: true, force: true }); +}); + +test("a hook with saving off still exits 0 and never blocks the session", () => { + const dir = freshCredsDir(); + writeSettings(dir, { autoSave: false }); + for (const script of ["on_session_start.mjs", "on_user_prompt.mjs", "on_post_tool.mjs"]) { + const result = spawnSync(process.execPath, [join(SCRIPTS, script)], { + input: JSON.stringify({ prompt: "a reasonably long prompt about pnpm" }), + encoding: "utf8", + env: { ...process.env, MEMWAL_CREDS_DIR: dir, MEMWAL_AUTO_SAVE: "" }, + }); + assert.equal(result.status, 0, `${script}: ${result.stderr}`); + } + rmSync(dir, { recursive: true, force: true }); +}); + +// ── the surface a user actually turns it on from ──────────────────────────── + +// ── the consent question ──────────────────────────────────────────────────── + +test("the prompt names the consequence, leads with permanence, and hedges the redaction", () => { + // These are the three wording rules the change exists for, so they are + // asserted rather than left to a reviewer's memory. + assert.match(CONSENT_PROMPT, /writes it to your memory without asking each time/); + assert.match(CONSENT_PROMPT, /Saved memories are permanent/); + assert.match(CONSENT_PROMPT, /immutable/); + assert.match(CONSENT_PROMPT, /cannot\s+delete one that is already saved/); + assert.match(CONSENT_PROMPT, /safety net, not a\s+guarantee/); + // Permanence comes first among the bullets — it is the fact that changes + // the answer. + const bullets = CONSENT_PROMPT.split("\n").filter((l) => l.trim().startsWith("- ")); + assert.equal(bullets.length, 3); + assert.match(bullets[0], /permanent/); + // Both options are offered plainly; declining is not dressed as a warning. + assert.match(CONSENT_PROMPT, /\[1\] Save automatically/); + assert.match(CONSENT_PROMPT, /\[2\] Only save when I ask/); + assert.match(CONSENT_PROMPT, /auto-save on\|off/); +}); + +test("Enter takes option 1, and anything unrecognised re-asks rather than assuming", () => { + assert.equal(interpretConsentAnswer(""), true); + assert.equal(interpretConsentAnswer(" "), true); + assert.equal(interpretConsentAnswer("1"), true); + assert.equal(interpretConsentAnswer("2"), false); + // A typo is not an answer to a question about permanent storage. + for (const raw of ["y", "n", "3", "yes", "maybe", "11"]) { + assert.equal(interpretConsentAnswer(raw), null, `"${raw}" must re-ask`); + } +}); + +/** Drive the prompt with scripted lines and collect what it wrote. */ +async function runPrompt(lines, { isTTY = true } = {}) { + const written = []; + const input = Readable.from(lines.map((l) => `${l}\n`)); + const output = new Writable({ + write(chunk, _enc, cb) { + written.push(chunk.toString()); + cb(); + }, + }); + const answer = await askAutoSaveConsent({ input, output, isTTY }); + return { answer, output: written.join("") }; +} + +test("an answer is taken from the terminal, and a bad one is re-asked", async () => { + assert.equal((await runPrompt(["1"])).answer, true); + assert.equal((await runPrompt([""])).answer, true); + assert.equal((await runPrompt(["2"])).answer, false); + + const retried = await runPrompt(["banana", "2"]); + assert.equal(retried.answer, false); + assert.match(retried.output, /Please answer 1 or 2/); + // The question is put again, not assumed away. + assert.equal(retried.output.split("Your choice").length - 1, 2); +}); + +test("a closed stream is not an answer, and does not hang", async () => { + // Ctrl-D, a killed terminal, a closed pipe. Recording a choice here would + // be recording one the user never made. + const { answer } = await runPrompt([]); + assert.equal(answer, null); +}); + +test("the prompt refuses to run without a TTY", async () => { + // Belt and braces with main()'s own check: a prompt with nobody in front of + // it is a hang, and this one would hang an MCP server's startup. + const { answer, output } = await runPrompt(["1"], { isTTY: false }); + assert.equal(answer, null); + assert.equal(output, "", "nothing may be written to a non-interactive stream"); +}); + +test("`auto-save` parses as a subcommand, and a bare one only reports", () => { + // A bare `auto-save` must not be read as consent to turn it ON — the + // difference between reporting a setting and changing one. + assert.equal(parseArgs(["auto-save"]).autoSave, "status"); + assert.equal(parseArgs(["auto-save", "on"]).autoSave, "on"); + assert.equal(parseArgs(["auto-save", "off"]).autoSave, "off"); + assert.equal(parseArgs(["auto-save", "status"]).autoSave, "status"); + assert.equal(parseArgs(["--auto-save", "on"]).autoSave, "on"); + + // Not a flag, so it must not be reported as an unrecognised one. + assert.deepEqual(parseArgs(["auto-save", "on"]).unknown, []); + // ...and it must not collide with the other subcommand. + assert.equal(parseArgs(["login"]).autoSave, undefined); + assert.equal(parseArgs(["auto-save", "on"]).forceLogin, false); +}); + +test("a typo'd flag does not swallow the subcommand or its value", () => { + // `--typo` consumes one following token as its value — unless that token + // is a command. Without the exemption `memwal-mcp --typo auto-save on` + // would silently run the server instead. + const parsed = parseArgs(["--typo", "auto-save", "on"]); + assert.deepEqual(parsed.unknown, ["--typo"]); + assert.equal(parsed.autoSave, "on"); +}); + +test("--help tells the user the setting exists and that login asks for it", () => { + // The plugin install path never shows a terminal, so --help and the + // post-login summary are where the choice reaches a person. + const help = helpText(); + assert.match(help, /auto-save on\|off/); + assert.match(help, /login` asks/); + assert.match(help, /nothing is saved unprompted until you/); + assert.match(help, /MEMWAL_AUTO_SAVE/); + // Permanence is stated here too — it is the fact that changes the answer. + assert.match(help, /permanent/); + // And that credentials are stripped regardless of which way it is set. + assert.match(help, /Credentials/); + assert.match(help, /not a guarantee/); +}); + +// ── the non-interactive path ──────────────────────────────────────────────── + +test("a non-TTY run never prompts, never hangs, and says where things stand", () => { + // Every MCP client spawn lands here. The failure this guards against is not + // a wrong answer, it is a server that never finishes starting because + // something is waiting on a stdin no human is attached to. + const dir = freshCredsDir(); + const result = spawnSync( + process.execPath, + [resolve(__dirname, "../dist/bin/memwal-mcp.js"), "auto-save"], + { + // Piped, not inherited: `process.stdin.isTTY` is undefined here, + // exactly as it is under an MCP client. + input: "", + encoding: "utf8", + timeout: 10_000, + env: { ...process.env, MEMWAL_CREDS_DIR: dir, MEMWAL_AUTO_SAVE: "" }, + }, + ); + assert.equal(result.status, 0, result.stderr); + assert.notEqual(result.signal, "SIGTERM", "the process hung waiting on stdin"); + // Reports, does not ask. + assert.doesNotMatch(result.stderr, /Your choice/); + assert.doesNotMatch(result.stderr, /Save automatically/); + assert.match(result.stderr, /Automatic memory: (ON|OFF)/); + rmSync(dir, { recursive: true, force: true }); +}); + +test("consent is not reachable from anything the model can call", () => { + // The single most important constraint in this change: a model answering + // on the user's behalf is not consent. The question lives in consent.ts, + // is called only from main() behind `process.stdin.isTTY`, and must never + // appear in a tool list, a tool description or the instructions. + const toolSurfaces = [ + readFileSync(resolve(__dirname, "../dist/auth-required.js"), "utf8"), + readFileSync(resolve(__dirname, "../dist/instructions.js"), "utf8"), + readFileSync(resolve(__dirname, "../dist/bridge.js"), "utf8"), + ].join("\n"); + assert.doesNotMatch(toolSurfaces, /askAutoSaveConsent/); + assert.doesNotMatch(toolSurfaces, /CONSENT_PROMPT/); + assert.doesNotMatch(toolSurfaces, /Your choice \[1\/2\]/); + + // And no tool is named for it. + const names = TOOL_DEFINITIONS.map((t) => t.name); + assert.ok(!names.some((n) => /consent|auto_?save/i.test(n)), names.join(", ")); +}); diff --git a/packages/mcp/test/auto-save-repo-trust.test.mjs b/packages/mcp/test/auto-save-repo-trust.test.mjs new file mode 100644 index 000000000..53d128783 --- /dev/null +++ b/packages/mcp/test/auto-save-repo-trust.test.mjs @@ -0,0 +1,349 @@ +/** + * A repository cannot turn automatic memory on (WALM-642, review findings 1-2). + * + * Two holes, one root cause: the consent decision was being re-derived from + * whatever `.memwal` directory the working directory resolved to, by two + * separate implementations. + * + * 1. The hook-side resolver (`plugin/scripts/lib/auto-save.mjs`) still used + * the pre-WALM-639 presence rule — nearest project `credentials.json` + * wins, contents never parsed. A repo carrying a file whose entire content + * was `not even valid json` read as proof of a long-standing install and + * injected the full proactive-save rubric; adding a committed + * `settings.json` pinned `autoSave: true` outright and silenced the + * consent-pending warning too. + * 2. Server-side, `settingsPath()` followed `credsPath()`, so approving a + * project (WALM-639) moved the consent answer into the repo — where a + * recorded "no" was not consulted and a fresh directory fell through to + * "legacy, keep saving". + * + * THESE TESTS MUST NOT SET `MEMWAL_CREDS_DIR`. It is the trusted override: it + * short-circuits project resolution in both resolvers, which is exactly why the + * existing opt-in suite could not see either bug. Everything here drives a real + * HOME and a real working directory instead, and finding 1 is asserted by + * spawning the actual hook the client spawns. + */ +import { test } from "node:test"; +import assert from "node:assert/strict"; +import { spawnSync } from "node:child_process"; +import { mkdtempSync, mkdirSync, writeFileSync, rmSync, existsSync, realpathSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { dirname, join, resolve } from "node:path"; +import { fileURLToPath } from "node:url"; + +const __dirname = dirname(fileURLToPath(import.meta.url)); +const SCRIPTS = resolve(__dirname, "../plugin/scripts"); + +const PRIVATE_KEY = "c".repeat(64); + +/** + * A session id nothing has seen before. `firstTime()` in lib/hook-io.mjs keeps + * a marker per (name, session) under the OS temp dir, so a fixed id makes the + * FIRST run of a test inject the full rubric and every run after it the + * one-line nudge. + */ +function freshSession(tag) { + return `${tag}-${Date.now()}-${Math.random().toString(16).slice(2)}`; +} + +function makeCreds(overrides = {}) { + return { + delegatePrivateKey: PRIVATE_KEY, + delegatePublicKeyHex: "d".repeat(64), + delegateAddress: "0x" + "e".repeat(64), + walletAddress: "0x" + "f".repeat(64), + accountId: "0x" + "a".repeat(64), + packageId: "0x" + "1".repeat(64), + relayerUrl: "https://relayer.example", + createdAt: new Date(0).toISOString(), + version: 1, + ...overrides, + }; +} + +function writeJson(path, value) { + mkdirSync(dirname(path), { recursive: true }); + writeFileSync(path, typeof value === "string" ? value : JSON.stringify(value)); + return path; +} + +/** + * A throwaway HOME and repo, with `MEMWAL_CREDS_DIR` cleared. + * + * Canonicalised because `homedir()` and `process.cwd()` both report resolved + * paths and `/tmp` is a symlink on macOS — an uncanonicalised HOME makes the + * project walk's "stop at the home directory" test miss. + */ +function sandbox(t) { + const home = realpathSync(mkdtempSync(join(tmpdir(), "memwal-trust-home-"))); + const repo = realpathSync(mkdtempSync(join(tmpdir(), "memwal-trust-repo-"))); + // A `.git` marker, because that is what a checkout has and what the project + // walk stops at. + mkdirSync(join(repo, ".git"), { recursive: true }); + t.after(() => { + rmSync(home, { recursive: true, force: true }); + rmSync(repo, { recursive: true, force: true }); + }); + return { home, repo }; +} + +/** Run one hook the way a client does: its own process, a cwd, and a HOME. */ +function runHookIn(script, { cwd, home, input = {}, env = {} }) { + const result = spawnSync(process.execPath, [join(SCRIPTS, script)], { + cwd, + input: JSON.stringify(input), + encoding: "utf8", + env: { + ...process.env, + HOME: home, + USERPROFILE: home, + MEMWAL_CREDS_DIR: "", + MEMWAL_AUTO_SAVE: "", + ...env, + }, + }); + assert.equal(result.status, 0, `${script} exited ${result.status}: ${result.stderr}`); + if (!result.stdout.trim()) return ""; + return JSON.parse(result.stdout).hookSpecificOutput?.additionalContext ?? ""; +} + +// `MEMWAL_CREDS_DIR: ""` must actually clear the override, or every assertion +// below is vacuous. Node keeps an empty string in the child's environment, and +// both resolvers read it as unset because `""` is falsy — pinned here so a +// future change to that check cannot quietly hollow this file out. +test("an empty MEMWAL_CREDS_DIR does not stand in for a real one", async () => { + const previous = process.env.MEMWAL_CREDS_DIR; + process.env.MEMWAL_CREDS_DIR = ""; + try { + const hook = await import("../plugin/scripts/lib/auto-save.mjs"); + assert.ok( + hook.hookStatePath().startsWith(join(realpathSync(process.env.HOME ?? tmpdir()))) || + hook.hookStatePath().includes(".memwal"), + "an empty override must fall back to the home directory", + ); + } finally { + if (previous === undefined) delete process.env.MEMWAL_CREDS_DIR; + else process.env.MEMWAL_CREDS_DIR = previous; + } +}); + +// ── finding 1: the hook repro, on the real hook ───────────────────────────── + +test("a committed credentials file cannot turn the save rubric on", (t) => { + const { home, repo } = sandbox(t); + // The reporter's repro, byte for byte: the contents are never parsed, so + // the file did not even have to be credentials. + writeJson(join(repo, ".memwal", "credentials.json"), "not even valid json"); + // The user's own install has been asked and has not answered. + writeJson(join(home, ".memwal", "settings.json"), { autoSaveConsent: "pending" }); + + const injected = runHookIn("on_user_prompt.mjs", { + cwd: repo, + home, + input: { session_id: freshSession("s1"), prompt: "I prefer pnpm" }, + }); + + assert.doesNotMatch( + injected, + /call memwal_remember \(or memwal_remember_bulk for several\)/, + "a repo file injected the proactive-save rubric", + ); + assert.match(injected, /Automatic saving is OFF|save ONLY what they ask you to save/); + + // And the control: the same HOME from a clean directory says the same + // thing, which is the point — the repo changed nothing at all. + const clean = realpathSync(mkdtempSync(join(tmpdir(), "memwal-trust-clean-"))); + t.after(() => rmSync(clean, { recursive: true, force: true })); + const control = runHookIn("on_user_prompt.mjs", { + cwd: clean, + home, + input: { session_id: freshSession("s2"), prompt: "I prefer pnpm" }, + }); + assert.equal(injected, control, "the repo steered the hook away from the control"); +}); + +test("a committed settings.json cannot pin autoSave on, or silence the warning", (t) => { + const { home, repo } = sandbox(t); + writeJson(join(repo, ".memwal", "credentials.json"), "not even valid json"); + writeJson(join(repo, ".memwal", "settings.json"), { autoSave: true }); + writeJson(join(home, ".memwal", "settings.json"), { autoSaveConsent: "pending" }); + + const prompt = runHookIn("on_user_prompt.mjs", { + cwd: repo, + home, + input: { session_id: freshSession("s3"), prompt: "I prefer pnpm" }, + }); + assert.doesNotMatch( + prompt, + /call memwal_remember \(or memwal_remember_bulk for several\)/, + "a repo settings.json switched automatic saving on", + ); + + const start = runHookIn("on_session_start.mjs", { cwd: repo, home }); + assert.match(start, /Automatic memory is OFF/, "a repo settings.json flipped the banner"); + assert.doesNotMatch(start, /do not ask whether to save it/i); +}); + +test("the hook reads nothing from the working directory at all", (t) => { + // Stronger than the two repros: not "this particular file is ignored" but + // "there is no file a checkout can add". The published state is what + // decides, and it lives where a repo cannot write. + const { home, repo } = sandbox(t); + writeJson(join(home, ".memwal", "auto-save-state.json"), { + version: 1, + enabled: true, + state: "on", + source: "settings", + pendingConsent: false, + settingsPath: join(home, ".memwal", "settings.json"), + updatedAt: new Date().toISOString(), + }); + // Every shape the old resolver would have followed, all saying "off". + writeJson(join(repo, ".memwal", "credentials.json"), makeCreds()); + writeJson(join(repo, ".memwal", "settings.json"), { autoSave: false }); + writeJson(join(repo, ".memwal", "auto-save-state.json"), { + version: 1, + enabled: false, + state: "off", + source: "settings", + pendingConsent: false, + settingsPath: join(repo, ".memwal", "settings.json"), + updatedAt: new Date().toISOString(), + }); + + const injected = runHookIn("on_user_prompt.mjs", { + cwd: repo, + home, + input: { session_id: freshSession("s4"), prompt: "I prefer pnpm" }, + }); + assert.match( + injected, + /call memwal_remember \(or memwal_remember_bulk for several\)/, + "the repo overrode the user's own published answer", + ); +}); + +test("a hook with no published state saves nothing and still exits 0", (t) => { + const { home, repo } = sandbox(t); + assert.equal(existsSync(join(home, ".memwal", "auto-save-state.json")), false); + for (const script of ["on_session_start.mjs", "on_user_prompt.mjs", "on_post_tool.mjs"]) { + const text = runHookIn(script, { + cwd: repo, + home, + input: { session_id: freshSession("s5"), prompt: "a reasonably long prompt about pnpm" }, + }); + assert.doesNotMatch( + text, + /call memwal_remember \(or memwal_remember_bulk for several\)/, + `${script} saved on an unknown state`, + ); + } +}); + +// ── finding 2: approving a project must not move the consent answer ───────── + +test("approving a project does not re-enable a declined auto-save", async (t) => { + const { home, repo } = sandbox(t); + const previous = { + home: process.env.HOME, + profile: process.env.USERPROFILE, + credsDir: process.env.MEMWAL_CREDS_DIR, + cwd: process.cwd(), + }; + process.env.HOME = home; + process.env.USERPROFILE = home; + delete process.env.MEMWAL_CREDS_DIR; + delete process.env.MEMWAL_AUTO_SAVE; + process.chdir(repo); + t.after(() => { + process.chdir(previous.cwd); + process.env.HOME = previous.home; + process.env.USERPROFILE = previous.profile; + if (previous.credsDir === undefined) delete process.env.MEMWAL_CREDS_DIR; + else process.env.MEMWAL_CREDS_DIR = previous.credsDir; + }); + + // Resolved per call, but the modules are imported once per process, so a + // cache-busting query is what makes them see this HOME. + const stamp = `${Date.now()}-${Math.random()}`; + const auth = await import(`../dist/auth.js?walm642=${stamp}`); + const autoSave = await import(`../dist/auto-save.js?walm642=${stamp}`); + + writeJson(join(home, ".memwal", "credentials.json"), makeCreds()); + // The user answers "[2] Only save when I ask". + autoSave.setAutoSave(false); + const answerPath = autoSave.settingsPath(); + assert.equal(answerPath, join(home, ".memwal", "settings.json")); + assert.equal(autoSave.isAutoSaveEnabled(), false); + + // Later, in a team repo, they approve that repo's credentials — a decision + // about WHERE memory is written, and nothing else. + writeJson( + join(repo, ".memwal", "credentials.json"), + makeCreds({ accountId: "0x" + "b".repeat(64) }), + ); + const approved = auth.approveProjectCreds(); + assert.equal(approved.outcome, "approved"); + assert.equal(auth.credsPath(), join(repo, ".memwal", "credentials.json")); + + // The recorded "no" is still the answer... + assert.equal(autoSave.isAutoSaveEnabled(), false, "approval re-enabled a declined auto-save"); + assert.equal(autoSave.autoSaveStatus().source, "settings"); + assert.equal(autoSave.autoSaveStatus().state, "off"); + // ...read from the same place it was written, not from the repo. + assert.equal(autoSave.settingsPath(), answerPath); + assert.equal( + existsSync(join(repo, ".memwal", "settings.json")), + false, + "the consent answer was written into the repository", + ); + + // Writing an answer while a project is approved must not put one there + // either — that file is committable. + autoSave.setAutoSave(true); + autoSave.markConsentPending(); + assert.equal( + existsSync(join(repo, ".memwal", "settings.json")), + false, + "setAutoSave wrote the consent answer into the repository", + ); + assert.equal( + existsSync(join(repo, ".memwal", "auto-save-state.json")), + false, + "the published hook state was written into the repository", + ); + assert.equal(existsSync(join(home, ".memwal", "auto-save-state.json")), true); +}); + +test("a project credentials file is not evidence of a long-standing install", async (t) => { + // The legacy rule — "no answer, but credentials on disk, so keep saving" — + // has to be asked of the user's own install. Answered with a repo file it + // is just the presence rule again, wearing a different hat. + const { home, repo } = sandbox(t); + const previous = { home: process.env.HOME, profile: process.env.USERPROFILE, cwd: process.cwd() }; + process.env.HOME = home; + process.env.USERPROFILE = home; + delete process.env.MEMWAL_CREDS_DIR; + delete process.env.MEMWAL_AUTO_SAVE; + process.chdir(repo); + t.after(() => { + process.chdir(previous.cwd); + process.env.HOME = previous.home; + process.env.USERPROFILE = previous.profile; + }); + + const stamp = `${Date.now()}-${Math.random()}`; + const auth = await import(`../dist/auth.js?walm642b=${stamp}`); + const autoSave = await import(`../dist/auto-save.js?walm642b=${stamp}`); + + // No global credentials, no answer anywhere — a machine that has never used + // MemWal, opening a repo that carries an approved credentials file. + writeJson(join(repo, ".memwal", "credentials.json"), makeCreds()); + auth.approveProjectCreds(); + assert.equal(auth.credsPath(), join(repo, ".memwal", "credentials.json")); + + const status = autoSave.autoSaveStatus(); + assert.equal(status.enabled, false, "a repo file was read as a long-standing install"); + assert.equal(status.source, "unanswered"); + assert.equal(status.pendingConsent, true); +}); diff --git a/packages/mcp/test/codex-installer-ownership.test.mjs b/packages/mcp/test/codex-installer-ownership.test.mjs new file mode 100644 index 000000000..41b6ed024 --- /dev/null +++ b/packages/mcp/test/codex-installer-ownership.test.mjs @@ -0,0 +1,216 @@ +/** + * WALM-643: the Codex hooks installer must remove its own hooks and nothing + * else. + * + * Ownership used to be decided by substring — any command containing + * `on_user_prompt.mjs` (or another of our generic hook filenames) counted as + * ours — and the match then removed the whole hook group. A company's own + * `/company/security/on_user_prompt.mjs` therefore disappeared on install and + * on uninstall, and took its unrelated siblings in the same group with it. + * + * These tests drive the real installer against a ~/.codex/hooks.json seeded + * with foreign hooks, across install, reinstall and uninstall. + */ +import { test } from "node:test"; +import assert from "node:assert/strict"; +import { spawnSync } from "node:child_process"; +import { mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { dirname, join, resolve } from "node:path"; +import { fileURLToPath } from "node:url"; + +const __dirname = dirname(fileURLToPath(import.meta.url)); +const PLUGIN_ROOT = resolve(__dirname, "../plugin"); +const INSTALLER = join(PLUGIN_ROOT, "scripts", "install_codex_hooks.mjs"); +const TEMPLATE_FILE = join(PLUGIN_ROOT, "hooks", "codex-hooks.json"); + +const MARKER_KEY = "_memwal"; + +/** A foreign tool whose hook filename happens to collide with one of ours. */ +const FOREIGN_PROMPT_HOOK = { + type: "command", + command: 'node "/company/security/on_user_prompt.mjs"', + timeout: 30, +}; +const FOREIGN_SIBLING_HOOK = { + type: "command", + command: 'node "/company/security/enforce-policy.mjs"', + timeout: 30, +}; +const FOREIGN_GROUP = { + matcher: "*", + statusMessage: "Running company policy checks...", + hooks: [FOREIGN_PROMPT_HOOK, FOREIGN_SIBLING_HOOK], +}; +const FOREIGN_CONFIG = { + hooks: { + UserPromptSubmit: [structuredClone(FOREIGN_GROUP)], + SessionStart: [ + { + matcher: "startup", + hooks: [ + { + type: "command", + command: 'node "/company/security/on_session_start.mjs"', + }, + ], + }, + ], + }, +}; + +/** + * The command a pre-WALM-643 build wrote for this plugin directory: raw + * placeholder substitution over the template text. + */ +function legacyCommandFor(event) { + const raw = readFileSync(TEMPLATE_FILE, "utf8").replaceAll("${PLUGIN_ROOT}", PLUGIN_ROOT); + return JSON.parse(raw).hooks[event][0].hooks[0].command; +} + +function makeHome(t, config) { + const home = mkdtempSync(join(process.env.TMPDIR || tmpdir(), "codex-ownership-")); + t.after(() => rmSync(home, { recursive: true, force: true })); + mkdirSync(join(home, ".codex"), { recursive: true }); + if (config) { + writeFileSync(join(home, ".codex", "hooks.json"), JSON.stringify(config, null, 2) + "\n"); + } + return home; +} + +function runInstaller(home, ...args) { + const result = spawnSync(process.execPath, [INSTALLER, ...args], { + env: { ...process.env, HOME: home, USERPROFILE: home }, + encoding: "utf8", + }); + assert.equal(result.status, 0, `${result.stdout}\n${result.stderr}`); + return result; +} + +function readHooks(home) { + return JSON.parse(readFileSync(join(home, ".codex", "hooks.json"), "utf8")); +} + +/** Groups in `event` that hold at least one hook this installer claims. */ +function memwalGroups(config, event) { + return (config.hooks[event] || []).filter((entry) => + (entry.hooks || []).some((hook) => hook[MARKER_KEY] !== undefined) + ); +} + +function foreignGroups(config, event) { + return (config.hooks[event] || []).filter((entry) => + (entry.hooks || []).every((hook) => hook[MARKER_KEY] === undefined) + ); +} + +test("install keeps a foreign hook whose filename collides with ours", (t) => { + const home = makeHome(t, FOREIGN_CONFIG); + runInstaller(home); + const config = readHooks(home); + + assert.deepEqual(foreignGroups(config, "UserPromptSubmit"), [FOREIGN_GROUP]); + assert.deepEqual(foreignGroups(config, "SessionStart"), FOREIGN_CONFIG.hooks.SessionStart); + assert.equal(memwalGroups(config, "UserPromptSubmit").length, 1); + assert.equal(memwalGroups(config, "SessionStart").length, 1); +}); + +test("install marks its own hooks", (t) => { + const home = makeHome(t); + runInstaller(home); + const config = readHooks(home); + for (const event of ["SessionStart", "UserPromptSubmit", "PostToolUse"]) { + const groups = config.hooks[event]; + assert.equal(groups.length, 1, event); + for (const hook of groups[0].hooks) assert.equal(hook[MARKER_KEY], "memwal-plugin-hooks"); + } +}); + +test("reinstalling is idempotent and leaves foreign hooks alone", (t) => { + const home = makeHome(t, FOREIGN_CONFIG); + runInstaller(home); + const first = readHooks(home); + runInstaller(home); + const second = readHooks(home); + + assert.deepEqual(second, first); + for (const event of ["SessionStart", "UserPromptSubmit", "PostToolUse"]) { + assert.equal(memwalGroups(second, event).length, 1, `duplicate MemWal group in ${event}`); + } + assert.deepEqual(foreignGroups(second, "UserPromptSubmit"), [FOREIGN_GROUP]); +}); + +test("uninstall removes only our hooks and restores the file to its old state", (t) => { + const home = makeHome(t, FOREIGN_CONFIG); + runInstaller(home); + runInstaller(home, "--uninstall"); + assert.deepEqual(readHooks(home), FOREIGN_CONFIG); +}); + +test("a foreign sibling survives when our hook is removed from its group", (t) => { + // A group holding a foreign hook, our hook, and another foreign hook. Only + // the middle one is ours; the group and its settings must survive. + const ours = { + type: "command", + command: legacyCommandFor("UserPromptSubmit"), + timeout: 12, + }; + const shared = { + matcher: "*", + statusMessage: "Running company policy checks...", + hooks: [FOREIGN_PROMPT_HOOK, ours, FOREIGN_SIBLING_HOOK], + }; + const home = makeHome(t, { hooks: { UserPromptSubmit: [shared] } }); + + runInstaller(home, "--uninstall"); + const config = readHooks(home); + assert.deepEqual(config.hooks.UserPromptSubmit, [ + { + matcher: "*", + statusMessage: "Running company policy checks...", + hooks: [FOREIGN_PROMPT_HOOK, FOREIGN_SIBLING_HOOK], + }, + ]); +}); + +test("an install predating the ownership marker is still replaced, not duplicated", (t) => { + const legacy = { + hooks: { + UserPromptSubmit: [ + { + hooks: [ + { + type: "command", + command: legacyCommandFor("UserPromptSubmit"), + timeout: 12, + }, + ], + }, + ], + }, + }; + const home = makeHome(t, legacy); + + runInstaller(home); + const config = readHooks(home); + assert.equal(config.hooks.UserPromptSubmit.length, 1); + assert.equal(config.hooks.UserPromptSubmit[0].hooks.length, 1); + assert.equal( + config.hooks.UserPromptSubmit[0].hooks[0][MARKER_KEY], + "memwal-plugin-hooks" + ); + + runInstaller(home, "--uninstall"); + assert.deepEqual(readHooks(home), { hooks: {} }); +}); + +test("an emptied group is dropped, but an event keeping foreign groups is not", (t) => { + const home = makeHome(t, FOREIGN_CONFIG); + runInstaller(home); + runInstaller(home, "--uninstall"); + const config = readHooks(home); + + assert.deepEqual(Object.keys(config.hooks).sort(), ["SessionStart", "UserPromptSubmit"]); + // PostToolUse held only our group, so the event is gone entirely. + assert.equal(config.hooks.PostToolUse, undefined); +}); diff --git a/packages/mcp/test/codex-installer-shell-safety.test.mjs b/packages/mcp/test/codex-installer-shell-safety.test.mjs new file mode 100644 index 000000000..cc8812ca3 --- /dev/null +++ b/packages/mcp/test/codex-installer-shell-safety.test.mjs @@ -0,0 +1,215 @@ +/** + * WALM-641: the Codex hooks installer must not let the install path reach a + * shell as syntax. + * + * The installer substitutes its own directory into the hook commands it writes + * to ~/.codex/hooks.json. That substitution used to run over the template + * *text* before JSON.parse, so a plugin directory containing `$(...)`, a + * backtick, quotes or a backslash landed unescaped in both the JSON document + * and the generated shell command -- running the hook executed whatever the + * path said. + * + * The end-to-end tests install from a deliberately hostile directory and run + * the generated commands with a stub `node` that only prints its argv, so the + * path is checked for round-trip fidelity without executing a real hook. The + * substitution unit tests then cover paths Node itself cannot host, notably + * backslashes (the ESM loader rejects any module specifier containing one). + */ +import { test, after } from "node:test"; +import assert from "node:assert/strict"; +import { spawnSync } from "node:child_process"; +import { + cpSync, + existsSync, + mkdirSync, + mkdtempSync, + readFileSync, + realpathSync, + rmSync, + writeFileSync, +} from "node:fs"; +import { tmpdir } from "node:os"; +import { dirname, join, resolve } from "node:path"; +import { fileURLToPath } from "node:url"; +import { + shellQuote, + substituteHookPlaceholder, +} from "../plugin/scripts/lib/hook-template.mjs"; + +const __dirname = dirname(fileURLToPath(import.meta.url)); +const PLUGIN_SOURCE = resolve(__dirname, "../plugin"); + +/** + * One path segment carrying every construct a shell acts on that a directory + * name may legally contain *and* Node can still load a module from: a command + * substitution, a backtick substitution, a command separator, a single quote, a + * double quote and spaces. Both substitutions create a canary file, so an + * escape leaves evidence even if its output goes nowhere. + * + * Backslashes are covered by the unit tests below instead: Node's ESM loader + * refuses every module specifier containing one, so no plugin can be installed + * from such a directory in the first place. + */ +const HOSTILE_SEGMENT = [ + "memwal", + "$(touch subst-canary; echo PATH_SUBSTITUTION_EXECUTED)", + "`touch backtick-canary; echo BACKTICK_EXECUTED`", + "it's", + '"quoted"', + "end", +].join(" "); + +const HOOK_SCRIPTS = ["on_session_start.mjs", "on_user_prompt.mjs", "on_post_tool.mjs"]; + +// The stub separates arguments with an ASCII record separator rather than a +// newline, so a value containing a newline is still read back exactly. +const RS = "\u001e"; + +const root = realpathSync( + mkdtempSync(join(process.env.TMPDIR || tmpdir(), "codex-shell-safety-")) +); +after(() => rmSync(root, { recursive: true, force: true })); + +const pluginRoot = join(root, HOSTILE_SEGMENT); +const home = join(root, "home"); +const canaryDir = join(root, "canaries"); +const fakeBin = join(root, "bin"); + +cpSync(PLUGIN_SOURCE, pluginRoot, { recursive: true }); +for (const dir of [home, canaryDir, fakeBin]) mkdirSync(dir, { recursive: true }); + +// A stub `node` that prints its arguments instead of running a hook. +writeFileSync( + join(fakeBin, "node"), + '#!/bin/sh\nfor arg in "$@"; do printf "%s\\036" "$arg"; done\n', + { mode: 0o755 } +); + +const install = spawnSync( + process.execPath, + [join(pluginRoot, "scripts", "install_codex_hooks.mjs")], + { + env: { ...process.env, HOME: home, USERPROFILE: home }, + cwd: canaryDir, + encoding: "utf8", + } +); + +/** Run a shell command with the stub node on PATH and read back its argv. */ +function argvFor(command) { + const result = spawnSync("/bin/sh", ["-c", command], { + env: { + ...process.env, + HOME: home, + USERPROFILE: home, + PATH: `${fakeBin}:${process.env.PATH}`, + }, + cwd: canaryDir, + encoding: "utf8", + }); + assert.equal(result.status, 0, `${command}\n${result.stderr}`); + return result.stdout.split(RS).slice(0, -1); +} + +/** Every `command` in a hooks file, in document order. */ +function hookCommands(file = join(home, ".codex", "hooks.json")) { + const config = JSON.parse(readFileSync(file, "utf8")); + const commands = []; + for (const entries of Object.values(config.hooks || {})) { + for (const entry of entries) { + for (const hook of entry.hooks || []) commands.push(hook.command); + } + } + return commands; +} + +test("installing from a hostile path succeeds and writes parseable JSON", () => { + assert.equal(install.status, 0, `${install.stdout}\n${install.stderr}`); + const commands = hookCommands(); + assert.equal(commands.length, HOOK_SCRIPTS.length); + for (const script of HOOK_SCRIPTS) { + assert.ok( + commands.some((command) => command.includes(script)), + `no hook command references ${script}: ${JSON.stringify(commands)}` + ); + } +}); + +test("generated hook commands do not execute anything the path spells out", () => { + for (const command of hookCommands()) { + const argv = argvFor(command); + assert.equal(argv.length, 1, `expected one argument, got ${JSON.stringify(argv)}`); + // The substitutions arrive as inert text, not as their output. + assert.ok(argv[0].includes("$(touch subst-canary;"), argv[0]); + assert.ok(argv[0].includes("`touch backtick-canary;"), argv[0]); + } + assert.ok(!existsSync(join(canaryDir, "subst-canary")), "command substitution ran"); + assert.ok(!existsSync(join(canaryDir, "backtick-canary")), "backtick substitution ran"); +}); + +test("the hostile path survives the round trip verbatim as a single argument", () => { + const seen = new Set(); + for (const command of hookCommands()) { + const argv = argvFor(command); + assert.equal(argv.length, 1, `expected one argument, got ${JSON.stringify(argv)}`); + assert.equal(dirname(dirname(argv[0])), pluginRoot); + seen.add(argv[0]); + } + assert.deepEqual( + [...seen].sort(), + HOOK_SCRIPTS.map((script) => join(pluginRoot, "scripts", script)).sort() + ); +}); + +test("shellQuote survives every shell metacharacter, backslashes included", () => { + const values = [ + "/plain/path", + "/with space/dir", + "/with/$(echo SUBST)", + "/with/`echo TICK`", + "/with/it's", + '/with/"double"', + "/with/back\\slash", + "/with/back\\\\slash", + "/with/$HOME and ${HOME}", + "/with/;rm -rf .;", + "/with/new\nline", + "/with/'''", + "/with/*?[a-z]", + 'C:\\Program Files\\mem"wal\\$(x)', + ]; + for (const value of values) { + assert.deepEqual(argvFor(`printf '%s\\036' ${shellQuote(value)}`), [value]); + } +}); + +test("substitution quotes the plugin root into commands and leaves JSON intact", () => { + const template = JSON.parse( + readFileSync(join(PLUGIN_SOURCE, "hooks", "codex-hooks.json"), "utf8") + ); + const hostileRoot = 'C:\\mem"wal\\$(touch pwned)\\`id`\\it\'s here'; + const substituted = substituteHookPlaceholder(template, "${PLUGIN_ROOT}", hostileRoot); + + // A path full of JSON escapes survives a write/read cycle unchanged. + const file = join(root, "substituted.json"); + writeFileSync(file, JSON.stringify(substituted, null, 2) + "\n"); + assert.deepEqual(JSON.parse(readFileSync(file, "utf8")), substituted); + + const commands = hookCommands(file); + assert.equal(commands.length, HOOK_SCRIPTS.length); + for (const command of commands) { + const argv = argvFor(command); + assert.equal(argv.length, 1, `expected one argument, got ${JSON.stringify(argv)}`); + assert.ok(argv[0].startsWith(`${hostileRoot}/scripts/`), argv[0]); + } + assert.ok(!existsSync(join(canaryDir, "pwned")), "command substitution ran"); +}); + +test("an argv array under `command` is substituted literally, not quoted", () => { + const substituted = substituteHookPlaceholder( + { hooks: { E: [{ hooks: [{ command: ["node", "${PLUGIN_ROOT}/x.mjs"] }] }] } }, + "${PLUGIN_ROOT}", + "/it's here" + ); + assert.deepEqual(substituted.hooks.E[0].hooks[0].command, ["node", "/it's here/x.mjs"]); +}); diff --git a/packages/mcp/test/codex-registration.test.mjs b/packages/mcp/test/codex-registration.test.mjs new file mode 100644 index 000000000..594962313 --- /dev/null +++ b/packages/mcp/test/codex-registration.test.mjs @@ -0,0 +1,194 @@ +/** + * The Codex fallback installer must MIGRATE an existing `[mcp_servers.memwal]` + * block, not skip it (WALM-640). + * + * The first cut of the fix only changed what a fresh install writes: it returned + * early on `content.includes("[mcp_servers.memwal]")`. Everyone who had already run + * the installer therefore kept `command = "npx"` / `args = ["-y", "…@"]` in + * ~/.codex/config.toml for ever, while re-running the installer printed "already + * present" — which reads as success. That is the exact population the ticket was + * filed for, so it is the case these tests cover: the block is rewritten, whatever + * else the user put in it survives, and the run says what it changed. + */ +import { test } from "node:test"; +import assert from "node:assert/strict"; +import { spawnSync } from "node:child_process"; +import { mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { dirname, join, resolve } from "node:path"; +import { fileURLToPath } from "node:url"; + +import { planMcpRegistration } from "../plugin/scripts/lib/codex-config.mjs"; + +const __dirname = dirname(fileURLToPath(import.meta.url)); +const PLUGIN_DIR = resolve(__dirname, "../plugin"); +const INSTALLER = join(PLUGIN_DIR, "scripts", "install_codex_hooks.mjs"); +const LAUNCHER = join(PLUGIN_DIR, "scripts", "launch_mcp.mjs"); +const PIN = JSON.parse(readFileSync(join(PLUGIN_DIR, "plugin.json"), "utf8")).version; + +/** Exactly what the pre-WALM-640 installer left on disk. */ +const LEGACY_BLOCK = [ + "[features]", + "codex_hooks = true", + "", + "# MemWal memory server", + "[mcp_servers.memwal]", + 'command = "npx"', + `args = ["-y", "@mysten-incubation/memwal-mcp@${PIN}"]`, + 'env = { MEMWAL_NAMESPACE = "work" }', + "startup_timeout_ms = 30000", + "", + "[mcp_servers.something_else]", + 'command = "other-server"', + "", +].join("\n"); + +function makeHome(t, configToml) { + const home = mkdtempSync(join(tmpdir(), "memwal-codex-home-")); + t.after(() => rmSync(home, { recursive: true, force: true })); + mkdirSync(join(home, ".codex"), { recursive: true }); + if (configToml !== undefined) writeFileSync(join(home, ".codex", "config.toml"), configToml); + return home; +} + +function runInstaller(home) { + const result = spawnSync(process.execPath, [INSTALLER], { + encoding: "utf8", + env: { ...process.env, HOME: home, USERPROFILE: home }, + }); + return { + ...result, + config: readFileSync(join(home, ".codex", "config.toml"), "utf8"), + }; +} + +test("a legacy npx registration is migrated when the installer is re-run", (t) => { + const home = makeHome(t, LEGACY_BLOCK); + const result = runInstaller(home); + + assert.equal(result.status, 0, result.stderr); + assert.doesNotMatch(result.config, /command\s*=\s*"npx"/, "npx must be gone from config.toml"); + assert.doesNotMatch(result.config, /memwal-mcp@/, "the npx spec must be gone too"); + assert.match(result.config, /command = "node"/); + assert.ok( + result.config.includes(JSON.stringify(LAUNCHER)), + `config.toml does not point at ${LAUNCHER}:\n${result.config}`, + ); + + // The user's own keys survive, inside and outside our block. + assert.match(result.config, /env = \{ MEMWAL_NAMESPACE = "work" \}/); + assert.match(result.config, /startup_timeout_ms = 30000/); + assert.match(result.config, /# MemWal memory server/); + assert.match(result.config, /\[mcp_servers\.something_else\]/); + assert.match(result.config, /command = "other-server"/); + assert.match(result.config, /codex_hooks = true/); + + // And it says so, rather than "already present". + assert.match(result.stdout, /Migrated \[mcp_servers\.memwal\]/); + assert.doesNotMatch(result.stdout, /already present/); + assert.match(result.stdout, /kept your other keys: env, startup_timeout_ms/); +}); + +test("a second run after the migration reports no change and rewrites nothing", (t) => { + const home = makeHome(t, LEGACY_BLOCK); + const first = runInstaller(home); + const second = runInstaller(home); + + assert.equal(second.status, 0, second.stderr); + assert.equal(second.config, first.config, "a settled config.toml must not keep churning"); + assert.match(second.stdout, /already runs the launcher/); +}); + +test("a fresh config.toml gets the launcher registration", (t) => { + const home = makeHome(t, ""); + const result = runInstaller(home); + + assert.equal(result.status, 0, result.stderr); + assert.match(result.stdout, /Registered \[mcp_servers\.memwal\]/); + assert.match(result.config, /\[mcp_servers\.memwal\]/); + assert.match(result.config, /command = "node"/); + assert.ok(result.config.includes(JSON.stringify(LAUNCHER))); +}); + +test("planMcpRegistration keeps every unrelated key and only rewrites command/args", () => { + const plan = planMcpRegistration(LEGACY_BLOCK, "/abs/launch_mcp.mjs"); + + assert.equal(plan.action, "migrated"); + assert.deepEqual(plan.preserved, ["env", "startup_timeout_ms"]); + assert.equal(plan.previous.command, '"npx"'); + assert.match(plan.previous.args, /memwal-mcp@/); + assert.equal( + plan.content, + [ + "[features]", + "codex_hooks = true", + "", + "# MemWal memory server", + "[mcp_servers.memwal]", + 'command = "node"', + 'args = ["/abs/launch_mcp.mjs"]', + 'env = { MEMWAL_NAMESPACE = "work" }', + "startup_timeout_ms = 30000", + "", + "[mcp_servers.something_else]", + 'command = "other-server"', + "", + ].join("\n"), + ); +}); + +test("a multi-line args array is replaced whole, not line by line", () => { + const content = [ + "[mcp_servers.memwal]", + 'command = "npx"', + "args = [", + ' "-y",', + ' "@mysten-incubation/memwal-mcp@0.0.14",', + "]", + 'env = { A = "1" }', + "", + "[other]", + "x = 1", + "", + ].join("\n"); + + const plan = planMcpRegistration(content, "/abs/launch_mcp.mjs"); + assert.equal(plan.action, "migrated"); + assert.doesNotMatch(plan.content, /-y/); + assert.doesNotMatch(plan.content, /memwal-mcp@0\.0\.14/); + assert.match(plan.content, /args = \["\/abs\/launch_mcp\.mjs"\]/); + assert.match(plan.content, /env = \{ A = "1" \}/); + assert.match(plan.content, /\[other\]\nx = 1/); +}); + +test("a block with the keys missing still ends up launching the launcher", () => { + const plan = planMcpRegistration( + ['[mcp_servers.memwal]', 'env = { A = "1" }', ""].join("\n"), + "/abs/launch_mcp.mjs", + ); + assert.equal(plan.action, "migrated"); + assert.match(plan.content, /command = "node"/); + assert.match(plan.content, /args = \["\/abs\/launch_mcp\.mjs"\]/); + assert.match(plan.content, /env = \{ A = "1" \}/); +}); + +test("a block pointing at a different launcher path is re-pointed", () => { + const stale = [ + "[mcp_servers.memwal]", + 'command = "node"', + 'args = ["/old/plugin/scripts/launch_mcp.mjs"]', + "", + ].join("\n"); + const plan = planMcpRegistration(stale, "/new/plugin/scripts/launch_mcp.mjs"); + assert.equal(plan.action, "migrated"); + assert.match(plan.content, /args = \["\/new\/plugin\/scripts\/launch_mcp\.mjs"\]/); + assert.doesNotMatch(plan.content, /\/old\/plugin/); +}); + +test("the section name is matched exactly, not by substring", () => { + const other = ['[mcp_servers.memwal_other]', 'command = "npx"', ""].join("\n"); + const plan = planMcpRegistration(other, "/abs/launch_mcp.mjs"); + assert.equal(plan.action, "added", "an unrelated server must not be rewritten"); + assert.match(plan.content, /\[mcp_servers\.memwal_other\]\ncommand = "npx"/); + assert.match(plan.content, /\[mcp_servers\.memwal\]\ncommand = "node"/); +}); diff --git a/packages/mcp/test/coldstart-init.test.mjs b/packages/mcp/test/coldstart-init.test.mjs index ef400c5bc..b7d426ce2 100644 --- a/packages/mcp/test/coldstart-init.test.mjs +++ b/packages/mcp/test/coldstart-init.test.mjs @@ -27,6 +27,8 @@ import { tmpdir } from "node:os"; import { join, dirname, resolve } from "node:path"; import { fileURLToPath } from "node:url"; +import { BASELINE_RELAYER_TOOLS } from "../dist/auth-required.js"; + const __dirname = dirname(fileURLToPath(import.meta.url)); const BIN = resolve(__dirname, "../dist/bin/memwal-mcp.js"); const EXPECTED_BEARER = "a".repeat(64); @@ -37,13 +39,16 @@ const EXPECTED_ACCOUNT_ID = "0x" + "3".repeat(64); * `initialize` would blow the assertion deadlines below. */ const SSE_DELAY_MS = 3_000; -/** The tools the real relayer sidecar registers - * (services/server/scripts/mcp/tools/index.ts). The cold-start static list must - * cover exactly these (plus the locally-served login/logout), so the - * static→refreshed transition doesn't change the tool set under the client. */ +/** The tools the CURRENT relayer sidecar registers + * (services/server/scripts/mcp/tools/index.ts) — i.e. a relayer as new as this + * bridge. The cold-start static list must be a SUBSET of this: it may lag the + * sidecar (a newer tool simply shows up on the post-connect re-list), but it + * must never advertise a name the relayer does not serve. A bridge is routinely + * newer than the relayer it dials, and over-advertising is GH #928. */ const UPSTREAM_TOOL_NAMES = [ "memwal_remember", "memwal_remember_bulk", + "memwal_remember_status", "memwal_recall", "memwal_analyze", "memwal_restore", @@ -314,6 +319,11 @@ test("initialize is answered locally during a slow relayer cold start; tools/cal ); assert.match(init.result.instructions, /memwal_recall/); assert.match(init.result.instructions, /memwal_remember/); + assert.doesNotMatch( + init.result.instructions, + /memwal_remember_status/, + "initialize instructions must not name a tool cold start does not advertise", + ); assert.notEqual( init.result.serverInfo.version, "0.0.1", @@ -324,13 +334,13 @@ test("initialize is answered locally during a slow relayer cold start; tools/cal `initialize took ${initElapsed}ms — expected it answered locally, well before the ${SSE_DELAY_MS}ms relayer connect`, ); - // 2) tools/list at cold start is served locally and instantly with the - // static list, which must be EXACTLY the upstream tool set plus the - // locally-served login/logout — each name once. + // 2) tools/list at cold start is served locally and instantly from the + // static list. It must fit INSIDE the upstream tool set (plus the + // locally-served login/logout) and cover the baseline — each name once. send({ jsonrpc: "2.0", id: 2, method: "tools/list", params: {} }); const list = await waitFor((m) => m.id === 2 && m.result, SSE_DELAY_MS); const coldNames = list.result.tools.map((t) => t.name); - const expectedNames = new Set([...UPSTREAM_TOOL_NAMES, "memwal_login", "memwal_logout"]); + const servableNames = new Set([...UPSTREAM_TOOL_NAMES, "memwal_login", "memwal_logout"]); // Unique names (TOOL_DEFINITIONS bundles its own memwal_login; a blind concat // with the local login/logout defs would list it twice). assert.equal( @@ -338,12 +348,21 @@ test("initialize is answered locally during a slow relayer cold start; tools/cal coldNames.length, `cold tools/list has duplicate tool names: ${coldNames}`, ); - // Exact set match — guards against the cold list drifting from the real - // upstream registration (e.g. missing memwal_remember_bulk / memwal_health). + // Over-advertising is the GH #928 failure: the agent is handed a name the + // relayer cannot answer, and the call waits out the orphan deadline. + const overAdvertised = coldNames.filter((n) => !servableNames.has(n)); + assert.deepEqual( + overAdvertised, + [], + `cold tools/list advertises tools no relayer serves: ${overAdvertised}`, + ); + // Under-advertising below the baseline is the opposite drift: a tool every + // supported relayer has, missing for the whole cold-start window. + const missingBaseline = [...BASELINE_RELAYER_TOOLS].filter((n) => !coldNames.includes(n)); assert.deepEqual( - new Set(coldNames), - expectedNames, - `cold tools/list set mismatch. got ${[...coldNames].sort()}, expected ${[...expectedNames].sort()}`, + missingBaseline, + [], + `cold tools/list omits baseline tools: ${missingBaseline}`, ); // 3) tools/call sent BEFORE the stream is up must be buffered and served @@ -374,9 +393,10 @@ test("initialize is answered locally during a slow relayer cold start; tools/cal assert.equal(initReplies[0].msg.result.serverInfo.name, "memwal"); // 6) After connect, a re-list is forwarded upstream and spliced with - // login/logout. That authoritative set must EQUAL the cold static set — - // the static→refreshed transition must not change the tool set (each - // name once, no dup even if upstream ever served login). + // login/logout. That authoritative set is what the client acts on, and + // every cold-start name must still be in it — the transition may ADD + // tools (a relayer newer than the baseline) but must never take one + // away under a client that already read the cold list. send({ jsonrpc: "2.0", id: 4, method: "tools/list", params: {} }); const relist = await waitFor((m) => m.id === 4 && m.result, 10_000); const splicedNames = relist.result.tools.map((t) => t.name); @@ -385,10 +405,16 @@ test("initialize is answered locally during a slow relayer cold start; tools/cal splicedNames.length, `post-connect tools/list has duplicate tool names: ${splicedNames}`, ); + const withdrawn = coldNames.filter((n) => !splicedNames.includes(n)); + assert.deepEqual( + withdrawn, + [], + `post-connect tools/list withdrew cold-start tools: ${withdrawn}. cold=${[...coldNames].sort()} spliced=${[...splicedNames].sort()}`, + ); assert.deepEqual( new Set(splicedNames), - new Set(coldNames), - `cold and post-connect tool sets differ. cold=${[...coldNames].sort()} spliced=${[...splicedNames].sort()}`, + servableNames, + `post-connect tools/list must mirror the relayer. got ${[...splicedNames].sort()}, expected ${[...servableNames].sort()}`, ); assert.ok(mock.getSseGetCount() >= 1, "expected at least one SSE handshake"); diff --git a/packages/mcp/test/credential-file-permissions.test.mjs b/packages/mcp/test/credential-file-permissions.test.mjs index f85f583e2..e0e0b55e6 100644 --- a/packages/mcp/test/credential-file-permissions.test.mjs +++ b/packages/mcp/test/credential-file-permissions.test.mjs @@ -172,6 +172,36 @@ test("the backup of a displaced account is written at 0600", async (t) => { assert.equal(JSON.parse(readFileSync(saved.backedUpTo, "utf8")).delegatePrivateKey, OLD_KEY); }); +// The login write-ahead record (WALM-332) holds the same plaintext key before +// the browser ever sees its public half, so it needs the same property. +test("savePendingLogin never writes the key through a pre-existing permissive file", POSIX_ONLY, async (t) => { + const { auth } = await sandbox(t); + const path = auth.pendingLoginPath(); + const makePending = (delegatePrivateKey) => ({ + delegatePrivateKey, + delegatePublicKeyHex: "d".repeat(64), + delegateAddress: "0x" + "e".repeat(64), + relayerUrl: "https://relayer.example", + label: "Test", + createdAt: new Date().toISOString(), + version: 1, + }); + mkdirSync(dirname(path), { recursive: true, mode: 0o700 }); + writeFileSync(path, JSON.stringify(makePending(OLD_KEY)), { mode: 0o644 }); + + const attackerFd = openSync(path, "r"); + t.after(() => closeSync(attackerFd)); + + auth.savePendingLogin(makePending(NEW_KEY)); + + assert.ok( + !readThroughOpenFd(attackerFd).includes(NEW_KEY), + "the pending private key must never be readable through the pre-existing 0644 inode", + ); + assert.equal(JSON.parse(readFileSync(path, "utf8")).delegatePrivateKey, NEW_KEY); + assert.equal(modeOf(path), 0o600); +}); + test("saveCreds leaves no temporary file behind", async (t) => { const { auth, home } = await sandbox(t, { existingFileMode: 0o644 }); diff --git a/packages/mcp/test/credential-resolution.test.mjs b/packages/mcp/test/credential-resolution.test.mjs index 7f8ca153b..80d210f58 100644 --- a/packages/mcp/test/credential-resolution.test.mjs +++ b/packages/mcp/test/credential-resolution.test.mjs @@ -11,6 +11,12 @@ * takes precedence over `~/.memwal/credentials.json`, npmrc/git-style. Purely * additive — a machine with no project-local file behaves exactly as before. * + * WALM-639 then made that precedence conditional: the project file only wins + * once the user has approved it on this machine, because the file itself lives + * inside the repository and could otherwise choose the destination on its own. + * The tests below that exercise precedence therefore approve first; the gate + * itself is covered in project-creds-approval.test.mjs. + * * `auth.js` resolves paths at call time, so each test sets HOME and cwd first * and then imports with a cache-busting query, the pattern used by * login-preflight.test.mjs. @@ -48,7 +54,7 @@ function writeCredsAt(root, accountId, label) { /** Fresh sandbox: a HOME and a working directory, with the module re-imported * so it observes them. Returns the module plus both roots. */ -async function sandbox(t, { global: globalAccount, project: projectAccount }) { +async function sandbox(t, { global: globalAccount, project: projectAccount, approve }) { // Canonicalise both roots: `process.cwd()` and `homedir()` report resolved // paths, so a raw mkdtemp path would not compare equal to what the module // computes. Needed on macOS (`/var` is a symlink to `/private/var`) and @@ -78,13 +84,20 @@ async function sandbox(t, { global: globalAccount, project: projectAccount }) { }); const auth = await import(`../dist/auth.js?walm361=${Date.now()}-${Math.random()}`); + // A project file is inert until approved (WALM-639). Tests about which file + // WINS approve it here so they keep testing precedence rather than the gate. + if (approve) { + const result = auth.approveProjectCreds(); + assert.equal(result.outcome, "approved", `approval failed: ${result.outcome}`); + } return { auth, home, cwd }; } -test("a project-local credentials file takes precedence over the global one", async (t) => { +test("an approved project-local credentials file takes precedence over the global one", async (t) => { const { auth, cwd } = await sandbox(t, { global: GLOBAL_ACCOUNT, project: PROJECT_ACCOUNT, + approve: true, }); assert.equal( @@ -110,6 +123,7 @@ test("saveCreds writes back to the project-local file when that is the one in us const { auth, home, cwd } = await sandbox(t, { global: GLOBAL_ACCOUNT, project: PROJECT_ACCOUNT, + approve: true, }); const updated = makeCreds(PROJECT_ACCOUNT, "Renamed"); @@ -249,6 +263,7 @@ test("a subdirectory of the project resolves to the project's credentials", asyn const { auth, cwd } = await sandbox(t, { global: GLOBAL_ACCOUNT, project: PROJECT_ACCOUNT, + approve: true, }); chdirBelow(cwd, "src", "nested"); @@ -295,6 +310,7 @@ test("removing a project file reports the global one that takes over", async (t) const { auth, home, cwd } = await sandbox(t, { global: GLOBAL_ACCOUNT, project: PROJECT_ACCOUNT, + approve: true, }); const result = auth.clearCreds(); diff --git a/packages/mcp/test/creds-dir-trust.test.mjs b/packages/mcp/test/creds-dir-trust.test.mjs new file mode 100644 index 000000000..81b2235c2 --- /dev/null +++ b/packages/mcp/test/creds-dir-trust.test.mjs @@ -0,0 +1,505 @@ +/** + * `MEMWAL_CREDS_DIR` cannot relocate the approval store into a repository, and + * an approved project never receives a plaintext delegate key it did not ask + * for (review of WALM-639). + * + * Two holes the gate itself left open: + * + * A. `trustedStateDir()` was `process.env.MEMWAL_CREDS_DIR ?? …`. `??` keeps + * an empty string, and nothing checked that the value was absolute — so + * `MEMWAL_CREDS_DIR=""` made `projectApprovalsPath()` the bare relative + * string "project-approvals.json", and `MEMWAL_CREDS_DIR=.memwal` made it + * `/.memwal/project-approvals.json`. Both resolve against the + * working directory, which is the repository. And because the override + * branch of `resolveCreds()` runs BEFORE any approval lookup, a committed + * `.cursor/mcp.json` carrying that one env var got the repo's credentials + * used with no approval at all — the exact silent redirect WALM-639 is + * about, through the escape hatch instead of around it. + * + * B. `pendingLoginPath()` was `join(dirname(credsPath()), …)`. Once a project + * was approved that is `/.memwal/`, so every sign-in wrote a 64-hex + * Ed25519 seed, in plaintext, into the working tree. + * + * Same sandbox pattern as project-creds-approval.test.mjs: HOME, cwd and the + * override are set first, then `auth.js` is imported with a cache-busting query. + */ +import { test } from "node:test"; +import assert from "node:assert/strict"; +import { + mkdtempSync, + mkdirSync, + writeFileSync, + readFileSync, + readdirSync, + rmSync, + existsSync, + realpathSync, +} from "node:fs"; +import { tmpdir } from "node:os"; +import { join, dirname, isAbsolute, relative } from "node:path"; + +const GLOBAL_ACCOUNT = "0x" + "a".repeat(64); +const PROJECT_ACCOUNT = "0x" + "b".repeat(64); +const CREDS_KEY = "c".repeat(64); +/** Distinct from CREDS_KEY — and from every other filler above, `packageId` + * included — so "this key never reaches the repo" is a claim about the pending + * record specifically and cannot be satisfied or broken by another field. */ +const PENDING_KEY = "4".repeat(64); +const GLOBAL_RELAYER = "https://relayer.example"; +const PROJECT_RELAYER = "https://project-relayer.example"; + +function makeCreds(overrides = {}) { + return { + delegatePrivateKey: CREDS_KEY, + delegatePublicKeyHex: "d".repeat(64), + delegateAddress: "0x" + "e".repeat(64), + walletAddress: "0x" + "f".repeat(64), + accountId: GLOBAL_ACCOUNT, + packageId: "0x" + "1".repeat(64), + relayerUrl: GLOBAL_RELAYER, + createdAt: new Date(0).toISOString(), + version: 1, + ...overrides, + }; +} + +function makePending(overrides = {}) { + return { + delegatePrivateKey: PENDING_KEY, + delegatePublicKeyHex: "2".repeat(64), + delegateAddress: "0x" + "3".repeat(64), + relayerUrl: PROJECT_RELAYER, + createdAt: new Date().toISOString(), + version: 1, + ...overrides, + }; +} + +function writeCredsAt(root, creds) { + const path = join(root, ".memwal", "credentials.json"); + mkdirSync(dirname(path), { recursive: true }); + writeFileSync(path, JSON.stringify(creds), { mode: 0o600 }); + return path; +} + +/** Same containment question the module asks: `relative` rather than a string + * prefix, so `/repo` and `/repo-2` do not look like the same directory. */ +function isInside(root, path) { + const rel = relative(root, path); + return rel === "" || (!rel.startsWith("..") && !isAbsolute(rel)); +} + +/** Every file at or under `root`, so a secret can be searched for across a + * whole working tree rather than at the one path a test remembered to check. */ +function filesUnder(root) { + const out = []; + for (const entry of readdirSync(root, { withFileTypes: true })) { + const path = join(root, entry.name); + if (entry.isDirectory()) out.push(...filesUnder(path)); + else if (entry.isFile()) out.push(path); + } + return out; +} + +/** Paths under `root` whose bytes contain `needle`. */ +function filesContaining(root, needle) { + return filesUnder(root).filter((path) => { + try { + return readFileSync(path, "utf8").includes(needle); + } catch { + return false; + } + }); +} + +/** + * A HOME, a working directory, a fresh module, and MEMWAL_CREDS_DIR cleared. + * + * Canonicalised because `process.cwd()` and `homedir()` report resolved paths + * and macOS routes /tmp through /private/tmp — the containment check under test + * has to hold for the same directory spelled either way. + */ +async function sandbox(t, { global: globalCreds, project: projectCreds, git = false } = {}) { + const home = realpathSync(mkdtempSync(join(tmpdir(), "memwal-trust-home-"))); + const cwd = realpathSync(mkdtempSync(join(tmpdir(), "memwal-trust-cwd-"))); + const previous = { + home: process.env.HOME, + profile: process.env.USERPROFILE, + credsDir: process.env.MEMWAL_CREDS_DIR, + cwd: process.cwd(), + }; + + process.env.HOME = home; + process.env.USERPROFILE = home; + delete process.env.MEMWAL_CREDS_DIR; + process.chdir(cwd); + + if (git) mkdirSync(join(cwd, ".git"), { recursive: true }); + if (globalCreds) writeCredsAt(home, globalCreds); + if (projectCreds) writeCredsAt(cwd, projectCreds); + + t.after(() => { + process.chdir(previous.cwd); + process.env.HOME = previous.home; + process.env.USERPROFILE = previous.profile; + if (previous.credsDir === undefined) delete process.env.MEMWAL_CREDS_DIR; + else process.env.MEMWAL_CREDS_DIR = previous.credsDir; + rmSync(home, { recursive: true, force: true }); + rmSync(cwd, { recursive: true, force: true }); + }); + + const bust = `${Date.now()}-${Math.random()}`; + const auth = await import(`../dist/auth.js?walm639trust=${bust}`); + return { auth, home, cwd, bust }; +} + +const globalFile = (home) => join(home, ".memwal", "credentials.json"); +const projectFile = (cwd) => join(cwd, ".memwal", "credentials.json"); + +/* --------------------------------------------------------------------- * + * A. An empty value means unset — never "the working directory". + * --------------------------------------------------------------------- */ + +test('MEMWAL_CREDS_DIR="" does not relocate the approval store into the repo', async (t) => { + const { auth, home, cwd } = await sandbox(t, { + global: makeCreds(), + project: makeCreds({ accountId: PROJECT_ACCOUNT, relayerUrl: PROJECT_RELAYER }), + git: true, + }); + process.env.MEMWAL_CREDS_DIR = ""; + + const approvals = auth.projectApprovalsPath(); + + // The reported repro: `??` kept the empty string, so `join("", FILE)` came + // back as the bare relative name and resolved against the repo. + assert.equal(isAbsolute(approvals), true, `approvals path is relative: ${approvals}`); + assert.notEqual(approvals, "project-approvals.json"); + assert.equal(approvals, join(home, ".memwal", "project-approvals.json")); + assert.equal(isInside(cwd, approvals), false, "the store must stay out of the repository"); +}); + +test('MEMWAL_CREDS_DIR="" does not silently approve the repo credentials file', async (t) => { + // The override branch returns before any approval lookup, so an empty value + // that counted as "set" adopted the repo's account with nothing asked. + const { auth, home } = await sandbox(t, { + global: makeCreds(), + project: makeCreds({ accountId: PROJECT_ACCOUNT, relayerUrl: PROJECT_RELAYER }), + git: true, + }); + process.env.MEMWAL_CREDS_DIR = ""; + + assert.equal(auth.resolveCreds().source, "global"); + assert.equal(auth.credsPath(), globalFile(home)); + assert.equal(auth.loadCreds()?.accountId, GLOBAL_ACCOUNT); + assert.equal(auth.resolveCreds().project?.decision, "unapproved"); + assert.ok(auth.formatProjectCredsNotice(), "the ignored repo file is still reported"); +}); + +test('MEMWAL_CREDS_DIR="" lets no repository approve its own credentials', async (t) => { + // The sharp end of the empty-string bug. `join("", APPROVALS_FILE)` is the + // bare name "project-approvals.json", which resolves against the working + // directory — so a repo that commits that one file at its root IS the + // approval store, and approves the credentials it also ships. Everything + // the gate does is decided by a file the attacker wrote. + const { auth, home, cwd } = await sandbox(t, { + global: makeCreds(), + project: makeCreds({ accountId: PROJECT_ACCOUNT, relayerUrl: PROJECT_RELAYER }), + git: true, + }); + const delegateAddress = "0x" + "e".repeat(64); + writeFileSync( + join(cwd, "project-approvals.json"), + JSON.stringify({ + version: 1, + approvals: [ + { + path: projectFile(cwd), + fingerprint: auth.credentialsFingerprint({ + accountId: PROJECT_ACCOUNT, + delegateAddress, + relayerUrl: PROJECT_RELAYER, + }), + accountId: PROJECT_ACCOUNT, + delegateAddress, + relayerUrl: PROJECT_RELAYER, + approvedAt: new Date().toISOString(), + }, + ], + }), + ); + process.env.MEMWAL_CREDS_DIR = ""; + + assert.equal(auth.credsPath(), globalFile(home), "a repo approved itself"); + assert.equal(auth.loadCreds()?.accountId, GLOBAL_ACCOUNT); + assert.equal(auth.resolveCreds().project?.decision, "unapproved"); +}); + +/* --------------------------------------------------------------------- * + * A. A relative value is refused, loudly. + * --------------------------------------------------------------------- */ + +for (const value of [".memwal", "creds", "./.memwal", "../elsewhere", ".claude/state"]) { + test(`a relative MEMWAL_CREDS_DIR (${value}) is refused, not followed`, async (t) => { + const { auth, cwd } = await sandbox(t, { + global: makeCreds(), + project: makeCreds({ accountId: PROJECT_ACCOUNT, relayerUrl: PROJECT_RELAYER }), + git: true, + }); + process.env.MEMWAL_CREDS_DIR = value; + + const refuses = /MEMWAL_CREDS_DIR is a relative path/; + assert.throws(() => auth.projectApprovalsPath(), refuses); + assert.throws(() => auth.credsPath(), refuses); + assert.throws(() => auth.resolveCreds(), refuses); + assert.throws(() => auth.loadCreds(), refuses); + // Refusing rather than falling back: a silent fallback would hide a + // misconfigured — or planted — client config. + assert.throws(() => auth.approveProjectCreds(), refuses); + assert.throws(() => auth.saveCreds(makeCreds()), refuses); + + // And nothing landed in the working tree on the way to refusing. + assert.deepEqual( + filesUnder(cwd).filter((p) => p.includes("project-approvals")), + [], + "an approval record was written inside the repository", + ); + }); +} + +test("a relative MEMWAL_CREDS_DIR leaves the repo file unapproved once cleared", async (t) => { + // The failure mode worth pinning: refusing must not be a disguised approval + // that shows up the moment the variable goes away. + const { auth, home, cwd } = await sandbox(t, { + global: makeCreds(), + project: makeCreds({ accountId: PROJECT_ACCOUNT, relayerUrl: PROJECT_RELAYER }), + git: true, + }); + process.env.MEMWAL_CREDS_DIR = ".memwal"; + assert.throws(() => auth.approveProjectCreds()); + + delete process.env.MEMWAL_CREDS_DIR; + + assert.equal(auth.credsPath(), globalFile(home)); + assert.equal(auth.resolveCreds().project?.decision, "unapproved"); + assert.equal(existsSync(join(cwd, ".memwal", "project-approvals.json")), false); +}); + +test("the refusal is loud at the CLI, not swallowed into a fallback", async (t) => { + const { cwd, bust } = await sandbox(t, { + global: makeCreds(), + project: makeCreds({ accountId: PROJECT_ACCOUNT }), + git: true, + }); + const { main } = await import(`../dist/index.js?walm639trust=${bust}`); + + const realTty = process.stdin.isTTY; + t.after(() => { + process.stdin.isTTY = realTty; + }); + process.stdin.isTTY = true; + process.env.MEMWAL_CREDS_DIR = join(cwd, ".memwal"); + + // `bin/memwal-mcp.ts` turns this into `[memwal-mcp] fatal: …` and exit 1, + // which is the behaviour the launcher already has for a relative runtime + // directory. + await assert.rejects(main(["approve-project"]), /MEMWAL_CREDS_DIR/); +}); + +/* --------------------------------------------------------------------- * + * A. An absolute value inside the project is refused the same way. + * --------------------------------------------------------------------- */ + +test("an absolute MEMWAL_CREDS_DIR inside the project is refused", async (t) => { + // `${workspaceFolder}` is expanded by editors inside the very config files + // an attacker can commit, so "absolute" is not evidence a human typed it. + const { auth, cwd } = await sandbox(t, { + global: makeCreds(), + project: makeCreds({ accountId: PROJECT_ACCOUNT }), + git: true, + }); + process.env.MEMWAL_CREDS_DIR = join(cwd, ".memwal"); + + assert.throws(() => auth.resolveCreds(), /points inside the current project/); + assert.throws(() => auth.projectApprovalsPath(), /points inside the current project/); +}); + +test("the refusal is against the project root, not just the working directory", async (t) => { + // Running from `src/nested` must not launder an override that points at the + // repository root — the creds walk climbs, so this check has to as well. + const { auth, cwd } = await sandbox(t, { + global: makeCreds(), + project: makeCreds({ accountId: PROJECT_ACCOUNT }), + git: true, + }); + const deep = join(cwd, "src", "nested"); + mkdirSync(deep, { recursive: true }); + process.chdir(deep); + process.env.MEMWAL_CREDS_DIR = join(cwd, ".memwal"); + + assert.throws(() => auth.resolveCreds(), /points inside the current project/); +}); + +/* --------------------------------------------------------------------- * + * A. The legitimate use keeps working, byte for byte. + * --------------------------------------------------------------------- */ + +test("an absolute MEMWAL_CREDS_DIR outside the project still decides everything", async (t) => { + const { auth, cwd } = await sandbox(t, { + global: makeCreds(), + project: makeCreds({ accountId: PROJECT_ACCOUNT, relayerUrl: PROJECT_RELAYER }), + git: true, + }); + const override = realpathSync(mkdtempSync(join(tmpdir(), "memwal-trust-override-"))); + t.after(() => rmSync(override, { recursive: true, force: true })); + process.env.MEMWAL_CREDS_DIR = override; + + assert.equal(auth.credsPath(), join(override, "credentials.json")); + assert.equal(auth.resolveCreds().source, "override"); + assert.equal(auth.projectApprovalsPath(), join(override, "project-approvals.json")); + assert.equal(auth.pendingLoginPath(), join(override, "login-pending.json")); + assert.equal(auth.formatProjectCredsNotice(), null, "an override has nothing to warn about"); + assert.equal(auth.approveProjectCreds().outcome, "overridden"); + + // And it is genuinely usable, not merely accepted. + auth.saveCreds(makeCreds({ label: "sandboxed" })); + assert.equal(auth.loadCreds()?.label, "sandboxed"); + assert.equal( + JSON.parse(readFileSync(projectFile(cwd), "utf8")).label, + undefined, + "the repo file must not have been written", + ); +}); + +/* --------------------------------------------------------------------- * + * B. An approved project never receives a plaintext delegate key. + * --------------------------------------------------------------------- */ + +test("the pending-login record for an approved project is not written into it", async (t) => { + const { auth, home, cwd } = await sandbox(t, { + global: makeCreds(), + project: makeCreds({ accountId: PROJECT_ACCOUNT, relayerUrl: PROJECT_RELAYER }), + git: true, + }); + auth.approveProjectCreds(); + assert.equal(auth.credsPath(), projectFile(cwd), "precondition: the project file is in use"); + + auth.savePendingLogin(makePending()); + + const path = auth.pendingLoginPath(); + assert.equal(isInside(cwd, path), false, `the write-ahead record landed in the repo: ${path}`); + assert.equal(isInside(join(home, ".memwal"), path), true, "it belongs in the trusted dir"); + assert.equal(existsSync(join(cwd, ".memwal", "login-pending.json")), false); + assert.deepEqual( + filesContaining(cwd, PENDING_KEY), + [], + "a plaintext delegate seed was written somewhere inside the repository", + ); + // Still a working write-ahead record, which is the whole point of it. + assert.equal(auth.loadPendingLogin()?.delegatePrivateKey, PENDING_KEY); + assert.equal(auth.reusablePendingLogin(PROJECT_RELAYER)?.delegatePrivateKey, PENDING_KEY); + auth.clearPendingLogin(); + assert.equal(auth.loadPendingLogin(), null); +}); + +test("one project's pending record is not another project's", async (t) => { + // Keying it by project is why it could live in the repo at all; moving it + // out must not turn it into one shared record. + const { auth, cwd } = await sandbox(t, { + global: makeCreds(), + project: makeCreds({ accountId: PROJECT_ACCOUNT, relayerUrl: PROJECT_RELAYER }), + git: true, + }); + auth.approveProjectCreds(); + auth.savePendingLogin(makePending()); + const first = auth.pendingLoginPath(); + + const other = realpathSync(mkdtempSync(join(tmpdir(), "memwal-trust-other-"))); + t.after(() => rmSync(other, { recursive: true, force: true })); + mkdirSync(join(other, ".git"), { recursive: true }); + writeCredsAt(other, makeCreds({ accountId: PROJECT_ACCOUNT, relayerUrl: PROJECT_RELAYER })); + process.chdir(other); + auth.approveProjectCreds(); + + assert.notEqual(auth.pendingLoginPath(), first, "two projects share one record"); + assert.equal(auth.loadPendingLogin(), null, "a sign-in leaked across projects"); +}); + +test("the global pending-login path is exactly where it always was", async (t) => { + const { auth, home } = await sandbox(t, { global: makeCreds() }); + + assert.equal(auth.pendingLoginPath(), join(home, ".memwal", "login-pending.json")); + auth.savePendingLogin(makePending()); + assert.equal(existsSync(join(home, ".memwal", "login-pending.json")), true); +}); + +/* --------------------------------------------------------------------- * + * B. And the user is told, rather than left to find the key in a diff. + * --------------------------------------------------------------------- */ + +test("approving says a private key will be written inside the repository", async (t) => { + const { auth, cwd, bust } = await sandbox(t, { + global: makeCreds(), + project: makeCreds({ accountId: PROJECT_ACCOUNT, relayerUrl: PROJECT_RELAYER }), + git: true, + }); + const { main } = await import(`../dist/index.js?walm639trust=${bust}`); + + const realTty = process.stdin.isTTY; + const realWrite = process.stderr.write.bind(process.stderr); + let output = ""; + t.after(() => { + process.stdin.isTTY = realTty; + process.stderr.write = realWrite; + }); + process.stdin.isTTY = true; + process.stderr.write = (chunk) => { + output += chunk; + return true; + }; + + await main(["approve-project"]); + process.stderr.write = realWrite; + + assert.match(output, /Approved /, "precondition: it approved"); + assert.match(output, /PRIVATE KEY/, `approval said nothing about the key:\n${output}`); + assert.match(output, /\.gitignore/, "no suggestion for keeping it out of the repo"); + assert.ok(output.includes(join(cwd, ".memwal")), "must name the directory in the repo"); +}); + +test("the warning is repeated at the last moment before a key is written", async (t) => { + // Approval may have happened months ago, or on someone else's shift. The + // sign-in warning is the last point the user can back out for free. + const { auth, cwd } = await sandbox(t, { + global: makeCreds(), + project: makeCreds({ accountId: PROJECT_ACCOUNT, relayerUrl: PROJECT_RELAYER }), + git: true, + }); + auth.approveProjectCreds(); + + const warning = auth.formatPendingSignInWarning(); + + assert.ok(warning.includes(projectFile(cwd)), "must name the file being replaced"); + assert.match(warning, /PRIVATE KEY/); + assert.match(warning, /\.gitignore/); +}); + +test("a global sign-in is not nagged about a repository it is not in", async (t) => { + const { auth, home } = await sandbox(t, { global: makeCreds() }); + + const warning = auth.formatPendingSignInWarning(); + + assert.ok(warning.includes(globalFile(home))); + assert.doesNotMatch(warning, /gitignore/, "nothing repo-shaped to say about the global file"); +}); + +test("nothing the storage warning prints is key material", async (t) => { + const { auth, cwd } = await sandbox(t, { + global: makeCreds(), + project: makeCreds({ accountId: PROJECT_ACCOUNT }), + }); + + const warning = auth.formatProjectCredsStorageWarning(projectFile(cwd)); + + assert.ok(!warning.includes(CREDS_KEY)); + assert.ok(!warning.includes(PENDING_KEY)); +}); diff --git a/packages/mcp/test/health-probe.test.mjs b/packages/mcp/test/health-probe.test.mjs new file mode 100644 index 000000000..e6ab805d6 --- /dev/null +++ b/packages/mcp/test/health-probe.test.mjs @@ -0,0 +1,140 @@ +/** + * `probeRelayerHealth` / `describeHealthProbe` (WALM-396). + * + * When a sent call's reply never arrives, the bridge asks the relayer's + * `/health` before answering, so the message can say whether the relayer is + * down, unhealthy, unreachable from this machine, or up with this one call + * stuck — the difference between "retry", "wait", and "fix your config". + */ +import { test } from "node:test"; +import assert from "node:assert/strict"; +import http from "node:http"; +import { + describeHealthProbe, + probeRelayerHealth, + resolveHealthProbeMs, +} from "../dist/health-probe.js"; + +async function fakeRelayer(respond) { + const server = http.createServer((req, res) => { + if (req.url === "/health") respond(res); + else res.writeHead(404).end(); + }); + await new Promise((r) => server.listen(0, "127.0.0.1", r)); + const { port } = server.address(); + return { + url: `http://127.0.0.1:${port}`, + close: () => + new Promise((r) => { + server.closeAllConnections(); + server.close(() => r()); + }), + }; +} + +async function closedPortUrl() { + const relayer = await fakeRelayer(() => {}); + await relayer.close(); + return relayer.url; +} + +const healthy = (res) => + res + .writeHead(200, { "content-type": "application/json" }) + .end(JSON.stringify({ status: "ok", version: "1.4.2" })); + +test("a healthy relayer is reported with its version", async (t) => { + const relayer = await fakeRelayer(healthy); + t.after(relayer.close); + const probe = await probeRelayerHealth(`${relayer.url}/`, 2000); + assert.equal(probe.kind, "ok"); + assert.equal(probe.version, "1.4.2"); + const { health, reachable } = describeHealthProbe(probe, relayer.url); + assert.equal(reachable, true); + assert.match(health, /^ok \(\d+ms, v1\.4\.2\)$/); +}); + +test("an unhealthy relayer is reported by status", async (t) => { + const relayer = await fakeRelayer((res) => res.writeHead(503).end()); + t.after(relayer.close); + const probe = await probeRelayerHealth(relayer.url, 2000); + assert.equal(probe.kind, "http"); + assert.equal(probe.status, 503); + const { health, verdict, reachable } = describeHealthProbe(probe, relayer.url); + assert.equal(reachable, false); + assert.match(health, /HTTP 503/); + assert.match(verdict, /not healthy/); +}); + +test("a relayer that never answers is a timeout, reported on time", async (t) => { + const relayer = await fakeRelayer(() => {}); + t.after(relayer.close); + const started = Date.now(); + const probe = await probeRelayerHealth(relayer.url, 200); + assert.equal(probe.kind, "timeout"); + assert.ok(Date.now() - started < 2000, "the probe must not outlive its budget"); + assert.match(describeHealthProbe(probe, relayer.url).verdict, /down|overloaded|not reachable/); +}); + +test("a refused connection is unreachable, and says the relayer may be down", async () => { + const url = await closedPortUrl(); + const probe = await probeRelayerHealth(url, 2000); + assert.equal(probe.kind, "unreachable"); + assert.equal(probe.code, "ECONNREFUSED"); + const { health, verdict } = describeHealthProbe(probe, url); + assert.match(health, /unreachable \(ECONNREFUSED\)/); + assert.match(verdict, /down or not reachable/); + assert.ok(verdict.includes(url), "the verdict names the address it tried"); +}); + +test("an unresolvable host points at the configured relayer URL", () => { + // A typo'd --relayer is the one cause the user can fix on their side. + const { verdict } = describeHealthProbe( + { kind: "unreachable", ms: 5, code: "ENOTFOUND" }, + "https://relayer.memroy.walrus.xyz", + ); + assert.match(verdict, /MEMWAL_SERVER_URL/); + assert.match(verdict, /--relayer/); +}); + +test("the probe budget has a default and an override", () => { + const saved = process.env.MEMWAL_MCP_HEALTH_PROBE_MS; + try { + delete process.env.MEMWAL_MCP_HEALTH_PROBE_MS; + assert.equal(resolveHealthProbeMs(), 3000); + process.env.MEMWAL_MCP_HEALTH_PROBE_MS = "750"; + assert.equal(resolveHealthProbeMs(), 750); + process.env.MEMWAL_MCP_HEALTH_PROBE_MS = "nonsense"; + assert.equal(resolveHealthProbeMs(), 3000); + // `AbortSignal.timeout` throws on both of these; the resolver must + // never hand it one. + process.env.MEMWAL_MCP_HEALTH_PROBE_MS = "2500.5"; + assert.equal(resolveHealthProbeMs(), 2500); + process.env.MEMWAL_MCP_HEALTH_PROBE_MS = String(2 ** 40); + assert.equal(resolveHealthProbeMs(), 60_000); + } finally { + if (saved === undefined) delete process.env.MEMWAL_MCP_HEALTH_PROBE_MS; + else process.env.MEMWAL_MCP_HEALTH_PROBE_MS = saved; + } +}); + +test("a probe given a budget AbortSignal cannot take still resolves", async (t) => { + // The sweeper answers the call from this promise. A rejection here would + // crash the bridge, or leave the call unanswered for good. + const relayer = await fakeRelayer(healthy); + t.after(relayer.close); + const probe = await probeRelayerHealth(relayer.url, 2500.5); + assert.ok(["ok", "unreachable"].includes(probe.kind), `got ${probe.kind}`); +}); + +test("paused writes show in the health line, since /health still answers 200", async (t) => { + const relayer = await fakeRelayer((res) => + res + .writeHead(200, { "content-type": "application/json" }) + .end(JSON.stringify({ status: "ok", version: "1.4.2", write_ready: false })), + ); + t.after(relayer.close); + const probe = await probeRelayerHealth(relayer.url, 2000); + assert.equal(probe.writesUnavailable, true); + assert.match(describeHealthProbe(probe, relayer.url).health, /writes unavailable/); +}); diff --git a/packages/mcp/test/health-relayer-annotation.test.mjs b/packages/mcp/test/health-relayer-annotation.test.mjs index 9d8e7fc81..e1e2b32b6 100644 --- a/packages/mcp/test/health-relayer-annotation.test.mjs +++ b/packages/mcp/test/health-relayer-annotation.test.mjs @@ -34,6 +34,36 @@ test("replaces the relayer the reply already carried rather than adding a second assert.ok(text.includes("write_ready=true")); }); +// A project-local credentials file can point this process at a different +// account than the one the user signed in with (WALM-639), so the account is +// part of the destination `memwal_health` reports, not just the relayer. + +const ACCOUNT = "0x" + "b".repeat(64); + +test("names the account this session signs for", () => { + const result = healthResult("Walrus Memory is reachable. status=ok"); + annotateHealthResult(result, DEV, ACCOUNT); + assert.ok(result.content[0].text.includes(`account=${ACCOUNT}`)); + assert.ok(result.content[0].text.includes(`relayer=${DEV}`)); +}); + +test("replaces an account the reply already carried rather than adding a second", () => { + const result = healthResult("status=ok account=0xstale write_ready=true"); + annotateHealthResult(result, DEV, ACCOUNT); + const text = result.content[0].text; + assert.equal(text.match(/account=/g).length, 1, `two account fields:\n${text}`); + assert.ok(!text.includes("0xstale")); + assert.ok(text.includes("write_ready=true"), "the field after it must survive"); +}); + +test("says nothing about an account it does not know", () => { + // Signed out, or a cold start before credentials are adopted. An empty + // `account=` would read as an account rather than as an absence. + const result = healthResult("status=ok"); + annotateHealthResult(result, DEV); + assert.ok(!result.content[0].text.includes("account=")); +}); + test("leaves a failed health call alone", () => { // Naming a relayer beside an error reads as though that relayer answered. const result = { ...healthResult("relayer unreachable"), isError: true }; diff --git a/packages/mcp/test/hook-creds-dir-trust.test.mjs b/packages/mcp/test/hook-creds-dir-trust.test.mjs new file mode 100644 index 000000000..108227b55 --- /dev/null +++ b/packages/mcp/test/hook-creds-dir-trust.test.mjs @@ -0,0 +1,77 @@ +import { test } from "node:test"; +import assert from "node:assert/strict"; +import { spawnSync } from "node:child_process"; +import { mkdtempSync, mkdirSync, writeFileSync, rmSync, realpathSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { dirname, join, resolve } from "node:path"; +import { fileURLToPath } from "node:url"; + +const __dirname = dirname(fileURLToPath(import.meta.url)); +const SCRIPTS = resolve(__dirname, "../plugin/scripts"); + +const SAVE_RUBRIC = /call memwal_remember \(or memwal_remember_bulk for several\)/; + +test("a repo-relative MEMWAL_CREDS_DIR cannot publish the hook's consent answer", async (t) => { + const home = realpathSync(mkdtempSync(join(tmpdir(), "memwal-hookdir-home-"))); + const repo = realpathSync(mkdtempSync(join(tmpdir(), "memwal-hookdir-repo-"))); + t.after(() => { + rmSync(home, { recursive: true, force: true }); + rmSync(repo, { recursive: true, force: true }); + }); + mkdirSync(join(repo, ".git"), { recursive: true }); + mkdirSync(join(repo, ".memwal"), { recursive: true }); + writeFileSync( + join(repo, ".memwal", "auto-save-state.json"), + JSON.stringify({ + version: 1, + enabled: true, + state: "on", + source: "settings", + pendingConsent: false, + updatedAt: new Date().toISOString(), + }), + ); + + const result = spawnSync(process.execPath, [join(SCRIPTS, "on_user_prompt.mjs")], { + cwd: repo, + input: JSON.stringify({ + session_id: `walm642-credsdir-${Date.now()}-${Math.random().toString(16).slice(2)}`, + prompt: "I prefer pnpm", + }), + encoding: "utf8", + env: { + ...process.env, + HOME: home, + USERPROFILE: home, + MEMWAL_CREDS_DIR: ".memwal", + MEMWAL_AUTO_SAVE: "", + }, + }); + assert.equal(result.status, 0, `on_user_prompt.mjs exited ${result.status}: ${result.stderr}`); + const injected = result.stdout.trim() + ? (JSON.parse(result.stdout).hookSpecificOutput?.additionalContext ?? "") + : ""; + assert.doesNotMatch( + injected, + SAVE_RUBRIC, + "a repo-relative MEMWAL_CREDS_DIR turned automatic memory on", + ); + + const previous = { home: process.env.HOME, profile: process.env.USERPROFILE }; + const previousCredsDir = process.env.MEMWAL_CREDS_DIR; + const previousCwd = process.cwd(); + process.env.HOME = home; + process.env.USERPROFILE = home; + process.env.MEMWAL_CREDS_DIR = ".memwal"; + process.chdir(repo); + t.after(() => { + process.chdir(previousCwd); + process.env.HOME = previous.home; + process.env.USERPROFILE = previous.profile; + if (previousCredsDir === undefined) delete process.env.MEMWAL_CREDS_DIR; + else process.env.MEMWAL_CREDS_DIR = previousCredsDir; + }); + const hook = await import(`../plugin/scripts/lib/auto-save.mjs?walm642creds=${Date.now()}`); + assert.equal(hook.hookStatePath(), join(home, ".memwal", "auto-save-state.json")); + assert.equal(hook.isAutoSaveEnabled(), false); +}); diff --git a/packages/mcp/test/hook-state-symlink.test.mjs b/packages/mcp/test/hook-state-symlink.test.mjs new file mode 100644 index 000000000..123d715ad --- /dev/null +++ b/packages/mcp/test/hook-state-symlink.test.mjs @@ -0,0 +1,171 @@ +/** + * Hook state must not be redirectable through a symlink (WALM-644). + * + * The old `existsSync(marker) ? skip : writeFileSync(marker, "1")` pair let a + * dangling symlink planted at a marker path pass the existence check and then + * absorb the write, creating a file outside the state directory. Markers are + * now created exclusively (O_CREAT|O_EXCL, plus O_NOFOLLOW where it exists), + * so an occupied path is refused rather than followed — and the directory + * itself is private and verified before use. + */ +import { test } from "node:test"; +import assert from "node:assert/strict"; +import { spawnSync } from "node:child_process"; +import { + mkdtempSync, + mkdirSync, + rmSync, + symlinkSync, + lstatSync, + existsSync, + readdirSync, + readFileSync, +} from "node:fs"; +import { tmpdir } from "node:os"; +import { dirname, join, resolve } from "node:path"; +import { fileURLToPath } from "node:url"; + +const __dirname = dirname(fileURLToPath(import.meta.url)); +const HOOK = resolve(__dirname, "../plugin/scripts/on_user_prompt.mjs"); +const HOOK_IO = "../plugin/scripts/lib/hook-io.mjs"; + +// Windows needs a privilege or developer mode to create symlinks at all, and +// has no O_NOFOLLOW; the attack this guards is a POSIX temp-dir one. +const skipOnWindows = + process.platform === "win32" ? "symlink creation needs privileges on Windows" : false; + +/** A throwaway TMPDIR plus an "elsewhere" directory a symlink could point into. */ +function sandbox(t) { + const root = mkdtempSync(join(tmpdir(), "memwal-hookstate-")); + const temp = join(root, "tmp"); + const elsewhere = join(root, "elsewhere"); + mkdirSync(temp); + mkdirSync(elsewhere); + const previous = process.env.TMPDIR; + process.env.TMPDIR = temp; + t.after(() => { + if (previous === undefined) delete process.env.TMPDIR; + else process.env.TMPDIR = previous; + rmSync(root, { recursive: true, force: true }); + }); + return { + root, + temp, + elsewhere, + stateDir: join(temp, "memwal-hooks"), + target: join(elsewhere, "pwned"), + }; +} + +/** Every file under `dir`, relative to it — used to prove nothing escaped. */ +function treeOf(dir) { + return readdirSync(dir, { recursive: true, withFileTypes: true }) + .filter((entry) => !entry.isDirectory()) + .map((entry) => entry.name) + .sort(); +} + +test("a dangling symlink at a marker path cannot create the target", { skip: skipOnWindows }, async (t) => { + const box = sandbox(t); + const { firstTime, stateDir } = await import(HOOK_IO); + + assert.equal(stateDir(), box.stateDir); + const marker = join(box.stateDir, "rubric_walm644"); + symlinkSync(box.target, marker); + assert.equal(existsSync(box.target), false, "precondition: the target is dangling"); + + // The reproduction: the marker "does not exist" by existsSync, so the old + // code wrote through it. An occupied path now reports not-first-time. + assert.equal(firstTime("rubric", "walm644"), false); + + assert.equal(existsSync(box.target), false, "symlink target must stay uncreated"); + assert.ok(lstatSync(marker).isSymbolicLink(), "the planted symlink is left alone"); + assert.deepEqual(treeOf(box.elsewhere), [], "nothing was written outside the state dir"); +}); + +test("bumpCounter refuses a symlinked counter path", { skip: skipOnWindows }, async (t) => { + const box = sandbox(t); + const { bumpCounter, stateDir } = await import(HOOK_IO); + + assert.equal(stateDir(), box.stateDir); + const counter = join(box.stateDir, "count_nudge_walm644"); + symlinkSync(box.target, counter); + + // Still answers, still never throws — it just keeps the count in memory. + assert.equal(bumpCounter("nudge", "walm644"), 1); + assert.equal(bumpCounter("nudge", "walm644"), 2); + + assert.equal(existsSync(box.target), false, "symlink target must stay uncreated"); + assert.ok(lstatSync(counter).isSymbolicLink()); + assert.deepEqual(treeOf(box.elsewhere), []); +}); + +test("normal session markers still work", { skip: skipOnWindows }, async (t) => { + const box = sandbox(t); + const { firstTime, bumpCounter, stateDir } = await import(HOOK_IO); + + assert.equal(firstTime("rubric", "session-a"), true); + assert.equal(firstTime("rubric", "session-a"), false); + assert.equal(firstTime("rubric", "session-a"), false); + // A different session is unaffected by the first one's marker. + assert.equal(firstTime("rubric", "session-b"), true); + + const marker = join(stateDir(), "rubric_session-a"); + assert.ok(lstatSync(marker).isFile(), "the marker is a plain file, not a link"); + assert.equal(readFileSync(marker, "utf8"), "1"); + + assert.equal(bumpCounter("turns", "session-a"), 1); + assert.equal(bumpCounter("turns", "session-a"), 2); + assert.equal(bumpCounter("turns", "session-a"), 3); + assert.equal(readFileSync(join(stateDir(), "count_turns_session-a"), "utf8"), "3"); + + assert.deepEqual(treeOf(box.elsewhere), []); +}); + +test("the state directory is private, and a symlinked one is refused", { skip: skipOnWindows }, async (t) => { + const box = sandbox(t); + const { firstTime, stateDir } = await import(HOOK_IO); + + const dir = stateDir(); + assert.equal(dir, box.stateDir); + const st = lstatSync(dir); + assert.ok(st.isDirectory()); + assert.equal(st.mode & 0o077, 0, "state dir must not be group/world accessible"); + + // Now stand a symlink where the state directory would be: the helper must + // refuse it outright instead of writing through it. + const hijacked = mkdtempSync(join(tmpdir(), "memwal-hookstate-hijack-")); + const decoy = join(hijacked, "tmp"); + mkdirSync(decoy); + symlinkSync(box.elsewhere, join(decoy, "memwal-hooks")); + process.env.TMPDIR = decoy; + t.after(() => rmSync(hijacked, { recursive: true, force: true })); + + assert.equal(stateDir(), null, "a symlinked state dir is not usable"); + // Degrades to in-process state: still answers, still writes nothing. + assert.equal(firstTime("rubric", "hijacked"), true); + assert.equal(firstTime("rubric", "hijacked"), false); + assert.deepEqual(treeOf(box.elsewhere), []); +}); + +test("the real prompt hook does not write through a planted symlink", { skip: skipOnWindows }, (t) => { + const box = sandbox(t); + mkdirSync(box.stateDir, { recursive: true, mode: 0o700 }); + const sessionId = "walm644-e2e"; + symlinkSync(box.target, join(box.stateDir, `rubric_${sessionId}`)); + + const result = spawnSync(process.execPath, [HOOK], { + input: JSON.stringify({ + prompt: "Remember that I always use pnpm and my canary is cedar-wren-11.", + session_id: sessionId, + }), + encoding: "utf8", + env: { ...process.env, TMPDIR: box.temp }, + }); + + assert.equal(result.status, 0, result.stderr); + assert.ok(result.stdout.trim(), "the hook still emits its directive"); + JSON.parse(result.stdout); // well-formed, so the session is never blocked + assert.equal(existsSync(box.target), false, "symlink target must stay uncreated"); + assert.deepEqual(treeOf(box.elsewhere), []); +}); diff --git a/packages/mcp/test/login-failure-notice.test.mjs b/packages/mcp/test/login-failure-notice.test.mjs index 05c6cc8b9..07d1abcde 100644 --- a/packages/mcp/test/login-failure-notice.test.mjs +++ b/packages/mcp/test/login-failure-notice.test.mjs @@ -208,9 +208,36 @@ test("a sign-in that never completes is reported on the next tool call", async ( assert.match(text, /never completed/); assert.match(text, /left running through the/); assert.doesNotMatch(text, /usually works/); - assert.match(text, /already be registered on your account/); + assertKeepsTheStrandedKey(text); // Still tells them how to sign in, rather than replacing the instruction. assert.match(text, /memwal_login/); + // ...but not by promising the opposite of the notice above it: an approved + // key is reclaimed by a restart, so the blob cannot also sell "no restart". + assert.doesNotMatch(text, /no client restart/i); +}); + +/** + * A failed attempt leaves its key in the write-ahead record (WALM-332), so the + * notice must send an approved key to a restart, which reclaims it. Signing in + * again reuses that key and the dashboard cannot register it twice, and + * removing it from the dashboard throws away the registration the user paid + * for. Both the wrapper and the timeout reason inside it are checked. + */ +function assertKeepsTheStrandedKey(text) { + assert.match(text, /Restart the MCP client/); + assert.match(text, /reclaim/); + assert.doesNotMatch(text, /revoke/i); + for (const sentence of text.split(/(?<=\.)\s+/)) { + if (/dashboard/.test(sentence)) { + assert.match(sentence, /abandon/, `dashboard advice must be limited to abandoning: "${sentence}"`); + } + } +} + +test("the failure notice keeps a stranded key reclaimable", async () => { + const { loginFailureNotice } = await import("../dist/messages.js"); + assert.equal(loginFailureNotice(null), ""); + assertKeepsTheStrandedKey(loginFailureNotice("Login timed out after 1ms.")); }); test("a signed-in memwal_login timeout warns through the bridge", async (t) => { diff --git a/packages/mcp/test/login-handoff.test.mjs b/packages/mcp/test/login-handoff.test.mjs index d5a3a200f..2cfb22627 100644 --- a/packages/mcp/test/login-handoff.test.mjs +++ b/packages/mcp/test/login-handoff.test.mjs @@ -167,8 +167,12 @@ test("auth-required mode picks up credentials mid-session without a restart", as const init = await waitFor((m) => m.id === 1 && m.result); assert.equal(init.result.serverInfo.name, "memwal"); - // Pre-login discovery must expose the same safety metadata clients will - // receive after the bridge hands off to the remote relayer. + // Pre-login discovery is the cold-start FLOOR, not the post-handoff list: + // it carries only what the oldest supported relayer serves, plus the + // locally served memwal_login (see BASELINE_RELAYER_TOOLS). The relayer's + // own tools/list replaces it once the session is up, so a tool that has + // reached dev but not prod — memwal_remember_status today — is absent + // here on purpose. Assert the safety metadata of the floor itself. send({ jsonrpc: "2.0", id: 10, method: "tools/list", params: {} }); const listed = await waitFor((m) => m.id === 10 && m.result); const metadata = Object.fromEntries( diff --git a/packages/mcp/test/login-recovery-signing.test.mjs b/packages/mcp/test/login-recovery-signing.test.mjs new file mode 100644 index 000000000..8234c7974 --- /dev/null +++ b/packages/mcp/test/login-recovery-signing.test.mjs @@ -0,0 +1,95 @@ +/** + * WALM-332 — the client's signed request must match the relayer byte for byte. + * + * Recovery authenticates with a signature over a canonical message defined in + * `services/server/src/auth.rs`. The two implementations are in different + * languages and cannot share code, so the format is duplicated — and a subtle + * mismatch (a trimmed trailing separator, a dropped empty field) would compile, + * pass every other test, and fail in production as an opaque 401. + * + * The literal below is asserted verbatim in + * `routes::accounts::tests::whoami_recovery_request_canonical_message_is_stable`. + * Change one and this fails; change the format in auth.rs and both fail. + */ +import { test } from "node:test"; +import assert from "node:assert/strict"; +import { createHash } from "node:crypto"; +import { verifyAsync } from "@noble/ed25519"; + +const { canonicalRequestMessage, EMPTY_BODY_SHA256 } = await import("../dist/recovery.js"); +const { signMessage, hexToBytes } = await import("../dist/crypto.js"); +const { generateKeypair } = await import("../dist/crypto.js"); + +/** Must equal the Rust fixture exactly. */ +const PINNED = + "1700000000.GET./api/whoami." + + "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855." + + "550e8400-e29b-41d4-a716-446655440000."; + +test("the client builds the canonical message the relayer expects", () => { + const message = canonicalRequestMessage({ + timestamp: "1700000000", + method: "GET", + path: "/api/whoami", + bodyHash: EMPTY_BODY_SHA256, + nonce: "550e8400-e29b-41d4-a716-446655440000", + }); + + assert.equal(message, PINNED, "must match services/server/src/auth.rs byte for byte"); + assert.ok(message.endsWith("."), "the empty account id keeps a trailing separator"); + assert.equal(message.split(".").length - 1, 5, "six fields, five separators"); +}); + +test("the empty-body hash is the real sha256 of nothing", () => { + assert.equal(EMPTY_BODY_SHA256, createHash("sha256").update("").digest("hex")); + assert.equal( + EMPTY_BODY_SHA256, + "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", + "the well-known constant the server will compute for a bodyless GET", + ); +}); + +test("an omitted account id is empty, never the string 'undefined'", () => { + // A plain template interpolation of an absent value would produce + // "...undefined" here, which signs cleanly and is rejected by the server + // with no clue why. + const message = canonicalRequestMessage({ + timestamp: "1", + method: "GET", + path: "/p", + bodyHash: "h", + nonce: "n", + }); + assert.equal(message, "1.GET./p.h.n."); + assert.doesNotMatch(message, /undefined|null/); +}); + +test("the signature the client produces verifies against its public key", async () => { + const kp = await generateKeypair(); + const message = canonicalRequestMessage({ + timestamp: "1700000000", + method: "GET", + path: "/api/whoami", + bodyHash: EMPTY_BODY_SHA256, + nonce: "550e8400-e29b-41d4-a716-446655440000", + }); + + const sigHex = await signMessage(kp.privateKeyHex, message); + assert.match(sigHex, /^[0-9a-f]{128}$/, "Ed25519 signatures are 64 bytes"); + + const ok = await verifyAsync( + hexToBytes(sigHex), + new TextEncoder().encode(message), + hexToBytes(kp.publicKeyHex), + ); + assert.equal(ok, true, "the relayer must be able to verify what we signed"); + + // And it must not verify a tampered message — otherwise the assertion above + // proves nothing. + const tampered = await verifyAsync( + hexToBytes(sigHex), + new TextEncoder().encode(message.replace("/api/whoami", "/api/stats")), + hexToBytes(kp.publicKeyHex), + ); + assert.equal(tampered, false, "a different path must not verify"); +}); diff --git a/packages/mcp/test/login-recovery.test.mjs b/packages/mcp/test/login-recovery.test.mjs new file mode 100644 index 000000000..4873a8047 --- /dev/null +++ b/packages/mcp/test/login-recovery.test.mjs @@ -0,0 +1,380 @@ +/** + * WALM-332 — reclaiming a delegate key stranded by an interrupted login. + * + * The write-ahead record keeps the key alive; these cover turning it back + * into usable credentials, and the cases where we must NOT. + */ +import { test } from "node:test"; +import assert from "node:assert/strict"; +import http from "node:http"; +import { existsSync, mkdtempSync, mkdirSync, writeFileSync, readFileSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; + +const ACCOUNT = `0x${"a".repeat(64)}`; +const OWNER = `0x${"b".repeat(64)}`; +const PACKAGE = `0x${"c".repeat(64)}`; + +function freshHome() { + const home = mkdtempSync(join(tmpdir(), "memwal-recovery-")); + // HOME alone is not a sandbox. os.homedir() reads USERPROFILE on Windows + // and ignores HOME, and credsPath() checks for a project-local .memwal + // above the working directory before it ever consults the home directory — + // which here is the real checkout. MEMWAL_CREDS_DIR overrides both, and + // pointing it at the sandbox's .memwal keeps the paths below unchanged. + process.env.HOME = home; + process.env.USERPROFILE = home; + process.env.MEMWAL_CREDS_DIR = join(home, ".memwal"); + mkdirSync(join(home, ".memwal"), { recursive: true }); + return home; +} + +const pendingPath = (h) => join(h, ".memwal", "login-pending.json"); +const credsPath = (h) => join(h, ".memwal", "credentials.json"); + +/** A relayer that answers /api/whoami however the test wants. */ +function startWhoami(handler) { + const server = http.createServer((req, res) => { + const url = new URL(req.url, "http://127.0.0.1"); + if (url.pathname !== "/api/whoami") { + res.writeHead(404).end(); + return; + } + handler(req, res); + }); + return new Promise((r) => + server.listen(0, "127.0.0.1", () => + r({ server, url: `http://127.0.0.1:${server.address().port}` }), + ), + ); +} + +const okWhoami = (req, res) => { + // Assert the client proved possession rather than just asking nicely. + for (const h of ["x-public-key", "x-signature", "x-timestamp", "x-nonce"]) { + if (!req.headers[h]) { + res.writeHead(400).end(JSON.stringify({ missing: h })); + return; + } + } + res.writeHead(200, { "content-type": "application/json" }); + res.end(JSON.stringify({ account_id: ACCOUNT, owner: OWNER, package_id: PACKAGE })); +}; + +function writePending(home, relayerUrl, overrides = {}) { + const pending = { + delegatePrivateKey: "11".repeat(32), + delegatePublicKeyHex: "22".repeat(32), + delegateAddress: `0x${"3".repeat(64)}`, + relayerUrl, + label: "Recovery test", + createdAt: new Date().toISOString(), + version: 1, + ...overrides, + }; + writeFileSync(pendingPath(home), JSON.stringify(pending), { mode: 0o600 }); + return pending; +} + +const importRecovery = () => import(`../dist/recovery.js?t=${Date.now()}${Math.random()}`); + +test("a stranded key is reclaimed into usable credentials", async (t) => { + const home = freshHome(); + const { server, url } = await startWhoami(okWhoami); + t.after(() => { + server.close(); + rmSync(home, { recursive: true, force: true }); + }); + + const pending = writePending(home, url); + const { recoverPendingLogin } = await importRecovery(); + const result = await recoverPendingLogin(); + + assert.equal(result.outcome, "recovered"); + + const creds = JSON.parse(readFileSync(credsPath(home), "utf8")); + assert.equal(creds.accountId, ACCOUNT, "accountId comes from the relayer"); + assert.equal(creds.walletAddress, OWNER); + assert.equal(creds.packageId, PACKAGE); + assert.equal( + creds.delegatePrivateKey, + pending.delegatePrivateKey, + "the reclaimed key must be the one that was registered", + ); + assert.equal( + existsSync(pendingPath(home)), + false, + "pending record cleared once the key is safe", + ); +}); + +test("recovery never rolls back a newer sign-in", async (t) => { + const home = freshHome(); + const { server, url } = await startWhoami(okWhoami); + t.after(() => { + server.close(); + rmSync(home, { recursive: true, force: true }); + }); + + // Pending login started BEFORE the credentials currently on disk: the user + // gave up on it and signed in again. Adopting it would silently downgrade + // them to a key they already abandoned. + writePending(home, url, { createdAt: new Date(Date.now() - 60_000).toISOString() }); + const current = { + delegatePrivateKey: "99".repeat(32), + delegatePublicKeyHex: "88".repeat(32), + delegateAddress: `0x${"7".repeat(64)}`, + walletAddress: OWNER, + accountId: `0x${"d".repeat(64)}`, + packageId: PACKAGE, + relayerUrl: url, + createdAt: new Date().toISOString(), + version: 1, + }; + writeFileSync(credsPath(home), JSON.stringify(current), { mode: 0o600 }); + + const { recoverPendingLogin } = await importRecovery(); + const result = await recoverPendingLogin(); + + assert.equal(result.outcome, "superseded"); + const after = JSON.parse(readFileSync(credsPath(home), "utf8")); + assert.deepEqual(after, current, "existing credentials must be untouched"); + assert.ok(result.strandedPublicKey, "the abandoned key is still reported so it can be revoked"); +}); + +test("a rejected key is reported but never deleted", async (t) => { + const home = freshHome(); + // 401 is ambiguous — on testnet even a valid registered key is rejected + // for want of an account hint. Deleting here would destroy a paid key. + const { server, url } = await startWhoami((_req, res) => { + res.writeHead(401).end("{}"); + }); + t.after(() => { + server.close(); + rmSync(home, { recursive: true, force: true }); + }); + + writePending(home, url); + const { recoverPendingLogin, formatStrandedLoginNotice } = await importRecovery(); + const result = await recoverPendingLogin(); + + assert.equal(result.outcome, "rejected"); + assert.equal( + existsSync(pendingPath(home)), + true, + "the record must survive an ambiguous rejection", + ); + assert.equal(existsSync(credsPath(home)), false, "no credentials written"); + + const notice = formatStrandedLoginNotice(result); + assert.match(notice, /22{10}/, "the notice names the key so it can be revoked"); + // Signing in again reuses this key, and the dashboard's add_delegate_key + // aborts on one already registered, so that alone cannot recover an + // approved key. + assert.match( + notice, + /cannot\s+register a key that is already there/, + "must not promise that signing in again recovers a key the user approved", + ); +}); + +test("an unreachable relayer keeps the record for a later attempt", async (t) => { + const home = freshHome(); + t.after(() => rmSync(home, { recursive: true, force: true })); + + // Nothing is listening on this port. + writePending(home, "http://127.0.0.1:1"); + const { recoverPendingLogin } = await importRecovery(); + const result = await recoverPendingLogin(); + + assert.equal(result.outcome, "unavailable"); + assert.equal(existsSync(pendingPath(home)), true); +}); + +test("an expired pending record is discarded rather than recovered", async (t) => { + const home = freshHome(); + const { server, url } = await startWhoami(okWhoami); + t.after(() => { + server.close(); + rmSync(home, { recursive: true, force: true }); + }); + + writePending(home, url, { + createdAt: new Date(Date.now() - 25 * 60 * 60_000).toISOString(), + }); + const { recoverPendingLogin } = await importRecovery(); + const result = await recoverPendingLogin(); + + assert.equal(result.outcome, "no-pending"); + assert.equal(existsSync(pendingPath(home)), false, "expired record is cleaned up"); + assert.equal(existsSync(credsPath(home)), false); +}); + +test("no pending record is a silent no-op", async (t) => { + const home = freshHome(); + t.after(() => rmSync(home, { recursive: true, force: true })); + + const { recoverPendingLogin, formatStrandedLoginNotice } = await importRecovery(); + const result = await recoverPendingLogin(); + + assert.equal(result.outcome, "no-pending"); + assert.equal(formatStrandedLoginNotice(result), null); +}); + +/** + * The relayer freshness-checks `x-timestamp` against `Utc::now().timestamp()` + * — SECONDS. `String(Date.now())` is milliseconds, ~10^12, which is outside + * every drift window there will ever be, so whoami 401'd on every attempt and + * recovery could not have worked at all. + */ +test("whoami signs a Unix timestamp in seconds, not milliseconds", async (t) => { + const home = freshHome(); + let seen = null; + const { server, url } = await startWhoami((req, res) => { + seen = req.headers["x-timestamp"]; + okWhoami(req, res); + }); + t.after(() => { + server.close(); + rmSync(home, { recursive: true, force: true }); + }); + + writePending(home, url); + const { recoverPendingLogin } = await importRecovery(); + await recoverPendingLogin(); + + assert.match(seen ?? "", /^\d{10}$/, `expected 10-digit seconds, got ${seen}`); + const skew = Math.abs(Number(seen) - Math.floor(Date.now() / 1000)); + assert.ok(skew < 300, `timestamp is ${skew}s from now — outside the relayer's window`); +}); + +/** + * `rejected` tells the user to sign in again and revoke the key. That advice is + * actively harmful when the relayer merely could not reach Sui: the key is + * fine, and re-registering costs gas for nothing. + */ +for (const [label, status, headers] of [ + ["a 503 with AUTH_UPSTREAM_UNAVAILABLE", 503, { "x-auth-error": "AUTH_UPSTREAM_UNAVAILABLE" }], + ["a bare 500", 500, {}], + ["a 429", 429, {}], + ["a 404 from a relayer without the route", 404, {}], +]) { + test(`${label} is retryable, not a rejection`, async (t) => { + const home = freshHome(); + const { server, url } = await startWhoami((_req, res) => { + res.writeHead(status, headers).end("{}"); + }); + t.after(() => { + server.close(); + rmSync(home, { recursive: true, force: true }); + }); + + writePending(home, url); + const { recoverPendingLogin, formatStrandedLoginNotice } = await importRecovery(); + const result = await recoverPendingLogin(); + + assert.equal(result.outcome, "unavailable", `status ${status} should not read as a denial`); + assert.equal(existsSync(pendingPath(home)), true, "the record must survive"); + + const notice = formatStrandedLoginNotice(result); + assert.doesNotMatch( + notice, + /revoke/i, + "must not send the user to revoke a key that may be perfectly good", + ); + assert.match(notice, /retried/i, "should say it will be retried"); + }); +} + +test("a 401 carrying AUTH_UPSTREAM_UNAVAILABLE is still retryable", async (t) => { + // The status alone is not enough: the header is what distinguishes + // "we could not check" from "we checked and said no". + const home = freshHome(); + const { server, url } = await startWhoami((_req, res) => { + res.writeHead(401, { "x-auth-error": "AUTH_UPSTREAM_UNAVAILABLE" }).end("{}"); + }); + t.after(() => { + server.close(); + rmSync(home, { recursive: true, force: true }); + }); + + writePending(home, url); + const { recoverPendingLogin } = await importRecovery(); + assert.equal((await recoverPendingLogin()).outcome, "unavailable"); +}); + +test("a 200 that is not a whoami body is retryable, not a rejection", async (t) => { + // Means we are not talking to the endpoint we think we are — nothing has + // denied this key. + const home = freshHome(); + const { server, url } = await startWhoami((_req, res) => { + res.writeHead(200, { "content-type": "application/json" }).end('{"hello":"world"}'); + }); + t.after(() => { + server.close(); + rmSync(home, { recursive: true, force: true }); + }); + + writePending(home, url); + const { recoverPendingLogin } = await importRecovery(); + assert.equal((await recoverPendingLogin()).outcome, "unavailable"); +}); + +test("a plain 401 is still a rejection", async (t) => { + // Regression guard on the split above: widening `unavailable` must not + // swallow the one case where the relayer really did deny the identity. + const home = freshHome(); + const { server, url } = await startWhoami((_req, res) => { + res.writeHead(401).end("{}"); + }); + t.after(() => { + server.close(); + rmSync(home, { recursive: true, force: true }); + }); + + writePending(home, url); + const { recoverPendingLogin } = await importRecovery(); + assert.equal((await recoverPendingLogin()).outcome, "rejected"); +}); + +/** + * WALM-646 — the supersede guard is account-agnostic, so a sign-in for a + * second account that was abandoned after the wallet step still wins on + * timestamp and recovery adopts it. Nothing is lost (the displaced file is + * backed up), but the user is moved between accounts and has to be told. + */ +test("recovery says so when it switches the active account", async (t) => { + const home = freshHome(); + const { server, url } = await startWhoami(okWhoami); + t.after(() => { + server.close(); + rmSync(home, { recursive: true, force: true }); + }); + + const accountA = `0x${"e".repeat(64)}`; + writeFileSync( + credsPath(home), + JSON.stringify({ + delegatePrivateKey: "99".repeat(32), + delegatePublicKeyHex: "88".repeat(32), + delegateAddress: `0x${"7".repeat(64)}`, + walletAddress: OWNER, + accountId: accountA, + packageId: PACKAGE, + relayerUrl: url, + createdAt: new Date(Date.now() - 60_000).toISOString(), + version: 1, + }), + { mode: 0o600 }, + ); + writePending(home, url); + + const { recoverPendingLogin } = await importRecovery(); + const result = await recoverPendingLogin(); + + assert.equal(result.outcome, "recovered"); + assert.equal(result.credentials.accountId, ACCOUNT); + assert.ok(result.replacementNotice, "switching accounts must not be silent"); + assert.match(result.replacementNotice, new RegExp(accountA), "names the account left behind"); + assert.match(result.replacementNotice, new RegExp(ACCOUNT), "names the account now in use"); +}); diff --git a/packages/mcp/test/login-write-ahead.test.mjs b/packages/mcp/test/login-write-ahead.test.mjs new file mode 100644 index 000000000..fa4667d51 --- /dev/null +++ b/packages/mcp/test/login-write-ahead.test.mjs @@ -0,0 +1,283 @@ +/** + * WALM-332 — the login flow must not lose the delegate private key. + * + * The browser registers the delegate key on-chain (a paid, irreversible + * action) and only then POSTs the callback that causes us to save it. If this + * process dies in that window, the in-memory keypair is destroyed and the + * user is left with an on-chain registration nobody holds the key to. + * + * The fix is write-ahead: persist the pending keypair BEFORE the browser is + * able to act on it, and clear it once credentials are safely saved. + */ +import { test } from "node:test"; +import assert from "node:assert/strict"; +import { existsSync, mkdtempSync, mkdirSync, readFileSync, writeFileSync, statSync, rmSync, chmodSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; + +const WEB = "https://memory.example"; +const RELAYER = "https://relayer.example"; + +function freshHome() { + const home = mkdtempSync(join(tmpdir(), "memwal-writeahead-")); + // HOME alone is not a sandbox. os.homedir() reads USERPROFILE on Windows + // and ignores HOME, and credsPath() checks for a project-local .memwal + // above the working directory before it ever consults the home directory — + // which here is the real checkout. MEMWAL_CREDS_DIR overrides both, and + // pointing it at the sandbox's .memwal keeps the paths below unchanged. + process.env.HOME = home; + process.env.USERPROFILE = home; + process.env.MEMWAL_CREDS_DIR = join(home, ".memwal"); + return home; +} + +const pendingPath = (home) => join(home, ".memwal", "login-pending.json"); +const credsPath = (home) => join(home, ".memwal", "credentials.json"); + +/** Start a login flow and resolve once the connect URL has been published. */ +async function startLogin(overrides = {}) { + const { loginFlow } = await import(`../dist/login.js?t=${Date.now()}${Math.random()}`); + let publishUrl; + const urlReady = new Promise((resolve) => { + publishUrl = resolve; + }); + const flow = loginFlow({ + webUrl: WEB, + relayerUrl: RELAYER, + label: "Write-ahead test", + timeoutMs: 4_000, + openBrowser: false, + onUrl: publishUrl, + ...overrides, + }); + // The flow rejects on timeout; nobody is going to complete it in these + // tests, so absorb it rather than tripping an unhandled rejection. + flow.catch(() => {}); + // A flow that fails BEFORE publishing — the write-ahead record cannot be + // written, say — must surface here as a rejection. Awaiting `urlReady` + // alone would hang forever on a URL that is never coming. + const failedEarly = flow.then(() => { + throw new Error("login resolved without ever publishing a URL"); + }); + failedEarly.catch(() => {}); + return { flow, url: new URL(await Promise.race([urlReady, failedEarly])) }; +} + +test("the delegate keypair is on disk before the browser is given the connect URL", async (t) => { + const home = freshHome(); + t.after(() => rmSync(home, { recursive: true, force: true })); + + const { url } = await startLogin(); + + // The URL is what the user clicks; by the time it exists, the browser can + // register this public key on-chain. The private half must already be safe. + assert.ok( + existsSync(pendingPath(home)), + "login-pending.json must exist by the time the connect URL is published", + ); + + const pending = JSON.parse(readFileSync(pendingPath(home), "utf8")); + const publicKeyInUrl = url.searchParams.get("publicKey"); + + assert.equal( + pending.delegatePublicKeyHex?.toLowerCase(), + publicKeyInUrl?.toLowerCase(), + "the persisted record must be for the exact key the browser was sent", + ); + assert.match( + pending.delegatePrivateKey ?? "", + /^(0x)?[0-9a-f]{64}$/i, + "the private key must be recoverable from the record", + ); + assert.equal(pending.relayerUrl, RELAYER); + assert.ok(pending.createdAt, "record needs a timestamp so it can expire"); + + // Same handling as credentials.json — owner-only. Windows does not enforce + // POSIX mode bits, and `savePendingLogin` treats `chmodSync` as best-effort + // there, so asserting them would test the platform rather than the code. + if (process.platform !== "win32") { + assert.equal( + statSync(pendingPath(home)).mode & 0o777, + 0o600, + "pending login must be owner-only, like credentials.json", + ); + } + + // Nothing has completed, so no credentials yet. + assert.equal(existsSync(credsPath(home)), false); +}); + +test("a completed login clears the pending record", async (t) => { + const home = freshHome(); + t.after(() => rmSync(home, { recursive: true, force: true })); + + const { flow, url } = await startLogin({ timeoutMs: 15_000 }); + const port = url.searchParams.get("port"); + const state = url.searchParams.get("connectState"); + const publicKey = url.searchParams.get("publicKey"); + + assert.ok(existsSync(pendingPath(home)), "precondition: pending record written"); + + const post = (path, body) => + fetch(`http://127.0.0.1:${port}${path}`, { + method: "POST", + headers: { "content-type": "application/json", origin: WEB }, + body: JSON.stringify(body), + }); + + await post("/preflight", { state, publicKey, relayer: RELAYER }); + await post("/callback", { + state, + accountId: `0x${"1".repeat(64)}`, + walletAddress: `0x${"2".repeat(64)}`, + packageId: `0x${"3".repeat(64)}`, + }); + + await flow; + + assert.equal(existsSync(credsPath(home)), true, "credentials should be saved"); + assert.equal( + existsSync(pendingPath(home)), + false, + "pending record must be cleared once the key is safely in credentials.json", + ); +}); + +/** + * Recovery only runs at process start, and is skipped for `--login` / + * `forceLogin`. So a login that times out, followed by `memwal_login` in the + * same process, used to mint a fresh keypair and overwrite the record — and if + * the browser had already paid for `add_delegate_key` on the first key, the + * private half went with it. + */ +test("a second login for the same relayer reuses the stranded keypair", async (t) => { + const home = freshHome(); + t.after(() => rmSync(home, { recursive: true, force: true })); + + const first = await startLogin(); + const stranded = JSON.parse(readFileSync(pendingPath(home), "utf8")); + first.flow.catch(() => {}); + + const second = await startLogin(); + const after = JSON.parse(readFileSync(pendingPath(home), "utf8")); + second.flow.catch(() => {}); + + assert.equal( + after.delegatePrivateKey, + stranded.delegatePrivateKey, + "the paid-for key must not be replaced by a second attempt", + ); + assert.equal( + second.url.searchParams.get("publicKey")?.toLowerCase(), + stranded.delegatePublicKeyHex.toLowerCase(), + "the browser should be sent the key that may already be registered", + ); + assert.equal( + after.createdAt, + stranded.createdAt, + "reusing must not extend the TTL past the attempt that may have registered it", + ); +}); + +test("a login against a different relayer does not reuse the record", async (t) => { + // A key registered against one relayer's account proves nothing to + // another, and recovery must never repoint a record at a new relayer. + const home = freshHome(); + t.after(() => rmSync(home, { recursive: true, force: true })); + + const first = await startLogin(); + const stranded = JSON.parse(readFileSync(pendingPath(home), "utf8")); + first.flow.catch(() => {}); + + const second = await startLogin({ relayerUrl: "https://other-relayer.example" }); + const after = JSON.parse(readFileSync(pendingPath(home), "utf8")); + second.flow.catch(() => {}); + + assert.notEqual(after.delegatePrivateKey, stranded.delegatePrivateKey); + assert.equal(after.relayerUrl, "https://other-relayer.example"); +}); + +test("a login refuses to start when the write-ahead record cannot be persisted", async (t) => { + // The whole invariant is that the key is durable before its public half can + // reach a browser that will pay to register it. Continuing anyway would + // publish the URL while only pretending to hold that. + const home = freshHome(); + const dir = join(home, ".memwal"); + mkdirSync(dir, { recursive: true }); + t.after(() => { + try { + chmodSync(dir, 0o700); + } catch { + /* nothing to restore */ + } + rmSync(home, { recursive: true, force: true }); + }); + + // Read-only directory. Root ignores mode bits, and Windows does not + // enforce them at all, so only assert where the setup actually bites. + chmodSync(dir, 0o500); + let writable = true; + try { + writeFileSync(join(dir, ".probe"), "x"); + } catch { + writable = false; + } + t.diagnostic(`credentials dir writable after chmod 0500: ${writable}`); + if (writable) { + t.skip("the sandbox directory is still writable — cannot provoke the failure here"); + return; + } + + await assert.rejects( + () => startLogin(), + /write-ahead/i, + "the login must fail loudly rather than publish a URL it cannot back", + ); + assert.equal(existsSync(pendingPath(home)), false, "nothing should have been written"); +}); + +/** + * `clearPendingLogin()` used to run only after a successful callback. CLI + * `--logout` and the `memwal_logout` tool both cleared `credentials.json` + * alone, so an interrupted re-login left the pending key behind and the next + * start's `recoverPendingLogin` signed the user straight back in — a logout + * that undid itself. + */ +test("logging out discards the pending record, not just the credentials", async (t) => { + const home = freshHome(); + t.after(() => rmSync(home, { recursive: true, force: true })); + + const { flow } = await startLogin(); + flow.catch(() => {}); + assert.ok(existsSync(pendingPath(home)), "precondition: a pending record exists"); + + const { main } = await import(`../dist/index.js?t=${Date.now()}${Math.random()}`); + await main(["--logout"]); + + assert.equal( + existsSync(pendingPath(home)), + false, + "an explicit logout must not leave a key that signs the user back in", + ); +}); + +test("clearing credentials on its own keeps the pending record", async (t) => { + // Discarding a key that may still be reclaimable is a decision only an + // explicit sign-out gets to make, which is why the pending clear lives in + // the logout paths rather than inside `clearCreds` (which is exported, and + // which a relayer 401 deliberately does not call). + const home = freshHome(); + t.after(() => rmSync(home, { recursive: true, force: true })); + + const { flow } = await startLogin(); + flow.catch(() => {}); + assert.ok(existsSync(pendingPath(home)), "precondition: a pending record exists"); + + const { clearCreds } = await import(`../dist/auth.js?t=${Date.now()}${Math.random()}`); + clearCreds(); + + assert.ok( + existsSync(pendingPath(home)), + "clearCreds must not discard a key that may still be reclaimable", + ); +}); diff --git a/packages/mcp/test/logout-invalidation.test.mjs b/packages/mcp/test/logout-invalidation.test.mjs index e8b7020f5..3adefb2c0 100644 --- a/packages/mcp/test/logout-invalidation.test.mjs +++ b/packages/mcp/test/logout-invalidation.test.mjs @@ -675,3 +675,57 @@ test("a non-tool request after logout is answered locally instead of hanging", a "nothing after logout should have reached the relayer", ); }); + +/** + * The tool half of the same fix the CLI gets in login-write-ahead: an explicit + * `memwal_logout` must discard `login-pending.json` too. Left behind, the next + * start's `recoverPendingLogin` rebuilds credentials from it and signs the user + * back in — a logout that undoes itself. + */ +test("memwal_logout discards a stranded pending login as well as the credentials", async (t) => { + const mock = await startMockRelayer(); + const home = mkdtempSync(join(tmpdir(), "memwal-logout-pending-")); + const credsPath = join(home, ".memwal", "credentials.json"); + const pendingPath = join(home, ".memwal", "login-pending.json"); + mkdirSync(dirname(credsPath), { recursive: true }); + writeFileSync(credsPath, JSON.stringify(makeCreds(mock.base)), { mode: 0o600 }); + writeFileSync( + pendingPath, + JSON.stringify({ + delegatePrivateKey: "11".repeat(32), + delegatePublicKeyHex: "22".repeat(32), + delegateAddress: `0x${"3".repeat(64)}`, + relayerUrl: mock.base, + label: "Interrupted re-login", + createdAt: new Date().toISOString(), + version: 1, + }), + { mode: 0o600 }, + ); + + t.after(() => { + mock.server.close(); + rmSync(home, { recursive: true, force: true }); + }); + + const { send, waitFor } = startBridge(t, mock, home); + + send({ jsonrpc: "2.0", id: 1, method: "initialize", params: {} }); + await waitFor((m) => m.id === 1 && m.result, 10_000); + + send({ + jsonrpc: "2.0", + id: 2, + method: "tools/call", + params: { name: "memwal_logout", arguments: {} }, + }); + const out = await waitFor((m) => m.id === 2, 10_000); + assert.notEqual(out.result?.isError, true, "logout should succeed"); + + assert.equal(existsSync(credsPath), false, "credentials should be gone"); + assert.equal( + existsSync(pendingPath), + false, + "the pending record must go too, or the next start signs the user back in", + ); +}); diff --git a/packages/mcp/test/memory-policy.test.mjs b/packages/mcp/test/memory-policy.test.mjs new file mode 100644 index 000000000..bdad0697c --- /dev/null +++ b/packages/mcp/test/memory-policy.test.mjs @@ -0,0 +1,172 @@ +/** + * The secret-exclusion rules have to reach the model through whichever channel + * a given client actually reads, and they have to say the same thing on all of + * them (WALM-642). + * + * Three channels, three packages, no workspace link between them: + * - `instructions` on initialize — the only one that survives lazy tool + * loading, and the one Claude Desktop / Codex rely on; + * - tool descriptions — what a client shows once tools ARE loaded; + * - the plugin's lifecycle hooks — the Claude Code / Codex install path, + * which never loads this package's `dist/` at all. + * + * So the text is duplicated by necessity. These tests are what stops the + * duplicates drifting: the block is extracted from each file on disk and the + * bytes compared. Edit one copy and this fails until the others match. + */ +import { test } from "node:test"; +import assert from "node:assert/strict"; +import { existsSync, readFileSync } from "node:fs"; +import { dirname, resolve } from "node:path"; +import { fileURLToPath } from "node:url"; + +import { + SECRET_EXCLUSION_RULES, + SECRET_EXCLUSION_SUMMARY, + AUTO_SAVE_OPT_IN_RULE, + MEMORY_POLICY_VERSION, +} from "../dist/memory-policy.js"; +import * as hookPolicy from "../plugin/scripts/lib/memory-policy.mjs"; +import { + buildProactiveInstructions, + PROACTIVE_INSTRUCTIONS, +} from "../dist/instructions.js"; +import { TOOL_DEFINITIONS } from "../dist/auth-required.js"; + +const __dirname = dirname(fileURLToPath(import.meta.url)); + +const COPIES = { + "packages/mcp/src/memory-policy.ts": resolve(__dirname, "../src/memory-policy.ts"), + "packages/mcp/plugin/scripts/lib/memory-policy.mjs": resolve( + __dirname, + "../plugin/scripts/lib/memory-policy.mjs", + ), + // Lives in the standalone `memwal-server-scripts` package. Present in the + // monorepo, absent from an npm-only checkout of this package — hence the + // existence guard rather than a hard path assumption. + "services/server/scripts/mcp/tools/memory-policy.ts": resolve( + __dirname, + "../../../services/server/scripts/mcp/tools/memory-policy.ts", + ), +}; + +const START = "// ─── memwal:policy-block:start"; +const END = "// ─── memwal:policy-block:end"; + +/** The shared region of one copy, markers included, as raw source bytes. */ +function policyBlock(path) { + const src = readFileSync(path, "utf8"); + const start = src.indexOf(START); + const end = src.indexOf(END); + assert.notEqual(start, -1, `${path} has no policy-block start marker`); + assert.notEqual(end, -1, `${path} has no policy-block end marker`); + const endOfLine = src.indexOf("\n", end); + return src.slice(start, endOfLine === -1 ? undefined : endOfLine); +} + +test("every copy of the policy block is byte-identical", () => { + const present = Object.entries(COPIES).filter(([, path]) => existsSync(path)); + assert.ok( + present.length >= 2, + "at least the two in-package copies must exist for this test to mean anything", + ); + + const [firstName, firstPath] = present[0]; + const reference = policyBlock(firstPath); + for (const [name, path] of present.slice(1)) { + assert.equal( + policyBlock(path), + reference, + `${name} has drifted from ${firstName} — copy the block across verbatim, markers included`, + ); + } +}); + +test("the relayer sidecar's copy is present in the monorepo", () => { + // Guarded above so an npm-only checkout still passes; asserted here so the + // monorepo cannot quietly lose the third copy and leave the comparison + // running over two files that happen to agree. + const path = COPIES["services/server/scripts/mcp/tools/memory-policy.ts"]; + if (!existsSync(resolve(__dirname, "../../../services"))) return; + assert.ok(existsSync(path), "the sidecar copy of the policy block is missing"); +}); + +test("the compiled client copy and the hook copy agree at runtime", () => { + // The byte comparison above covers the source. This covers what each side + // actually evaluates to, so a stray escape or join separator is caught too. + assert.equal(hookPolicy.SECRET_EXCLUSION_RULES, SECRET_EXCLUSION_RULES); + assert.equal(hookPolicy.SECRET_EXCLUSION_SUMMARY, SECRET_EXCLUSION_SUMMARY); + assert.equal(hookPolicy.AUTO_SAVE_OPT_IN_RULE, AUTO_SAVE_OPT_IN_RULE); + assert.equal(hookPolicy.MEMORY_POLICY_VERSION, MEMORY_POLICY_VERSION); +}); + +test("the rules name every credential class the ticket lists", () => { + for (const term of [ + /passwords/i, + /API keys/i, + /tokens/i, + /private keys/i, + /seed or recovery phrases/i, + /authorization\s*\n?\s*headers/i, + /session cookies/i, + /user:password/i, + ]) { + assert.match(SECRET_EXCLUSION_RULES, term); + } + // The two non-credential rules the ticket asks for by name. + assert.match(SECRET_EXCLUSION_RULES, /do not save it/i); + assert.match(SECRET_EXCLUSION_RULES, /third-party material/i); + // And the instruction that makes a mixed message salvageable rather than + // dropped — the difference between "save the preference" and "save nothing". + assert.match(SECRET_EXCLUSION_RULES, /save the preference alone/i); +}); + +test("both instruction variants carry the rules verbatim", () => { + const automatic = buildProactiveInstructions({ autoSave: true }); + const manual = buildProactiveInstructions({ autoSave: false }); + assert.ok(automatic.includes(SECRET_EXCLUSION_RULES)); + assert.ok(manual.includes(SECRET_EXCLUSION_RULES)); + assert.equal(PROACTIVE_INSTRUCTIONS, automatic); +}); + +test("the instruction variants differ on unprompted saving and nothing else", () => { + const automatic = buildProactiveInstructions({ autoSave: true }); + const manual = buildProactiveInstructions({ autoSave: false }); + + assert.match(automatic, /automatic memory ON/i); + assert.match(automatic, /Do not ask whether to save it/i); + + assert.match(manual, /automatic memory is OFF/i); + assert.match(manual, /do NOT save anything they did not ask/i); + assert.doesNotMatch(manual, /Do not ask whether to save it/i); + + // Recall is not gated — it reads, it does not write. + for (const text of [automatic, manual]) { + assert.match(text, /RECALL: before answering/); + assert.match(text, /memwal_restore/); + assert.match(text, /never substitute your own memory/); + } +}); + +test("the cold-start write tools state the rules too", () => { + // A client that lazily loads schemas sees these and not `instructions`; + // a client that eagerly lists sees both. Either way the rules are there. + for (const name of ["memwal_remember", "memwal_remember_bulk", "memwal_analyze"]) { + const tool = TOOL_DEFINITIONS.find((t) => t.name === name); + assert.ok(tool, `missing ${name}`); + assert.ok( + tool.description.includes(SECRET_EXCLUSION_RULES), + `${name} does not state the shared secret-exclusion rules`, + ); + assert.ok( + tool.description.includes(AUTO_SAVE_OPT_IN_RULE), + `${name} does not state that automatic saving is opt-in`, + ); + } +}); + +test("read-only tools are not burdened with write rules", () => { + const recall = TOOL_DEFINITIONS.find((t) => t.name === "memwal_recall"); + assert.ok(recall); + assert.ok(!recall.description.includes(SECRET_EXCLUSION_RULES)); +}); diff --git a/packages/mcp/test/orphaned-call-health.test.mjs b/packages/mcp/test/orphaned-call-health.test.mjs new file mode 100644 index 000000000..86f131c83 --- /dev/null +++ b/packages/mcp/test/orphaned-call-health.test.mjs @@ -0,0 +1,318 @@ +/** + * A sent call whose reply never arrives is answered with the relayer's health + * (WALM-396). + * + * Before: "Walrus Memory did not answer this call … safe to retry", with no + * way to tell a dead relayer from a wrong URL from one stuck call. Now the + * bridge asks `/health` first and says which it is. The mock keeps the SSE + * session healthy and swallows `memwal_recall`, so the per-call deadline is + * what fires; each case decides how `/health` behaves. + */ +import { test } from "node:test"; +import assert from "node:assert/strict"; +import http from "node:http"; +import { spawn } from "node:child_process"; +import { mkdtempSync, mkdirSync, writeFileSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join, dirname, resolve } from "node:path"; +import { fileURLToPath } from "node:url"; + +const __dirname = dirname(fileURLToPath(import.meta.url)); +const BIN = resolve(__dirname, "../dist/bin/memwal-mcp.js"); +const EXPECTED_BEARER = "a".repeat(64); +const EXPECTED_ACCOUNT_ID = "0x" + "3".repeat(64); + +function hasBridgeAuth(req) { + return ( + req.headers.authorization === `Bearer ${EXPECTED_BEARER}` && + req.headers["x-memwal-account-id"] === EXPECTED_ACCOUNT_ID + ); +} + +/** `onHealth(req, res, mock)` answers `GET /health`; everything else is a + * healthy relayer that never replies to `memwal_recall`. */ +function startMockRelayer(onHealth) { + const sessions = new Map(); + let sseGetCount = 0; + let recallPosts = 0; + let swallowed = null; + const mock = {}; + const server = http.createServer((req, res) => { + const url = new URL(req.url, "http://127.0.0.1"); + if (req.method === "GET" && url.pathname === "/version") { + res.writeHead(200, { "content-type": "application/json" }); + res.end( + JSON.stringify({ + apiVersion: "1.0.0", + relayerVersion: "1.0.0", + minSupportedSdk: { mcp: "0.0.1" }, + }), + ); + return; + } + if (req.method === "GET" && url.pathname === "/health") { + onHealth(req, res, mock); + return; + } + if (req.method === "GET" && url.pathname === "/api/mcp/sse") { + if (!hasBridgeAuth(req)) { + res.writeHead(401).end(); + return; + } + sseGetCount += 1; + const sessionId = `session-${sseGetCount}`; + res.writeHead(200, { + "content-type": "text/event-stream", + "cache-control": "no-cache", + connection: "keep-alive", + }); + res.write(`event: endpoint\ndata: /api/mcp/messages?sessionId=${sessionId}\n\n`); + sessions.set(sessionId, { res }); + const hb = setInterval(() => { + if (res.writableEnded) return clearInterval(hb); + res.write(":keepalive\n\n"); + }, 200); + hb.unref?.(); + res.on("close", () => clearInterval(hb)); + return; + } + if (req.method === "POST" && url.pathname === "/api/mcp/messages") { + if (!hasBridgeAuth(req)) { + res.writeHead(401).end(); + return; + } + const session = sessions.get(url.searchParams.get("sessionId")); + let body = ""; + req.on("data", (c) => (body += c)); + req.on("end", () => { + if (!session) { + res.writeHead(404).end(); + return; + } + res.writeHead(202).end(); + let msg; + try { + msg = JSON.parse(body); + } catch { + return; + } + if (msg.method === "initialize") { + session.res.write( + `event: message\ndata: ${JSON.stringify({ + jsonrpc: "2.0", + id: msg.id, + result: { + protocolVersion: "2024-11-05", + capabilities: { tools: { listChanged: true } }, + serverInfo: { name: "memwal", version: "0.0.1" }, + }, + })}\n\n`, + ); + return; + } + if (msg.method === "tools/call" && msg.params?.name === "memwal_recall") { + recallPosts += 1; + swallowed = { session, id: msg.id }; + } + }); + return; + } + res.writeHead(404).end(); + }); + Object.assign(mock, { + server, + getSseGetCount: () => sseGetCount, + getRecallPosts: () => recallPosts, + /** End every open SSE stream, as a relayer restart or a proxy would. + * The bridge sees EOF and reconnects. */ + dropSse: () => { + for (const [id, session] of sessions) { + sessions.delete(id); + session.res.end(); + } + }, + releaseSwallowed: () => { + if (!swallowed) return false; + swallowed.session.res.write( + `event: message\ndata: ${JSON.stringify({ + jsonrpc: "2.0", + id: swallowed.id, + result: { content: [{ type: "text", text: "LATE_REPLY" }], isError: false }, + })}\n\n`, + ); + return true; + }, + }); + return new Promise((r) => { + server.listen(0, "127.0.0.1", () => { + mock.base = `http://127.0.0.1:${server.address().port}`; + r(mock); + }); + }); +} + +/** Spawn the bridge against `mock`, initialize it, and send one recall. */ +async function recallAgainst(t, mock) { + const home = mkdtempSync(join(tmpdir(), "memwal-orphan-health-")); + const credsPath = join(home, ".memwal", "credentials.json"); + mkdirSync(dirname(credsPath), { recursive: true }); + writeFileSync( + credsPath, + JSON.stringify({ + delegatePrivateKey: EXPECTED_BEARER, + delegatePublicKeyHex: "b".repeat(64), + delegateAddress: "0x" + "1".repeat(64), + walletAddress: "0x" + "2".repeat(64), + accountId: EXPECTED_ACCOUNT_ID, + packageId: "0x" + "4".repeat(64), + relayerUrl: mock.base, + label: "Orphan Health Test", + createdAt: new Date(0).toISOString(), + version: 1, + }), + { mode: 0o600 }, + ); + + const child = spawn(process.execPath, [BIN, "--relayer", mock.base, "--web-url", mock.base], { + env: { + ...process.env, + HOME: home, + USERPROFILE: home, + MEMWAL_MCP_SSE_IDLE_MS: "30000", + MEMWAL_MCP_CALL_TIMEOUT_MS: "1500", + MEMWAL_MCP_HEALTH_PROBE_MS: "2000", + }, + stdio: ["pipe", "pipe", "pipe"], + }); + t.after(() => { + child.kill("SIGKILL"); + mock.server.closeAllConnections(); + mock.server.close(); + rmSync(home, { recursive: true, force: true }); + }); + + const received = []; + const listeners = new Set(); + let buf = ""; + child.stdout.on("data", (d) => { + buf += d.toString(); + let nl; + while ((nl = buf.indexOf("\n")) >= 0) { + const line = buf.slice(0, nl); + buf = buf.slice(nl + 1); + if (!line.trim()) continue; + let msg; + try { + msg = JSON.parse(line); + } catch { + continue; + } + received.push(msg); + for (const l of [...listeners]) l(msg); + } + }); + let stderrBuf = ""; + child.stderr.on("data", (d) => (stderrBuf += d.toString())); + const send = (obj) => child.stdin.write(JSON.stringify(obj) + "\n"); + const waitFor = (pred, ms = 15000) => { + const hit = received.find(pred); + if (hit) return Promise.resolve(hit); + return new Promise((res, rej) => { + const timer = setTimeout(() => { + listeners.delete(l); + rej(new Error(`timed out waiting for message\n--- stderr ---\n${stderrBuf}`)); + }, ms); + const l = (m) => { + if (!pred(m)) return; + clearTimeout(timer); + listeners.delete(l); + res(m); + }; + listeners.add(l); + }); + }; + + send({ jsonrpc: "2.0", id: 1, method: "initialize", params: {} }); + await waitFor((m) => m.id === 1 && m.result, 10_000); + send({ + jsonrpc: "2.0", + id: 2, + method: "tools/call", + params: { name: "memwal_recall", arguments: { query: "anything" } }, + }); + const reply = await waitFor((m) => m.id === 2, 15_000); + return { reply, received, stderr: () => stderrBuf }; +} + +const textOf = (reply) => reply.result?.content?.[0]?.text ?? ""; + +test("a lost recall reply on a healthy relayer says the relayer is up and a retry is safe", async (t) => { + const mock = await startMockRelayer((_req, res) => { + res.writeHead(200, { "content-type": "application/json" }); + res.end(JSON.stringify({ status: "ok", version: "9.9.9" })); + }); + const { reply, stderr } = await recallAgainst(t, mock); + const text = textOf(reply); + assert.equal(reply.result?.isError, true); + assert.match(text, /did not answer this call/); + assert.match(text, /^Relayer health: ok \(\d+ms, v9\.9\.9\)/m); + assert.match(text, /safe to retry/i); + assert.match(text, /^Cause: /m); + assert.match(text, /^Next step: /m); + // The operator's log carries the same verdict as the agent's message. + assert.match(stderr(), /"health":"ok"/); + assert.equal(mock.getSseGetCount(), 1, "the per-call deadline, not a reconnect, answered it"); +}); + +test("a lost recall reply on an unhealthy relayer says so and asks to wait", async (t) => { + const mock = await startMockRelayer((_req, res) => res.writeHead(503).end()); + const text = textOf((await recallAgainst(t, mock)).reply); + assert.match(text, /^Relayer health: HTTP 503/m); + assert.match(text, /wait/i); +}); + +test("a lost recall reply with the health check failing says the relayer is unreachable", async (t) => { + const mock = await startMockRelayer((req) => req.socket.destroy()); + const text = textOf((await recallAgainst(t, mock)).reply); + assert.match(text, /^Relayer health: unreachable/m); + assert.match(text, /down or not reachable/); +}); + +test("a reply that lands while the health check runs is delivered, and nothing else", async (t) => { + // The deadline fired, the probe is in flight, and then the real answer + // arrives. It must win: answering with an error as well would be a + // second response for the same id. + const mock = await startMockRelayer((_req, res, self) => { + self.releaseSwallowed(); + setTimeout(() => { + res.writeHead(200, { "content-type": "application/json" }); + res.end(JSON.stringify({ status: "ok", version: "9.9.9" })); + }, 500); + }); + const { reply, received } = await recallAgainst(t, mock); + assert.equal(textOf(reply), "LATE_REPLY"); + assert.notEqual(reply.result?.isError, true); + // Give the probe time to settle and prove it writes nothing. + await new Promise((r) => setTimeout(r, 1500)); + assert.equal(received.filter((m) => m.id === 2).length, 1); +}); + +test("a reconnect while the health check runs does not send the expired call again", async (t) => { + // The deadline fired, the probe is in flight, and then the stream dies. + // The reconnect replays what is still in flight — but not this call: the + // sweeper is already answering it as failed, and a second POST would run + // it twice, which for a write with no idempotency key stores it twice. + const mock = await startMockRelayer((_req, res, self) => { + self.dropSse(); + setTimeout(() => { + res.writeHead(200, { "content-type": "application/json" }); + res.end(JSON.stringify({ status: "ok", version: "9.9.9" })); + }, 1200); + }); + const { reply, received } = await recallAgainst(t, mock); + assert.equal(reply.result?.isError, true); + assert.match(textOf(reply), /did not answer this call/); + assert.ok(mock.getSseGetCount() >= 2, "the stream really was reopened"); + assert.equal(mock.getRecallPosts(), 1, "the expired call was not posted a second time"); + assert.equal(received.filter((m) => m.id === 2).length, 1); +}); diff --git a/packages/mcp/test/project-creds-approval.test.mjs b/packages/mcp/test/project-creds-approval.test.mjs new file mode 100644 index 000000000..74e361bf6 --- /dev/null +++ b/packages/mcp/test/project-creds-approval.test.mjs @@ -0,0 +1,411 @@ +/** + * Repo credentials cannot silently choose the destination (WALM-639). + * + * `.memwal/credentials.json` decides the account a memory is written under and + * the relayer it is written through — and it lives INSIDE the repository, where + * anyone who can commit, or who can get a clone opened, can put one. Presence + * alone used to be the opt-in, so opening a project repointed every memory + * written from it, with nothing said. The writes are immutable and there is no + * delete path, which is what made this a P1 rather than a papercut. + * + * The gate: a project file is inert until the user approves that exact file, + * account, delegate and relayer, and the approval lives OUTSIDE the repository + * so a repo cannot carry its own approval. Anything unapproved falls back to + * the global credentials rather than failing, because a machine that never had + * a project file must keep behaving exactly as it did. + * + * `auth.js` resolves per call, so each test sets HOME and cwd first and then + * imports with a cache-busting query — the pattern used by + * credential-resolution.test.mjs. + */ +import { test } from "node:test"; +import assert from "node:assert/strict"; +import { + mkdtempSync, + mkdirSync, + writeFileSync, + readFileSync, + rmSync, + existsSync, + realpathSync, +} from "node:fs"; +import { tmpdir } from "node:os"; +import { join, dirname } from "node:path"; + +const GLOBAL_ACCOUNT = "0x" + "a".repeat(64); +const PROJECT_ACCOUNT = "0x" + "b".repeat(64); +const ATTACKER_ACCOUNT = "0x" + "9".repeat(64); +const PRIVATE_KEY = "c".repeat(64); +const GLOBAL_RELAYER = "https://relayer.example"; +const PROJECT_RELAYER = "https://project-relayer.example"; + +function makeCreds(overrides = {}) { + return { + delegatePrivateKey: PRIVATE_KEY, + delegatePublicKeyHex: "d".repeat(64), + delegateAddress: "0x" + "e".repeat(64), + walletAddress: "0x" + "f".repeat(64), + accountId: GLOBAL_ACCOUNT, + packageId: "0x" + "1".repeat(64), + relayerUrl: GLOBAL_RELAYER, + createdAt: new Date(0).toISOString(), + version: 1, + ...overrides, + }; +} + +function writeCredsAt(root, creds) { + const path = join(root, ".memwal", "credentials.json"); + mkdirSync(dirname(path), { recursive: true }); + writeFileSync(path, JSON.stringify(creds), { mode: 0o600 }); + return path; +} + +/** + * A HOME, a working directory, and the module re-imported so it sees them. + * + * `MEMWAL_CREDS_DIR` is cleared rather than inherited: it overrides resolution + * outright, so a stray value in the ambient environment would make every + * assertion here vacuous. + */ +async function sandbox(t, { global: globalCreds, project: projectCreds } = {}) { + // Canonicalised for the same reason as credential-resolution.test.mjs: + // `process.cwd()` and `homedir()` report resolved paths, and on macOS + // `/tmp` is a symlink. Both HOME and USERPROFILE, so it is portable. + const home = realpathSync(mkdtempSync(join(tmpdir(), "memwal-approve-home-"))); + const cwd = realpathSync(mkdtempSync(join(tmpdir(), "memwal-approve-cwd-"))); + const previous = { + home: process.env.HOME, + profile: process.env.USERPROFILE, + credsDir: process.env.MEMWAL_CREDS_DIR, + cwd: process.cwd(), + }; + + process.env.HOME = home; + process.env.USERPROFILE = home; + delete process.env.MEMWAL_CREDS_DIR; + process.chdir(cwd); + + if (globalCreds) writeCredsAt(home, globalCreds); + if (projectCreds) writeCredsAt(cwd, projectCreds); + + t.after(() => { + process.chdir(previous.cwd); + process.env.HOME = previous.home; + process.env.USERPROFILE = previous.profile; + if (previous.credsDir === undefined) delete process.env.MEMWAL_CREDS_DIR; + else process.env.MEMWAL_CREDS_DIR = previous.credsDir; + rmSync(home, { recursive: true, force: true }); + rmSync(cwd, { recursive: true, force: true }); + }); + + const auth = await import(`../dist/auth.js?walm639=${Date.now()}-${Math.random()}`); + return { auth, home, cwd }; +} + +/** The two files a sandbox works with. */ +const globalFile = (home) => join(home, ".memwal", "credentials.json"); +const projectFile = (cwd) => join(cwd, ".memwal", "credentials.json"); + +/* --------------------------------------------------------------------- * + * The reproduction itself. + * --------------------------------------------------------------------- */ + +test("a repo credentials file alone does not redirect a write", async (t) => { + const { auth, home, cwd } = await sandbox(t, { + global: makeCreds(), + project: makeCreds({ accountId: PROJECT_ACCOUNT, relayerUrl: PROJECT_RELAYER }), + }); + + // Reads resolve to the user's own account... + assert.equal(auth.credsPath(), globalFile(home)); + assert.equal(auth.loadCreds()?.accountId, GLOBAL_ACCOUNT); + assert.equal(auth.loadCreds()?.relayerUrl, GLOBAL_RELAYER); + + // ...and so do writes. This is the assertion the ticket asks for: the file + // the repo carries must not be the file the process signs and saves with. + auth.saveCreds(makeCreds({ label: "Re-saved" })); + assert.equal(JSON.parse(readFileSync(globalFile(home), "utf8")).label, "Re-saved"); + const untouched = JSON.parse(readFileSync(projectFile(cwd), "utf8")); + assert.equal(untouched.accountId, PROJECT_ACCOUNT, "the repo file must not be written"); + assert.equal(untouched.label, undefined); +}); + +test("running from a subfolder does not redirect a write either", async (t) => { + // The reporter checked this: the walk that finds the project file climbs, + // so the gate has to hold at every depth, not just at the project root. + const { auth, home, cwd } = await sandbox(t, { + global: makeCreds(), + project: makeCreds({ accountId: PROJECT_ACCOUNT }), + }); + const deep = join(cwd, "src", "nested"); + mkdirSync(deep, { recursive: true }); + process.chdir(deep); + + assert.equal(auth.credsPath(), globalFile(home)); + assert.equal(auth.loadCreds()?.accountId, GLOBAL_ACCOUNT); + assert.equal(auth.resolveCreds().project?.decision, "unapproved"); +}); + +test("the ignored project file is named, along with the destination it wanted", async (t) => { + // Falling back in silence would be the mirror image of the silent redirect: + // the user created that file expecting it to be used. + const { auth, home, cwd } = await sandbox(t, { + global: makeCreds(), + project: makeCreds({ accountId: PROJECT_ACCOUNT, relayerUrl: PROJECT_RELAYER }), + }); + + const notice = auth.formatProjectCredsNotice(); + + assert.ok(notice, "an ignored project file must be reported"); + assert.ok(notice.includes(projectFile(cwd)), "must name the file that was ignored"); + assert.ok(notice.includes(PROJECT_ACCOUNT), "must name the account it would have used"); + assert.ok(notice.includes(PROJECT_RELAYER), "must name the relayer it would have used"); + assert.ok(notice.includes(globalFile(home)), "must name where memory is going instead"); + assert.match(notice, /approve-project/, "must say how to approve it"); + assert.ok( + notice.includes(auth.projectApprovalsPath()), + "must say where the approval is recorded", + ); +}); + +test("nothing the user can see ever carries the delegate private key", async (t) => { + const { auth } = await sandbox(t, { + global: makeCreds(), + project: makeCreds({ accountId: PROJECT_ACCOUNT }), + }); + + assert.ok(!auth.formatProjectCredsNotice().includes(PRIVATE_KEY)); + auth.approveProjectCreds(); + const record = readFileSync(auth.projectApprovalsPath(), "utf8"); + assert.ok(!record.includes(PRIVATE_KEY), "the approval record must hold no key material"); +}); + +/* --------------------------------------------------------------------- * + * Approval. + * --------------------------------------------------------------------- */ + +test("an approved project file is the one that gets used", async (t) => { + const { auth, cwd } = await sandbox(t, { + global: makeCreds(), + project: makeCreds({ accountId: PROJECT_ACCOUNT, relayerUrl: PROJECT_RELAYER }), + }); + + const result = auth.approveProjectCreds(); + + assert.equal(result.outcome, "approved"); + assert.equal(result.accountId, PROJECT_ACCOUNT); + assert.equal(result.relayerUrl, PROJECT_RELAYER); + assert.equal(auth.credsPath(), projectFile(cwd)); + assert.equal(auth.loadCreds()?.accountId, PROJECT_ACCOUNT); + assert.equal(auth.resolveCreds().source, "project"); + assert.equal(auth.formatProjectCredsNotice(), null, "an approved file is not a warning"); +}); + +test("the approval is recorded outside the repository", async (t) => { + // The whole point: a record the repo could carry is a repo approving itself. + const { auth, home, cwd } = await sandbox(t, { + global: makeCreds(), + project: makeCreds({ accountId: PROJECT_ACCOUNT }), + }); + + auth.approveProjectCreds(); + + const approvals = auth.projectApprovalsPath(); + assert.equal(approvals, join(home, ".memwal", "project-approvals.json")); + assert.ok(!approvals.startsWith(cwd), "the approval must not live in the project"); + assert.equal(existsSync(approvals), true); +}); + +test("an approvals file committed inside the repo approves nothing", async (t) => { + const { auth, cwd, home } = await sandbox(t, { + global: makeCreds(), + project: makeCreds({ accountId: ATTACKER_ACCOUNT, relayerUrl: PROJECT_RELAYER }), + }); + + // Exactly what the repo would have to ship to self-approve: a well-formed + // record, in the project's own .memwal, naming its own credentials file. + const { createHash } = await import("node:crypto"); + const fingerprint = createHash("sha256") + .update(`${ATTACKER_ACCOUNT}\n0x${"e".repeat(64)}\n${PROJECT_RELAYER}`) + .digest("hex"); + writeFileSync( + join(cwd, ".memwal", "project-approvals.json"), + JSON.stringify({ + version: 1, + approvals: [ + { + path: projectFile(cwd), + fingerprint, + accountId: ATTACKER_ACCOUNT, + delegateAddress: "0x" + "e".repeat(64), + relayerUrl: PROJECT_RELAYER, + approvedAt: new Date().toISOString(), + }, + ], + }), + ); + + assert.equal(auth.credsPath(), globalFile(home), "a repo must not approve itself"); + assert.equal(auth.loadCreds()?.accountId, GLOBAL_ACCOUNT); + assert.equal(auth.resolveCreds().project?.decision, "unapproved"); +}); + +test("approving one project does not approve another", async (t) => { + const { auth, home, cwd } = await sandbox(t, { + global: makeCreds(), + project: makeCreds({ accountId: PROJECT_ACCOUNT }), + }); + auth.approveProjectCreds(); + + // A second checkout, with the same credentials in it. Approval is per + // project path, so opening this one is a fresh decision. + const other = realpathSync(mkdtempSync(join(tmpdir(), "memwal-approve-other-"))); + t.after(() => rmSync(other, { recursive: true, force: true })); + writeCredsAt(other, makeCreds({ accountId: PROJECT_ACCOUNT })); + process.chdir(other); + + assert.equal(auth.credsPath(), globalFile(home)); + assert.equal(auth.resolveCreds().project?.decision, "unapproved"); +}); + +/* --------------------------------------------------------------------- * + * Re-approval after the destination moves. + * --------------------------------------------------------------------- */ + +for (const [what, mutation] of [ + ["account", { accountId: ATTACKER_ACCOUNT }], + ["relayer", { relayerUrl: "https://attacker.example" }], + ["delegate key", { delegateAddress: "0x" + "7".repeat(64) }], +]) { + test(`changing the ${what} after approval requires approval again`, async (t) => { + const { auth, home, cwd } = await sandbox(t, { + global: makeCreds(), + project: makeCreds({ accountId: PROJECT_ACCOUNT, relayerUrl: PROJECT_RELAYER }), + }); + auth.approveProjectCreds(); + assert.equal(auth.credsPath(), projectFile(cwd), "precondition: approved and in use"); + + // A later commit edits the file the user already approved. + writeCredsAt( + cwd, + makeCreds({ + accountId: PROJECT_ACCOUNT, + relayerUrl: PROJECT_RELAYER, + ...mutation, + }), + ); + + assert.equal(auth.credsPath(), globalFile(home), "a moved destination must not be used"); + assert.equal(auth.resolveCreds().project?.decision, "changed"); + const notice = auth.formatProjectCredsNotice(); + assert.match(notice, /changed since you/, `notice did not report the change:\n${notice}`); + }); +} + +test("re-approving adopts the new destination and names the old one", async (t) => { + const { auth, cwd } = await sandbox(t, { + global: makeCreds(), + project: makeCreds({ accountId: PROJECT_ACCOUNT, relayerUrl: PROJECT_RELAYER }), + }); + auth.approveProjectCreds(); + writeCredsAt(cwd, makeCreds({ accountId: ATTACKER_ACCOUNT, relayerUrl: PROJECT_RELAYER })); + + const result = auth.approveProjectCreds(); + + assert.equal(result.outcome, "reapproved"); + assert.equal(result.previousAccountId, PROJECT_ACCOUNT, "must name what it replaced"); + assert.equal(result.accountId, ATTACKER_ACCOUNT); + assert.equal(auth.credsPath(), projectFile(cwd)); +}); + +test("approving the same destination twice writes nothing new", async (t) => { + const { auth } = await sandbox(t, { + global: makeCreds(), + project: makeCreds({ accountId: PROJECT_ACCOUNT }), + }); + auth.approveProjectCreds(); + + assert.equal(auth.approveProjectCreds().outcome, "already-approved"); + const stored = JSON.parse(readFileSync(auth.projectApprovalsPath(), "utf8")); + assert.equal(stored.approvals.length, 1, "approvals must not accumulate duplicates"); +}); + +test("revoking sends memory back to the global account", async (t) => { + const { auth, home, cwd } = await sandbox(t, { + global: makeCreds(), + project: makeCreds({ accountId: PROJECT_ACCOUNT }), + }); + auth.approveProjectCreds(); + assert.equal(auth.credsPath(), projectFile(cwd)); + + const revoked = auth.revokeProjectCredsApproval(); + + assert.equal(revoked.outcome, "revoked"); + assert.equal(auth.credsPath(), globalFile(home)); + assert.equal(auth.revokeProjectCredsApproval().outcome, "none", "revoking twice is a no-op"); +}); + +/* --------------------------------------------------------------------- * + * The escape hatch, and the cases that must stay quiet. + * --------------------------------------------------------------------- */ + +test("MEMWAL_CREDS_DIR still overrides both files, approved or not", async (t) => { + const { auth, cwd } = await sandbox(t, { + global: makeCreds(), + project: makeCreds({ accountId: PROJECT_ACCOUNT }), + }); + auth.approveProjectCreds(); + assert.equal(auth.credsPath(), projectFile(cwd), "precondition: the project file is in use"); + + const override = realpathSync(mkdtempSync(join(tmpdir(), "memwal-approve-override-"))); + t.after(() => { + delete process.env.MEMWAL_CREDS_DIR; + rmSync(override, { recursive: true, force: true }); + }); + process.env.MEMWAL_CREDS_DIR = override; + + assert.equal(auth.credsPath(), join(override, "credentials.json")); + assert.equal(auth.resolveCreds().source, "override"); + assert.equal(auth.formatProjectCredsNotice(), null, "an override has nothing to warn about"); + assert.equal( + auth.approveProjectCreds().outcome, + "overridden", + "there is nothing to approve while the override decides", + ); +}); + +test("a malformed project credentials file falls back to the global one", async (t) => { + const { auth, home, cwd } = await sandbox(t, { global: makeCreds() }); + mkdirSync(join(cwd, ".memwal"), { recursive: true }); + writeFileSync(projectFile(cwd), "{ not json"); + + assert.equal(auth.credsPath(), globalFile(home)); + assert.equal(auth.loadCreds()?.accountId, GLOBAL_ACCOUNT, "a broken repo file is not a logout"); + assert.equal(auth.resolveCreds().project?.decision, "unreadable"); + assert.match(auth.formatProjectCredsNotice(), /not a valid/); + assert.equal(auth.approveProjectCreds().outcome, "unreadable"); +}); + +test("a machine with no project file behaves exactly as it always did", async (t) => { + const { auth, home } = await sandbox(t, { global: makeCreds() }); + + assert.equal(auth.credsPath(), globalFile(home)); + assert.equal(auth.loadCreds()?.accountId, GLOBAL_ACCOUNT); + assert.equal(auth.resolveCreds().source, "global"); + assert.equal(auth.formatProjectCredsNotice(), null, "nothing to report, so nothing is said"); + assert.equal(auth.approveProjectCreds().outcome, "none"); + assert.equal(existsSync(auth.projectApprovalsPath()), false, "no file is created for nothing"); +}); + +test("an unapproved project file does not leave the user signed out", async (t) => { + // Falling back must not be a fail-closed: a user with only a repo file and + // no global one still gets the normal signed-out sign-in path, not an error. + const { auth, home } = await sandbox(t, { + project: makeCreds({ accountId: PROJECT_ACCOUNT }), + }); + + assert.equal(auth.credsPath(), globalFile(home)); + assert.equal(auth.loadCreds(), null); + assert.ok(auth.formatProjectCredsNotice(), "and the ignored file is still explained"); +}); diff --git a/packages/mcp/test/streamable-transport.test.mjs b/packages/mcp/test/streamable-transport.test.mjs new file mode 100644 index 000000000..03838106e --- /dev/null +++ b/packages/mcp/test/streamable-transport.test.mjs @@ -0,0 +1,53 @@ +import assert from "node:assert/strict"; +import test from "node:test"; + +import { resolveTransport, streamableUrl } from "../dist/streamable.js"; + +/** + * Transport selection and endpoint derivation only — the parts that hold + * without a socket. The session itself needs a live relayer, so it is covered + * by the live suite rather than here. + */ + +test("the default transport stays SSE", () => { + // Every released bridge dials SSE. Streamable HTTP is opt-in until it has + // production mileage, so an unset variable must not move users onto it. + assert.equal(resolveTransport(undefined), "sse"); + assert.equal(resolveTransport(""), "sse"); + assert.equal(resolveTransport(" "), "sse"); +}); + +test("http is selected by any of its spellings", () => { + for (const value of ["http", "streamable", "streamable-http"]) { + assert.equal(resolveTransport(value), "http", value); + } +}); + +test("selection ignores case and surrounding whitespace", () => { + assert.equal(resolveTransport(" HTTP "), "http"); + assert.equal(resolveTransport("SSE"), "sse"); +}); + +test("an unrecognised value falls back instead of throwing", () => { + // A typo in a user's MCP config must not stop their memory from working. + assert.equal(resolveTransport("htpp"), "sse"); + assert.equal(resolveTransport("websocket"), "sse"); +}); + +test("the streamable endpoint sits on the same base as the SSE pair", () => { + assert.equal( + streamableUrl("https://relayer.memory.walrus.xyz"), + "https://relayer.memory.walrus.xyz/api/mcp" + ); +}); + +test("a trailing slash on the relayer URL does not double up", () => { + // Users paste URLs with and without it; a `//api/mcp` path 404s. + assert.equal(streamableUrl("https://relayer.example/"), "https://relayer.example/api/mcp"); + assert.equal(streamableUrl("https://relayer.example///"), "https://relayer.example/api/mcp"); +}); + +test("a relayer on a port or subpath keeps it", () => { + assert.equal(streamableUrl("http://127.0.0.1:8000"), "http://127.0.0.1:8000/api/mcp"); + assert.equal(streamableUrl("https://host/base"), "https://host/base/api/mcp"); +}); diff --git a/packages/mcp/test/tool-definitions.test.mjs b/packages/mcp/test/tool-definitions.test.mjs index 6bbdbdb6c..57f3600a3 100644 --- a/packages/mcp/test/tool-definitions.test.mjs +++ b/packages/mcp/test/tool-definitions.test.mjs @@ -8,7 +8,10 @@ import assert from "node:assert/strict"; import { TOOL_DEFINITIONS, SIGNED_OUT_TOOL_DEFINITIONS, + ALL_TOOL_DEFINITIONS, + BASELINE_RELAYER_TOOLS, } from "../dist/auth-required.js"; +import { PROACTIVE_INSTRUCTIONS } from "../dist/instructions.js"; function desc(list, name) { const tool = list.find((t) => t.name === name); @@ -65,3 +68,127 @@ test("memwal_recall is advertised as a read-only search", () => { destructiveHint: false, }); }); + +/** + * The bridge carries its own copy of the tool list for the cold-start window, + * and nothing compares it to the sidecar's — the tests above only pin it + * against literals. That is how the two drifted: the sidecar grew batch + * settling and a pending-result contract, the bridge's copy did not, and the + * mismatch is invisible until a real session hits it. + * + * These pin the parts an agent acts on, so the next divergence fails here + * instead of in someone's first save of the session. + */ + +test("memwal_remember_status accepts a whole batch", () => { + // Not advertised at cold start (see the baseline tests below) — the shape + // is still pinned here, so it stays reviewed while it waits for the prod + // release that lets it into the cold-start list. + for (const list of [ALL_TOOL_DEFINITIONS]) { + const tool = list.find((t) => t.name === "memwal_remember_status"); + assert.ok(tool, "missing memwal_remember_status"); + + // `memwal_remember_bulk` hands back a pending body telling the agent to + // call this with job_ids. Under additionalProperties:false an absent + // property makes that instruction unfollowable. + const ids = tool.inputSchema.properties.job_ids; + assert.ok(ids, "job_ids must be advertised, or a batch cannot be settled"); + assert.equal(ids.type, "array"); + assert.equal(ids.items.type, "string"); + // Matches MAX_BULK_ITEMS, so a full batch settles in one call. + assert.equal(ids.maxItems, 20); + + // Requiring job_id would reject the batch form outright. + assert.ok( + !(tool.inputSchema.required ?? []).includes("job_id"), + "job_id must not be required — the batch form passes job_ids instead", + ); + } +}); + +test("memwal_remember_status waitMs bound matches the sidecar's ceiling", () => { + // The sidecar validates waitMs with zod and rejects anything above its own + // cap. Advertising a larger maximum invites the agent to send a value that + // comes straight back as an MCP validation error — observed live at 60000 + // once the sidecar lowered its ceiling to 45000. + for (const list of [ALL_TOOL_DEFINITIONS]) { + const tool = list.find((t) => t.name === "memwal_remember_status"); + assert.equal(tool.inputSchema.properties.waitMs.maximum, 45000); + } +}); + +test("cold-start write tools warn that a result may not be saved yet", () => { + // Both write tools return at accept now. An agent that was never told a + // pending result is normal reports it to the user as stored. + for (const name of ["memwal_remember", "memwal_remember_bulk"]) { + const d = desc(TOOL_DEFINITIONS, name); + assert.match(d, /NOT (yet )?(saved|stored)/i, `${name} omits the pending warning`); + assert.match(d, /settle it|settle them/i, `${name} does not say to settle the job`); + } +}); + +/** + * GH #928. The bridge ships on npm and updates itself; a relayer ships per + * environment and does not, so 0.0.14-dev.0 dialled prod and staging still on + * 0.0.13. Its cold-start list named `memwal_remember_status`, which neither + * serves, and the description told the agent to go call it — one live run + * spent 90.67s on a tool that does not exist there before erroring. + * + * The cold-start list is served before any relayer capability is known, so it + * has to be a floor: only tools the oldest supported relayer serves, and no + * description pointing at anything outside it. + */ + +test("cold-start lists advertise nothing beyond the baseline relayer", () => { + for (const list of [TOOL_DEFINITIONS, SIGNED_OUT_TOOL_DEFINITIONS]) { + const names = list.map((t) => t.name); + const beyond = names.filter( + (n) => !BASELINE_RELAYER_TOOLS.has(n) && n !== "memwal_login", + ); + assert.deepEqual( + beyond, + [], + `cold start advertises tools the oldest supported relayer cannot serve: ${beyond}`, + ); + for (const baseline of BASELINE_RELAYER_TOOLS) { + assert.ok(names.includes(baseline), `cold start omits ${baseline}`); + } + } +}); + +test("no cold-start description names a tool cold start does not advertise", () => { + // The tool list and the prose have to agree. A description is an + // instruction the agent follows, so naming an unadvertised tool is the + // same defect as listing it — it just fails one step later. + for (const list of [TOOL_DEFINITIONS, SIGNED_OUT_TOOL_DEFINITIONS]) { + const advertised = new Set(list.map((t) => t.name)); + for (const tool of list) { + const named = tool.description.match(/memwal_[a-z_]+/g) ?? []; + const dangling = [...new Set(named)].filter((n) => !advertised.has(n)); + assert.deepEqual( + dangling, + [], + `${tool.name}'s description sends the agent to unadvertised tools: ${dangling}`, + ); + } + } +}); + +test("initialize instructions do not name a tool cold start does not advertise", () => { + // Injected at `initialize`, before any `tools/list`. Naming + // `memwal_remember_status` here is the same #928 defect as listing it: + // the agent follows the sentence during the cold-start window, the + // local refusal used to be skipped while `upstreamToolNames` was empty, + // and the call sat in `inFlight` until the orphan deadline. + const advertised = new Set(TOOL_DEFINITIONS.map((t) => t.name)); + advertised.add("memwal_login"); + advertised.add("memwal_logout"); + const named = PROACTIVE_INSTRUCTIONS.match(/memwal_[a-z_]+/g) ?? []; + const dangling = [...new Set(named)].filter((n) => !advertised.has(n)); + assert.deepEqual( + dangling, + [], + `initialize instructions send the agent to unadvertised tools: ${dangling}`, + ); + assert.doesNotMatch(PROACTIVE_INSTRUCTIONS, /memwal_remember_status/); +}); diff --git a/packages/mcp/test/tool-not-served.test.mjs b/packages/mcp/test/tool-not-served.test.mjs new file mode 100644 index 000000000..7412ceb2a --- /dev/null +++ b/packages/mcp/test/tool-not-served.test.mjs @@ -0,0 +1,641 @@ +/** + * Regression test for GH #928 — a `tools/call` for a tool the connected relayer + * does not serve must be refused locally and immediately. + * + * What happened live: the bridge ships on npm and updates itself, a relayer + * ships per environment and does not, so 0.0.14-dev.0 talked to prod and + * staging still on 0.0.13. Its cold-start list advertised + * `memwal_remember_status`, and the pending-write wording told the agent to go + * call it. Neither relayer registers that tool, so the call was forwarded into + * a session that would never answer it and sat in `inFlight` until the orphan + * sweeper's deadline: one run spent 90.67s before erroring. + * + * The relayer's own `tools/list` is the authority on what it serves. Once the + * bridge has seen it, a call for anything outside that set (plus the two tools + * the bridge serves itself) is answered here — a stale tool list is not a + * transport fault, and the agent can only act on it if it is told now. + */ +import { test } from "node:test"; +import assert from "node:assert/strict"; +import http from "node:http"; +import { spawn } from "node:child_process"; +import { mkdtempSync, mkdirSync, writeFileSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join, dirname, resolve } from "node:path"; +import { fileURLToPath } from "node:url"; + +import { unknownToolText } from "../dist/bridge.js"; + +const __dirname = dirname(fileURLToPath(import.meta.url)); +const BIN = resolve(__dirname, "../dist/bin/memwal-mcp.js"); +const EXPECTED_BEARER = "a".repeat(64); +const EXPECTED_ACCOUNT_ID = "0x" + "3".repeat(64); + +/** What a 0.0.13 relayer registers: no `memwal_remember_status`. */ +const OLD_RELAYER_TOOLS = [ + "memwal_remember", + "memwal_remember_bulk", + "memwal_recall", + "memwal_analyze", + "memwal_restore", + "memwal_health", +]; + +/** The tool the newer bridge knows about and this relayer has never heard of. */ +const MISSING_TOOL = "memwal_remember_status"; + +test("the refusal tells the agent what exists and that nothing ran", () => { + const text = unknownToolText(MISSING_TOOL, [...OLD_RELAYER_TOOLS, "memwal_login"].sort()); + + // Which tool was refused, or the agent cannot tell which of several calls + // this answers. + assert.match(text, new RegExp(MISSING_TOOL)); + // What to do instead: the list it holds is stale, and here is the real one. + assert.match(text, /tools\/list/); + for (const name of OLD_RELAYER_TOOLS) assert.match(text, new RegExp(name)); + // A write tool that errors is ambiguous about whether the write happened. + // Say it plainly: an agent that guesses here tells the user a fact is + // saved when it never left the process. + assert.match(text, /nothing was saved/i); +}); + +function hasBridgeAuth(req) { + return ( + req.headers.authorization === `Bearer ${EXPECTED_BEARER}` && + req.headers["x-memwal-account-id"] === EXPECTED_ACCOUNT_ID + ); +} + +/** Relayer on the OLD tool set. It answers what it knows and stays silent on + * anything else — which is what a call for an unregistered tool looks like from + * the bridge's side, and what made the old behaviour a 90s wait. */ +function startOldRelayer() { + const sessions = new Map(); + const callsSeen = []; + const server = http.createServer((req, res) => { + const url = new URL(req.url, "http://127.0.0.1"); + if (req.method === "GET" && url.pathname === "/version") { + res.writeHead(200, { "content-type": "application/json" }); + res.end( + JSON.stringify({ + apiVersion: "1.0.0", + relayerVersion: "0.0.13", + minSupportedSdk: { mcp: "0.0.1" }, + }), + ); + return; + } + if (req.method === "GET" && url.pathname === "/api/mcp/sse") { + if (!hasBridgeAuth(req)) { + res.writeHead(401); + res.end(); + return; + } + const sessionId = `session-${sessions.size + 1}`; + res.writeHead(200, { + "content-type": "text/event-stream", + "cache-control": "no-cache", + connection: "keep-alive", + }); + res.write(`event: endpoint\ndata: /api/mcp/messages?sessionId=${sessionId}\n\n`); + sessions.set(sessionId, { res }); + const hb = setInterval(() => { + if (res.writableEnded) { + clearInterval(hb); + return; + } + res.write(":\n\n"); + }, 200); + hb.unref?.(); + res.on("close", () => clearInterval(hb)); + return; + } + if (req.method === "POST" && url.pathname === "/api/mcp/messages") { + if (!hasBridgeAuth(req)) { + res.writeHead(401); + res.end(); + return; + } + const session = sessions.get(url.searchParams.get("sessionId")); + let body = ""; + req.on("data", (c) => (body += c)); + req.on("end", () => { + if (!session) { + res.writeHead(404); + res.end(); + return; + } + res.writeHead(202); + res.end(); + let msg; + try { + msg = JSON.parse(body); + } catch { + return; + } + const reply = (result) => + session.res.write( + `event: message\ndata: ${JSON.stringify({ jsonrpc: "2.0", id: msg.id, result })}\n\n`, + ); + if (msg.method === "initialize") { + reply({ + protocolVersion: "2024-11-05", + capabilities: { tools: { listChanged: true } }, + serverInfo: { name: "memwal-upstream", version: "0.0.13" }, + }); + return; + } + if (msg.method === "tools/list") { + reply({ + tools: OLD_RELAYER_TOOLS.map((name) => ({ + name, + description: `upstream ${name}`, + inputSchema: { type: "object" }, + })), + }); + return; + } + if (msg.method === "tools/call") { + callsSeen.push(msg.params?.name); + if (msg.params?.name === "memwal_health") { + reply({ + content: [{ type: "text", text: "status=ok version=0.0.13" }], + isError: false, + }); + } + // Anything else: silence, exactly as an unregistered tool + // would behave if the call ever got this far. + return; + } + }); + return; + } + res.writeHead(404); + res.end(); + }); + return new Promise((res) => { + server.listen(0, "127.0.0.1", () => { + const { port } = server.address(); + res({ server, base: `http://127.0.0.1:${port}`, callsSeen }); + }); + }); +} + +function makeCreds(relayerUrl) { + return { + delegatePrivateKey: EXPECTED_BEARER, + delegatePublicKeyHex: "b".repeat(64), + delegateAddress: "0x" + "1".repeat(64), + walletAddress: "0x" + "2".repeat(64), + accountId: EXPECTED_ACCOUNT_ID, + packageId: "0x" + "4".repeat(64), + relayerUrl, + label: "Tool-not-served Test", + createdAt: new Date(0).toISOString(), + version: 1, + }; +} + +test("a tool the relayer does not serve is refused locally, not waited out", async (t) => { + const mock = await startOldRelayer(); + const home = mkdtempSync(join(tmpdir(), "memwal-toolskew-test-")); + const credsPath = join(home, ".memwal", "credentials.json"); + mkdirSync(dirname(credsPath), { recursive: true }); + writeFileSync(credsPath, JSON.stringify(makeCreds(mock.base)), { mode: 0o600 }); + + const child = spawn(process.execPath, [BIN, "--relayer", mock.base, "--web-url", mock.base], { + env: { ...process.env, HOME: home, USERPROFILE: home }, + stdio: ["pipe", "pipe", "pipe"], + }); + + const received = []; + const listeners = new Set(); + let buf = ""; + child.stdout.on("data", (d) => { + buf += d.toString(); + let nl; + while ((nl = buf.indexOf("\n")) >= 0) { + const line = buf.slice(0, nl); + buf = buf.slice(nl + 1); + if (!line.trim()) continue; + let msg; + try { + msg = JSON.parse(line); + } catch { + continue; + } + received.push({ msg, at: Date.now() }); + for (const l of [...listeners]) l(msg); + } + }); + let stderrBuf = ""; + child.stderr.on("data", (d) => (stderrBuf += d.toString())); + + const send = (obj) => child.stdin.write(JSON.stringify(obj) + "\n"); + const waitFor = (pred, ms = 15000) => { + const hit = received.find((r) => pred(r.msg)); + if (hit) return Promise.resolve(hit.msg); + return new Promise((res, rej) => { + const timer = setTimeout(() => { + listeners.delete(l); + rej( + new Error( + `timed out waiting for message\n--- stderr ---\n${stderrBuf}\n--- received ---\n${received.map((r) => JSON.stringify(r.msg)).join("\n")}`, + ), + ); + }, ms); + const l = (m) => { + if (pred(m)) { + clearTimeout(timer); + listeners.delete(l); + res(m); + } + }; + listeners.add(l); + }); + }; + + t.after(() => { + child.kill("SIGKILL"); + mock.server.close(); + rmSync(home, { recursive: true, force: true }); + }); + + send({ + jsonrpc: "2.0", + id: 1, + method: "initialize", + params: { + protocolVersion: "2025-06-18", + capabilities: {}, + clientInfo: { name: "toolskew-test", version: "1.0.0" }, + }, + }); + await waitFor((m) => m.id === 1 && m.result, 10_000); + + // The cold-start list is a floor, so it must not name the tool this relayer + // lacks even before anything upstream is known. + send({ jsonrpc: "2.0", id: 2, method: "tools/list", params: {} }); + const cold = await waitFor((m) => m.id === 2 && m.result, 10_000); + assert.ok( + !cold.result.tools.some((tool) => tool.name === MISSING_TOOL), + `cold-start tools/list advertises ${MISSING_TOOL}, which no released relayer serves`, + ); + + // Cold-start window: initialize instructions used to name this tool, and + // the gate used to fail open while `upstreamToolNames` was empty (it is + // only filled from a forwarded `tools/list`; the cold-start list is + // answered locally). Call it here, before any forwarded list. + const coldCallStarted = Date.now(); + send({ + jsonrpc: "2.0", + id: "cold-missing", + method: "tools/call", + params: { name: MISSING_TOOL, arguments: { job_id: "job-1" } }, + }); + const coldRefusal = await waitFor((m) => m.id === "cold-missing", 10_000); + const coldElapsed = Date.now() - coldCallStarted; + assert.equal( + coldRefusal.result?.isError, + true, + `expected a local error during cold start, got ${JSON.stringify(coldRefusal)}`, + ); + assert.match(coldRefusal.result.content[0].text, new RegExp(MISSING_TOOL)); + assert.match(coldRefusal.result.content[0].text, /tools\/list/); + assert.match(coldRefusal.result.content[0].text, /nothing was saved/i); + assert.ok( + coldElapsed < 5_000, + `cold-start refusal took ${coldElapsed}ms — expected it answered locally, not at the orphan deadline`, + ); + assert.ok( + !mock.callsSeen.includes(MISSING_TOOL), + `bridge forwarded ${MISSING_TOOL} during cold start: ${mock.callsSeen}`, + ); + + // Re-list once connected so the bridge learns what this relayer actually + // registers. That reply is the authority the refusal below is based on. + await waitFor((m) => m.method === "notifications/tools/list_changed", 10_000); + send({ jsonrpc: "2.0", id: 3, method: "tools/list", params: {} }); + const upstream = await waitFor((m) => m.id === 3 && m.result, 10_000); + const upstreamNames = upstream.result.tools.map((tool) => tool.name); + assert.ok( + upstreamNames.includes("memwal_recall"), + `expected the relayer's own tool list, got ${upstreamNames}`, + ); + assert.ok(!upstreamNames.includes(MISSING_TOOL)); + + // The call the live run made. Before the fix it was forwarded and left to + // the orphan sweeper: 60s tool ceiling + 30s headroom. + const startedAt = Date.now(); + send({ + jsonrpc: "2.0", + id: 4, + method: "tools/call", + params: { name: MISSING_TOOL, arguments: { job_id: "job-1" } }, + }); + const refusal = await waitFor((m) => m.id === 4, 10_000); + const elapsed = Date.now() - startedAt; + + assert.equal(refusal.result?.isError, true, `expected an error result, got ${JSON.stringify(refusal)}`); + const text = refusal.result.content[0].text; + assert.match(text, new RegExp(MISSING_TOOL)); + // The agent's way out has to be in the message: which tools exist, and that + // its own list is what is wrong. + assert.match(text, /tools\/list/); + assert.match(text, /memwal_recall/); + // And it must be unambiguous that the write did not happen, or the agent + // reports a fact as saved on the strength of an error. + assert.match(text, /nothing was saved/i); + assert.ok( + elapsed < 5_000, + `refusal took ${elapsed}ms — expected it answered locally, not at the orphan deadline`, + ); + + // Proof it never left the process: the relayer saw the health call we made + // nothing of, and never the missing tool. + assert.ok( + !mock.callsSeen.includes(MISSING_TOOL), + `bridge forwarded ${MISSING_TOOL} upstream: ${mock.callsSeen}`, + ); + + // A tool the relayer DOES serve still goes through, so the gate is a filter + // and not a wall. + send({ + jsonrpc: "2.0", + id: 5, + method: "tools/call", + params: { name: "memwal_health", arguments: {} }, + }); + const health = await waitFor((m) => m.id === 5 && m.result, 10_000); + assert.notEqual(health.result?.isError, true); + assert.match(JSON.stringify(health.result), /status=ok/); +}); + +/** First session advertises the newer tool; later sessions look like 0.0.13. + * Reproducing a login / SSE reconnect that swaps the relayer under the bridge. */ +const NEW_RELAYER_TOOLS = [...OLD_RELAYER_TOOLS, MISSING_TOOL]; + +function startSkewRelayer() { + const sessions = []; + const callsSeen = []; + const server = http.createServer((req, res) => { + const url = new URL(req.url, "http://127.0.0.1"); + if (req.method === "GET" && url.pathname === "/version") { + res.writeHead(200, { "content-type": "application/json" }); + res.end( + JSON.stringify({ + apiVersion: "1.0.0", + relayerVersion: "0.0.14", + minSupportedSdk: { mcp: "0.0.1" }, + }), + ); + return; + } + if (req.method === "GET" && url.pathname === "/api/mcp/sse") { + if (!hasBridgeAuth(req)) { + res.writeHead(401); + res.end(); + return; + } + const first = sessions.length === 0; + const sessionId = `session-${sessions.length + 1}`; + res.writeHead(200, { + "content-type": "text/event-stream", + "cache-control": "no-cache", + connection: "keep-alive", + }); + res.write(`event: endpoint\ndata: /api/mcp/messages?sessionId=${sessionId}\n\n`); + sessions.push({ id: sessionId, res, first }); + const hb = setInterval(() => { + if (res.writableEnded) { + clearInterval(hb); + return; + } + res.write(":\n\n"); + }, 200); + hb.unref?.(); + res.on("close", () => clearInterval(hb)); + return; + } + if (req.method === "POST" && url.pathname === "/api/mcp/messages") { + if (!hasBridgeAuth(req)) { + res.writeHead(401); + res.end(); + return; + } + const session = sessions.find((s) => s.id === url.searchParams.get("sessionId")); + let body = ""; + req.on("data", (c) => (body += c)); + req.on("end", () => { + if (!session) { + res.writeHead(404); + res.end(); + return; + } + res.writeHead(202); + res.end(); + let msg; + try { + msg = JSON.parse(body); + } catch { + return; + } + const tools = session.first ? NEW_RELAYER_TOOLS : OLD_RELAYER_TOOLS; + const reply = (result) => { + if (session.res.writableEnded) return; + session.res.write( + `event: message\ndata: ${JSON.stringify({ jsonrpc: "2.0", id: msg.id, result })}\n\n`, + ); + }; + if (msg.method === "initialize") { + reply({ + protocolVersion: "2024-11-05", + capabilities: { tools: { listChanged: true } }, + serverInfo: { name: "memwal-upstream", version: session.first ? "0.0.14" : "0.0.13" }, + }); + return; + } + if (msg.method === "tools/list") { + reply({ + tools: tools.map((name) => ({ + name, + description: `upstream ${name}`, + inputSchema: { type: "object" }, + })), + }); + return; + } + if (msg.method === "tools/call") { + callsSeen.push({ session: session.id, name: msg.params?.name }); + if (msg.params?.name === "memwal_health") { + reply({ + content: [{ type: "text", text: "status=ok version=skew" }], + isError: false, + }); + return; + } + if (msg.params?.name === MISSING_TOOL && session.first) { + reply({ + content: [{ type: "text", text: "status=done blob_id=from-first-session" }], + isError: false, + }); + } + // Later sessions: silence on the missing tool, same as 0.0.13. + } + }); + return; + } + res.writeHead(404); + res.end(); + }); + return new Promise((res) => { + server.listen(0, "127.0.0.1", () => { + const { port } = server.address(); + res({ + server, + base: `http://127.0.0.1:${port}`, + callsSeen, + sessions, + closeFirst() { + const first = sessions[0]; + if (first && !first.res.writableEnded) first.res.end(); + }, + }); + }); + }); +} + +test("reconnect forgets the previous relayer's tool set instead of stale-allowing", async (t) => { + const mock = await startSkewRelayer(); + const home = mkdtempSync(join(tmpdir(), "memwal-toolskew-reconnect-")); + const credsPath = join(home, ".memwal", "credentials.json"); + mkdirSync(dirname(credsPath), { recursive: true }); + writeFileSync(credsPath, JSON.stringify(makeCreds(mock.base)), { mode: 0o600 }); + + const child = spawn(process.execPath, [BIN, "--relayer", mock.base, "--web-url", mock.base], { + env: { ...process.env, HOME: home, USERPROFILE: home }, + stdio: ["pipe", "pipe", "pipe"], + }); + + const received = []; + const listeners = new Set(); + let buf = ""; + child.stdout.on("data", (d) => { + buf += d.toString(); + let nl; + while ((nl = buf.indexOf("\n")) >= 0) { + const line = buf.slice(0, nl); + buf = buf.slice(nl + 1); + if (!line.trim()) continue; + let msg; + try { + msg = JSON.parse(line); + } catch { + continue; + } + received.push({ msg, at: Date.now() }); + for (const l of [...listeners]) l(msg); + } + }); + let stderrBuf = ""; + child.stderr.on("data", (d) => (stderrBuf += d.toString())); + + const send = (obj) => child.stdin.write(JSON.stringify(obj) + "\n"); + const waitFor = (pred, ms = 15000) => { + const hit = received.find((r) => pred(r.msg)); + if (hit) return Promise.resolve(hit.msg); + return new Promise((res, rej) => { + const timer = setTimeout(() => { + listeners.delete(l); + rej( + new Error( + `timed out waiting for message\n--- stderr ---\n${stderrBuf}\n--- received ---\n${received.map((r) => JSON.stringify(r.msg)).join("\n")}`, + ), + ); + }, ms); + const l = (m) => { + if (pred(m)) { + clearTimeout(timer); + listeners.delete(l); + res(m); + } + }; + listeners.add(l); + }); + }; + + t.after(() => { + child.kill("SIGKILL"); + mock.server.close(); + rmSync(home, { recursive: true, force: true }); + }); + + send({ + jsonrpc: "2.0", + id: 1, + method: "initialize", + params: { + protocolVersion: "2025-06-18", + capabilities: {}, + clientInfo: { name: "toolskew-reconnect", version: "1.0.0" }, + }, + }); + await waitFor((m) => m.id === 1 && m.result, 10_000); + + await waitFor((m) => m.method === "notifications/tools/list_changed", 10_000); + send({ jsonrpc: "2.0", id: 2, method: "tools/list", params: {} }); + const firstList = await waitFor((m) => m.id === 2 && m.result, 10_000); + assert.ok( + firstList.result.tools.some((tool) => tool.name === MISSING_TOOL), + "first session must advertise the newer tool so the stale-allow path is live", + ); + + send({ + jsonrpc: "2.0", + id: 3, + method: "tools/call", + params: { name: MISSING_TOOL, arguments: { job_id: "job-1" } }, + }); + const firstCall = await waitFor((m) => m.id === 3 && m.result, 10_000); + assert.notEqual(firstCall.result?.isError, true); + assert.match(JSON.stringify(firstCall.result), /from-first-session/); + + const changedBefore = received.filter( + (r) => r.msg.method === "notifications/tools/list_changed", + ).length; + mock.closeFirst(); + await waitFor( + () => + mock.sessions.length >= 2 && + received.filter((r) => r.msg.method === "notifications/tools/list_changed").length > + changedBefore, + 10_000, + ); + + // Do NOT re-list. The previous allow-set still named the tool; forwarding + // it into session-2 (0.0.13, silent on this name) is the #928 hang. + const startedAt = Date.now(); + send({ + jsonrpc: "2.0", + id: 4, + method: "tools/call", + params: { name: MISSING_TOOL, arguments: { job_id: "job-2" } }, + }); + const refusal = await waitFor((m) => m.id === 4, 10_000); + const elapsed = Date.now() - startedAt; + assert.equal( + refusal.result?.isError, + true, + `expected a local refusal after reconnect, got ${JSON.stringify(refusal)}`, + ); + assert.match(refusal.result.content[0].text, new RegExp(MISSING_TOOL)); + assert.match(refusal.result.content[0].text, /nothing was saved/i); + assert.ok( + elapsed < 5_000, + `post-reconnect refusal took ${elapsed}ms — expected it answered locally`, + ); + assert.ok( + !mock.callsSeen.some((c) => c.session !== "session-1" && c.name === MISSING_TOOL), + `bridge forwarded ${MISSING_TOOL} to the new session: ${JSON.stringify(mock.callsSeen)}`, + ); +}); diff --git a/packages/mcp/test/trusted-launcher.test.mjs b/packages/mcp/test/trusted-launcher.test.mjs new file mode 100644 index 000000000..f94a3d246 --- /dev/null +++ b/packages/mcp/test/trusted-launcher.test.mjs @@ -0,0 +1,693 @@ +/** + * The plugin must never run an MCP server that came out of the user's project + * (WALM-640). + * + * The plugin used to start the server with `npx -y @mysten-incubation/memwal-mcp@`. + * npx resolves that name against the working directory the MCP client was started + * in — the project the user has open — so a project carrying an installed package + * of the same name that *claims the pinned version* won, offline. The version pin + * is no defence: the fake package just says it is that version. + * + * Every test here plants exactly that: a fake `@mysten-incubation/memwal-mcp` in a + * temp project's node_modules, with the real pinned version in its manifest and a + * bin that prints LOCAL_PACKAGE_EXECUTED, plus the `node_modules/.bin` shim npx + * would have found. The launcher must resolve and run the trusted absolute path + * instead — and when the trusted install cannot be produced, it must fail rather + * than fall back to the name. + */ +import { test } from "node:test"; +import assert from "node:assert/strict"; +import { spawnSync } from "node:child_process"; +import { + chmodSync, + lstatSync, + mkdirSync, + mkdtempSync, + readFileSync, + rmSync, + symlinkSync, + writeFileSync, +} from "node:fs"; +import { tmpdir } from "node:os"; +import { dirname, join, resolve } from "node:path"; +import { fileURLToPath } from "node:url"; + +import { + MCP_PACKAGE_NAME, + MCP_REGISTRY, + assertRuntimeRootLocation, + canonicalPath, + expectedEntryPath, + installArguments, + installDir, + installEnvironment, + npmSpawnPlan, + pinnedVersion, + prepareRuntimeRoot, + quoteWindowsArgument, + resolveInstalledEntry, + resolveNpm, + runtimeRoot, + verifyTrustedDirectory, +} from "../plugin/scripts/lib/mcp-launch.mjs"; + +const __dirname = dirname(fileURLToPath(import.meta.url)); +const PLUGIN_DIR = resolve(__dirname, "../plugin"); +const LAUNCHER = join(PLUGIN_DIR, "scripts", "launch_mcp.mjs"); +const PIN = pinnedVersion(); + +const LOCAL_MARKER = "LOCAL_PACKAGE_EXECUTED"; +const TRUSTED_MARKER = "TRUSTED_ENTRY_EXECUTED"; + +/** + * Write a package that looks exactly like the real one to any resolver: same name, + * same `bin` layout, and whatever version the caller wants it to claim. + */ +function plantPackage(nodeModulesParent, { version, marker }) { + const packageDir = join(nodeModulesParent, "node_modules", ...MCP_PACKAGE_NAME.split("/")); + const binDir = join(packageDir, "dist", "bin"); + mkdirSync(binDir, { recursive: true }); + writeFileSync( + join(packageDir, "package.json"), + JSON.stringify({ + name: MCP_PACKAGE_NAME, + version, + bin: { "memwal-mcp": "dist/bin/memwal-mcp.js" }, + }), + ); + const entry = join(binDir, "memwal-mcp.js"); + writeFileSync( + entry, + `console.log(${JSON.stringify(marker)});\n` + + `console.log("ARGV:" + process.argv.slice(2).join(" "));\n`, + ); + return entry; +} + +/** The PATH-relative shim npx would have reached for inside the project. */ +function plantBinShim(projectDir, marker) { + const binDir = join(projectDir, "node_modules", ".bin"); + mkdirSync(binDir, { recursive: true }); + const shim = join(binDir, "memwal-mcp"); + writeFileSync(shim, `#!/bin/sh\necho ${marker}\n`); + chmodSync(shim, 0o755); + return shim; +} + +/** + * A project a user might have open, carrying a fake package that claims the exact + * version the plugin pins. + */ +function makeHostileProject(t, { version = PIN } = {}) { + const dir = mkdtempSync(join(tmpdir(), "memwal-hostile-project-")); + t.after(() => rmSync(dir, { recursive: true, force: true })); + const entry = plantPackage(dir, { version, marker: LOCAL_MARKER }); + plantBinShim(dir, LOCAL_MARKER); + return { dir, entry }; +} + +/** A trusted runtime root with the pinned version already installed in it. */ +function makeTrustedRuntime(t, { version = PIN, populate = true } = {}) { + const root = mkdtempSync(join(tmpdir(), "memwal-runtime-")); + t.after(() => rmSync(root, { recursive: true, force: true })); + if (!populate) return { root, entry: null }; + const target = installDir(version, root); + mkdirSync(target, { recursive: true }); + const entry = plantPackage(target, { version, marker: TRUSTED_MARKER }); + return { root, entry }; +} + +function runLauncher(args, { cwd, env }) { + return spawnSync(process.execPath, [LAUNCHER, ...args], { + cwd, + encoding: "utf8", + env: { ...process.env, ...env }, + }); +} + +test("the planted package claims the exact pinned version", (t) => { + const project = makeHostileProject(t); + const manifest = JSON.parse( + readFileSync( + join(project.dir, "node_modules", ...MCP_PACKAGE_NAME.split("/"), "package.json"), + "utf8", + ), + ); + // If this drifts the rest of the file stops testing the reported attack. + assert.equal(manifest.version, PIN); + assert.equal(manifest.name, MCP_PACKAGE_NAME); +}); + +test("resolution ignores the project entirely and lands in the trusted directory", (t) => { + const project = makeHostileProject(t); + const trusted = makeTrustedRuntime(t); + + const previousCwd = process.cwd(); + t.after(() => process.chdir(previousCwd)); + process.chdir(project.dir); + + const resolved = resolveInstalledEntry(installDir(PIN, trusted.root), PIN); + assert.equal(resolved, trusted.entry); + assert.equal(resolved, expectedEntryPath({ version: PIN, root: trusted.root })); + assert.ok(resolved.startsWith(trusted.root), `${resolved} is not under ${trusted.root}`); + assert.ok(!resolved.includes(project.dir), `${resolved} points into the project`); +}); + +test("--print-entry from inside the hostile project prints the trusted path", (t) => { + const project = makeHostileProject(t); + const trusted = makeTrustedRuntime(t); + + const result = runLauncher(["--print-entry"], { + cwd: project.dir, + env: { MEMWAL_MCP_RUNTIME_DIR: trusted.root }, + }); + + assert.equal(result.status, 0, result.stderr); + const printed = result.stdout.trim(); + assert.equal(printed, trusted.entry); + assert.notEqual(printed, project.entry); + assert.ok(!printed.includes(project.dir), `${printed} points into the project`); +}); + +test("launching from the hostile project runs the trusted entry, not the local one", (t) => { + const project = makeHostileProject(t); + const trusted = makeTrustedRuntime(t); + + const result = runLauncher([], { + cwd: project.dir, + env: { MEMWAL_MCP_RUNTIME_DIR: trusted.root }, + }); + + assert.equal(result.status, 0, result.stderr); + assert.match(result.stdout, new RegExp(TRUSTED_MARKER)); + assert.doesNotMatch(result.stdout, new RegExp(LOCAL_MARKER)); +}); + +test("server flags and env are forwarded to the trusted entry unchanged", (t) => { + const project = makeHostileProject(t); + const trusted = makeTrustedRuntime(t); + + const result = runLauncher(["--dev", "--namespace", "work"], { + cwd: project.dir, + env: { MEMWAL_MCP_RUNTIME_DIR: trusted.root }, + }); + + assert.equal(result.status, 0, result.stderr); + assert.match(result.stdout, /ARGV:--dev --namespace work/); +}); + +test("an install under a different version is not accepted for the pin", (t) => { + const project = makeHostileProject(t); + // Trusted directory holds a stale build; the pinned directory is empty. + const trusted = makeTrustedRuntime(t, { version: "0.0.0-stale" }); + + assert.equal(resolveInstalledEntry(installDir(PIN, trusted.root), PIN), null); + assert.equal( + resolveInstalledEntry(installDir("0.0.0-stale", trusted.root), PIN), + null, + "a directory whose manifest disagrees with the pin must not be used", + ); + assert.ok(!String(trusted.entry).includes(project.dir)); +}); + +test("with no trusted install and a failing install, the launcher fails instead of falling back", (t) => { + const project = makeHostileProject(t); + const trusted = makeTrustedRuntime(t, { populate: false }); + + // A read-only trusted root makes the install fail the way an offline or broken + // toolchain would, without going near the network. (An empty PATH no longer does + // it: npm is resolved from process.execPath, which is the point of that change.) + let result; + chmodSync(trusted.root, 0o500); + try { + result = runLauncher([], { + cwd: project.dir, + env: { MEMWAL_MCP_RUNTIME_DIR: trusted.root }, + }); + } finally { + chmodSync(trusted.root, 0o700); + } + + assert.notEqual(result.status, 0, "launcher must not succeed without a trusted install"); + assert.doesNotMatch(result.stdout, new RegExp(LOCAL_MARKER)); + assert.doesNotMatch(result.stderr, new RegExp(LOCAL_MARKER)); + assert.match(result.stderr, /refusing to fall back/); +}); + +test("a relative MEMWAL_MCP_RUNTIME_DIR is rejected, not resolved against the project", (t) => { + const project = makeHostileProject(t); + + const result = runLauncher(["--print-entry"], { + cwd: project.dir, + env: { MEMWAL_MCP_RUNTIME_DIR: "node_modules" }, + }); + + assert.notEqual(result.status, 0); + assert.match(result.stderr, /absolute path/); + assert.doesNotMatch(result.stdout, new RegExp(LOCAL_MARKER)); +}); + +test("the default trusted root is ~/.memwal/runtime", (t) => { + const project = makeHostileProject(t); + const home = mkdtempSync(join(tmpdir(), "memwal-home-")); + t.after(() => rmSync(home, { recursive: true, force: true })); + + const target = join(home, ".memwal", "runtime", `memwal-mcp@${PIN}`); + mkdirSync(target, { recursive: true }); + const entry = plantPackage(target, { version: PIN, marker: TRUSTED_MARKER }); + + // USERPROFILE alongside HOME: os.homedir() reads USERPROFILE on Windows and + // would otherwise escape the sandbox into the real home. + const result = runLauncher(["--print-entry"], { + cwd: project.dir, + env: { HOME: home, USERPROFILE: home, MEMWAL_MCP_RUNTIME_DIR: "" }, + }); + + assert.equal(result.status, 0, result.stderr); + assert.equal(result.stdout.trim(), entry); +}); + +test("runtimeRoot honours an absolute override and otherwise sits under the home dir", () => { + const previous = process.env.MEMWAL_MCP_RUNTIME_DIR; + try { + const absolute = join(tmpdir(), "memwal-runtime-override"); + process.env.MEMWAL_MCP_RUNTIME_DIR = absolute; + assert.equal(runtimeRoot(), absolute); + + delete process.env.MEMWAL_MCP_RUNTIME_DIR; + assert.match(runtimeRoot(), /[\\/]\.memwal[\\/]runtime$/); + } finally { + if (previous === undefined) delete process.env.MEMWAL_MCP_RUNTIME_DIR; + else process.env.MEMWAL_MCP_RUNTIME_DIR = previous; + } +}); + +test("no plugin launch manifest resolves the server through npx", () => { + const manifests = [ + [".mcp.json", "${CLAUDE_PLUGIN_ROOT}"], + [".cursor-mcp.json", "${CURSOR_PLUGIN_ROOT}"], + [".codex-mcp.json", "${PLUGIN_ROOT}"], + ]; + for (const [name, root] of manifests) { + const server = JSON.parse(readFileSync(join(PLUGIN_DIR, name), "utf8")).mcpServers.memwal; + assert.equal(server.command, "node", `${name} must not launch through npx`); + assert.deepEqual(server.args, [`${root}/scripts/launch_mcp.mjs`], name); + } + + const installer = readFileSync( + join(PLUGIN_DIR, "scripts", "install_codex_hooks.mjs"), + "utf8", + ); + assert.doesNotMatch( + installer, + /command\s*=\s*\\?"npx/, + "the Codex fallback installer must register the launcher, not npx", + ); + assert.match(installer, /launch_mcp\.mjs/); +}); + +test("the pinned MCP version is the release this plugin is, or a prerelease of it", () => { + const manifest = JSON.parse(readFileSync(join(PLUGIN_DIR, "plugin.json"), "utf8")); + const pin = manifest.mcpPackageVersion; + + assert.equal(typeof pin, "string", "plugin.json must carry mcpPackageVersion"); + assert.match( + pin, + /^\d+\.\d+\.\d+(?:-[0-9A-Za-z.-]+)?$/, + `mcpPackageVersion is not a semver: ${pin}`, + ); + + // The field exists so `version` can name the release being prepared while the + // launcher installs something npm actually carries — installing a version npm + // does not have fails the launch outright. That only works while the two stay + // tied. Bump `version` for the next release and leave `mcpPackageVersion` on a + // prerelease of the previous one and nothing breaks loudly: every install + // quietly keeps serving the older client, which is how 0.0.14-dev.0 outlived + // the relayer change that made its own instructions wrong. So: either the + // release itself, or a prerelease of it. Nothing else. + assert.ok( + pin === manifest.version || pin.startsWith(`${manifest.version}-`), + `mcpPackageVersion ${pin} is neither ${manifest.version} nor a prerelease of it`, + ); + + assert.equal( + pinnedVersion(PLUGIN_DIR), + pin, + "the launcher must install exactly what plugin.json pins", + ); +}); + +test("pinnedVersion prefers mcpPackageVersion, falls back to version, and refuses neither", (t) => { + const dir = mkdtempSync(join(tmpdir(), "memwal-plugin-manifest-")); + t.after(() => rmSync(dir, { recursive: true, force: true })); + const write = (manifest) => + writeFileSync(join(dir, "plugin.json"), JSON.stringify(manifest)); + + write({ version: "9.9.9", mcpPackageVersion: "9.9.9-dev.3" }); + assert.equal(pinnedVersion(dir), "9.9.9-dev.3"); + + write({ version: "9.9.9" }); + assert.equal(pinnedVersion(dir), "9.9.9", "an absent mcpPackageVersion falls back to version"); + + write({ version: " " }); + assert.throws(() => pinnedVersion(dir), /no usable/, "a blank version is not a pin"); +}); + +/* ------------------------------------------------------------------------- * + * Review follow-ups (WALM-640): an absolute runtime root is not a trusted + * runtime root, and the install step must not execute what it has not verified. + * ------------------------------------------------------------------------- */ + +const POSIX = process.platform !== "win32"; + +/** A project that a repository could really ship: markers, config, planted runtime. */ +function makeRepoProject(t) { + const dir = mkdtempSync(join(tmpdir(), "memwal-repo-project-")); + t.after(() => rmSync(dir, { recursive: true, force: true })); + mkdirSync(join(dir, ".git"), { recursive: true }); + writeFileSync(join(dir, "package.json"), JSON.stringify({ name: "victim", version: "1.0.0" })); + return dir; +} + +/** The whole WALM-640 tree, committed inside the repository under `name`. */ +function plantRuntimeRoot(projectDir, name = ".memwal-runtime") { + const root = join(projectDir, name); + const target = installDir(PIN, root); + mkdirSync(target, { recursive: true }); + plantPackage(target, { version: PIN, marker: LOCAL_MARKER }); + return root; +} + +test("an absolute MEMWAL_MCP_RUNTIME_DIR inside the project is refused", (t) => { + const project = makeHostileProject(t); + // What a client writes into ${workspaceFolder}/.cursor/mcp.json expands to: an + // absolute path, which the old check accepted, pointing straight back into the + // repository that supplied it. + const plantedRoot = plantRuntimeRoot(project.dir); + + const result = runLauncher(["--print-entry"], { + cwd: project.dir, + env: { MEMWAL_MCP_RUNTIME_DIR: plantedRoot }, + }); + + assert.notEqual(result.status, 0, "an in-project runtime root must not be usable"); + assert.match(result.stderr, /inside the project tree/); + assert.match(result.stderr, /refusing to fall back/); + assert.doesNotMatch(result.stdout, new RegExp(LOCAL_MARKER)); + assert.doesNotMatch(result.stdout, new RegExp(plantedRoot.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"))); +}); + +test("the planted in-project tree is otherwise a complete, usable install", (t) => { + // Without this the test above could pass for the wrong reason (a tree that would + // not have resolved anyway). It is refused because of where it is, not what it is. + const project = makeHostileProject(t); + const plantedRoot = plantRuntimeRoot(project.dir); + assert.ok( + resolveInstalledEntry(installDir(PIN, plantedRoot), PIN), + "the planted tree must be a resolvable install, so the refusal is about location", + ); +}); + +test("a runtime root inside the repository but above the cwd is refused too", (t) => { + const repo = makeRepoProject(t); + const plantedRoot = plantRuntimeRoot(repo); + const nested = join(repo, "packages", "app"); + mkdirSync(nested, { recursive: true }); + + const result = runLauncher(["--print-entry"], { + cwd: nested, + env: { MEMWAL_MCP_RUNTIME_DIR: plantedRoot }, + }); + + assert.notEqual(result.status, 0); + assert.match(result.stderr, /inside the project tree/); + assert.doesNotMatch(result.stdout, new RegExp(LOCAL_MARKER)); +}); + +test("a runtime root that only looks external is refused once symlinks resolve", (t) => { + if (!POSIX) return; + const project = makeHostileProject(t); + const plantedRoot = plantRuntimeRoot(project.dir); + const outside = mkdtempSync(join(tmpdir(), "memwal-link-")); + t.after(() => rmSync(outside, { recursive: true, force: true })); + const link = join(outside, "runtime"); + symlinkSync(plantedRoot, link); + + const result = runLauncher(["--print-entry"], { + cwd: project.dir, + env: { MEMWAL_MCP_RUNTIME_DIR: link }, + }); + + assert.notEqual(result.status, 0, "a symlink into the project is still the project"); + assert.match(result.stderr, /inside the project tree/); + assert.doesNotMatch(result.stdout, new RegExp(LOCAL_MARKER)); +}); + +test("assertRuntimeRootLocation accepts a root outside the project and rejects one inside", (t) => { + const repo = makeRepoProject(t); + const outside = mkdtempSync(join(tmpdir(), "memwal-outside-")); + t.after(() => rmSync(outside, { recursive: true, force: true })); + + assert.equal(assertRuntimeRootLocation(outside, { cwd: repo }), resolve(outside)); + assert.throws( + () => assertRuntimeRootLocation(join(repo, ".memwal-runtime"), { cwd: repo }), + /inside the project tree/, + ); + assert.throws( + () => assertRuntimeRootLocation(repo, { cwd: repo }), + /inside the project tree/, + ); + // The reverse containment is just as wrong: the project must not sit inside the + // directory whose contents we are about to execute. + assert.throws( + () => assertRuntimeRootLocation(dirname(repo), { cwd: repo }), + /inside the project tree/, + ); +}); + +test("a group-writable runtime root is refused", (t) => { + if (!POSIX) return; + const project = makeHostileProject(t); + const trusted = makeTrustedRuntime(t); + chmodSync(trusted.root, 0o777); + + const result = runLauncher(["--print-entry"], { + cwd: project.dir, + env: { MEMWAL_MCP_RUNTIME_DIR: trusted.root }, + }); + chmodSync(trusted.root, 0o700); + + assert.notEqual(result.status, 0, "a root anyone can write is not a trusted root"); + assert.match(result.stderr, /group- or world-writable/); + assert.match(result.stderr, /refusing to fall back/); + assert.doesNotMatch(result.stdout, new RegExp(LOCAL_MARKER)); +}); + +test("verifyTrustedDirectory rejects loose modes, symlinks and foreign owners", (t) => { + if (!POSIX) return; + const root = mkdtempSync(join(tmpdir(), "memwal-verify-")); + t.after(() => rmSync(root, { recursive: true, force: true })); + + const good = join(root, "good"); + mkdirSync(good, { mode: 0o700 }); + assert.equal(verifyTrustedDirectory(good), true); + assert.equal(verifyTrustedDirectory(join(root, "absent")), false); + + for (const mode of [0o770, 0o707, 0o777, 0o755 | 0o020]) { + chmodSync(good, mode); + assert.throws(() => verifyTrustedDirectory(good), /group- or world-writable/, `mode ${mode.toString(8)}`); + } + chmodSync(good, 0o700); + + const link = join(root, "link"); + symlinkSync(good, link); + assert.throws(() => verifyTrustedDirectory(link), /not a directory/); + + const file = join(root, "file"); + writeFileSync(file, ""); + assert.throws(() => verifyTrustedDirectory(file), /not a directory/); + + // A directory owned by somebody else. Skipped when the suite runs as that + // somebody else (root in a container), where the check cannot fire. + const foreign = "/usr"; + if (process.getuid() !== 0 && lstatSync(foreign).uid !== process.getuid()) { + assert.throws(() => verifyTrustedDirectory(foreign), /owned by uid/); + } +}); + +test("the runtime root and ~/.memwal are created 0700, not at the umask default", (t) => { + if (!POSIX) return; + const home = mkdtempSync(join(tmpdir(), "memwal-home-mode-")); + t.after(() => rmSync(home, { recursive: true, force: true })); + + const previousHome = process.env.HOME; + t.after(() => { + if (previousHome === undefined) delete process.env.HOME; + else process.env.HOME = previousHome; + }); + process.env.HOME = home; + + const root = prepareRuntimeRoot(join(home, ".memwal", "runtime")); + assert.equal(lstatSync(root).mode & 0o7777, 0o700); + // The launcher can run before the first login, so it is this call that creates + // ~/.memwal — the directory credentials.json later lands in. + assert.equal(lstatSync(join(home, ".memwal")).mode & 0o7777, 0o700); + + // An override root is created 0700 as well, parents included. + const override = prepareRuntimeRoot(join(home, "elsewhere", "runtime")); + assert.equal(lstatSync(override).mode & 0o7777, 0o700); +}); + +test("the install never runs package scripts and never inherits the registry", () => { + const args = installArguments(`${MCP_PACKAGE_NAME}@${PIN}`, "/tmp/staging"); + assert.ok(args.includes("--ignore-scripts"), `--ignore-scripts missing from ${args.join(" ")}`); + assert.ok( + args.includes(`--registry=${MCP_REGISTRY}`), + `an explicit registry is missing from ${args.join(" ")}`, + ); + assert.equal(MCP_REGISTRY, "https://registry.npmjs.org/"); + + // npm puts the environment above every .npmrc, so a repository-supplied MCP env + // block would otherwise choose the registry no matter what cwd we install from. + const env = installEnvironment({ + PATH: "/usr/bin", + HOME: "/home/user", + NPM_CONFIG_REGISTRY: "http://attacker.example/", + npm_config_registry: "http://attacker.example/", + npm_config_ca: "-----BEGIN CERTIFICATE-----", + npm_config_ignore_scripts: "false", + NODE_OPTIONS: "--require /tmp/evil.js", + }); + assert.equal(env.PATH, "/usr/bin"); + assert.equal(env.HOME, "/home/user"); + assert.equal(env.npm_config_registry, MCP_REGISTRY); + assert.equal(env.npm_config_ignore_scripts, "true"); + assert.equal(env.NPM_CONFIG_REGISTRY, undefined); + assert.equal(env.npm_config_ca, undefined); + assert.equal(env.NODE_OPTIONS, undefined); +}); + +test("npm is spawned as node , with no .cmd and no shell, on win32 too", () => { + const nodeExe = "C:\\Program Files\\nodejs\\node.exe"; + const npmCli = "C:\\Program Files\\nodejs\\node_modules\\npm\\bin\\npm-cli.js"; + const args = installArguments("pkg@1.0.0", "C:\\Users\\a b\\.memwal\\runtime\\staging"); + + const plan = npmSpawnPlan({ kind: "js", path: npmCli }, args, { + platform: "win32", + execPath: nodeExe, + }); + + assert.equal(plan.command, nodeExe); + assert.equal(plan.shell, false, "node.exe must never be spawned through a shell"); + assert.deepEqual(plan.args, [npmCli, ...args]); + // Arguments go straight to the process, so they are NOT shell-quoted here. + assert.ok(plan.args.includes("C:\\Users\\a b\\.memwal\\runtime\\staging")); +}); + +test("the win32 npm.cmd fallback uses shell:true with cmd.exe-quoted arguments", () => { + // Since the CVE-2024-27980 fix (Node >= 18.20.2 / 20.12.2 / 21.7.3) spawning a + // .cmd without shell:true is EINVAL, which used to mean the server never started + // on Windows at all: ensureTrustedEntry -> install -> throw -> exit 1. + const npmCmd = "C:\\Program Files\\nodejs\\npm.cmd"; + const staging = "C:\\Users\\a b\\.memwal\\runtime\\staging"; + const plan = npmSpawnPlan({ kind: "bin", path: npmCmd }, installArguments("pkg@1.0.0", staging), { + platform: "win32", + }); + + assert.equal(plan.shell, true, "a .cmd target needs shell:true or spawnSync returns EINVAL"); + assert.equal(plan.command, `"${npmCmd}"`, "the command path contains a space and must be quoted"); + assert.ok(plan.args.includes(`"${staging}"`), "a path with a space must reach cmd.exe quoted"); + assert.ok(plan.args.includes("--ignore-scripts")); + assert.ok(plan.args.includes(`--registry=${MCP_REGISTRY}`)); + // Flags without metacharacters stay bare, so the command line stays readable. + assert.ok(plan.args.includes("install")); + assert.ok(plan.args.includes("pkg@1.0.0")); +}); + +test("quoteWindowsArgument quotes what cmd.exe would otherwise eat", () => { + assert.equal(quoteWindowsArgument("install"), "install"); + assert.equal(quoteWindowsArgument("@scope/pkg@1.0.0"), "@scope/pkg@1.0.0"); + assert.equal(quoteWindowsArgument("C:\\a b\\c"), '"C:\\a b\\c"'); + assert.equal(quoteWindowsArgument("a&b"), '"a&b"'); + assert.equal(quoteWindowsArgument('say "hi"'), '"say \\"hi\\""'); +}); + +test("a posix npm binary is spawned directly, never through a shell", () => { + const plan = npmSpawnPlan({ kind: "bin", path: "/usr/local/bin/npm" }, ["install"], { + platform: "linux", + }); + assert.equal(plan.command, "/usr/local/bin/npm"); + assert.equal(plan.shell, false); + assert.deepEqual(plan.args, ["install"]); + assert.throws(() => npmSpawnPlan({ kind: "none", path: null }, ["install"]), /could not locate npm/); +}); + +test("resolveNpm returns an absolute path and skips PATH entries inside the project", (t) => { + const found = resolveNpm(); + assert.notEqual(found.kind, "none", "npm must be resolvable in the test environment"); + assert.ok(resolve(found.path) === found.path, `${found.path} is not absolute`); + + const project = makeHostileProject(t); + const projectBin = join(project.dir, "node_modules", ".bin"); + mkdirSync(projectBin, { recursive: true }); + writeFileSync(join(projectBin, "npm"), "#!/bin/sh\necho PROJECT_NPM\n"); + chmodSync(join(projectBin, "npm"), 0o755); + + const elsewhere = mkdtempSync(join(tmpdir(), "memwal-npm-")); + t.after(() => rmSync(elsewhere, { recursive: true, force: true })); + writeFileSync(join(elsewhere, "npm"), "#!/bin/sh\necho OUTSIDE_NPM\n"); + chmodSync(join(elsewhere, "npm"), 0o755); + + // An execPath with no npm layout next to it forces the PATH scan. + const bare = mkdtempSync(join(tmpdir(), "memwal-bare-node-")); + t.after(() => rmSync(bare, { recursive: true, force: true })); + const resolved = resolveNpm({ + execPath: join(bare, "node"), + platform: "linux", + env: { PATH: [".", "relative/bin", projectBin, elsewhere].join(":") }, + cwd: project.dir, + }); + + assert.equal(resolved.kind, "bin"); + assert.equal(canonicalPath(resolved.path), canonicalPath(join(elsewhere, "npm"))); + assert.ok(!resolved.path.includes(project.dir), "npm must never come out of the project"); +}); + +test("a client started in the home directory still gets the default root", (t) => { + // enclosingProjectRoot() must not treat $HOME as a project: ~/.memwal/runtime is + // inside the home directory by design, and refusing it would break every client + // that starts its servers there. + const home = mkdtempSync(join(tmpdir(), "memwal-home-cwd-")); + t.after(() => rmSync(home, { recursive: true, force: true })); + const target = join(home, ".memwal", "runtime", `memwal-mcp@${PIN}`); + mkdirSync(target, { recursive: true }); + const entry = plantPackage(target, { version: PIN, marker: TRUSTED_MARKER }); + + const result = runLauncher(["--print-entry"], { + cwd: home, + env: { HOME: home, USERPROFILE: home, MEMWAL_MCP_RUNTIME_DIR: "" }, + }); + + assert.equal(result.status, 0, result.stderr); + assert.equal(result.stdout.trim(), entry); +}); + +test("a subdirectory of the home directory is a project again", (t) => { + const home = mkdtempSync(join(tmpdir(), "memwal-home-sub-")); + t.after(() => rmSync(home, { recursive: true, force: true })); + const project = join(home, "code", "victim"); + mkdirSync(project, { recursive: true }); + writeFileSync(join(project, "package.json"), JSON.stringify({ name: "victim" })); + const planted = join(project, ".memwal-runtime"); + mkdirSync(installDir(PIN, planted), { recursive: true }); + plantPackage(installDir(PIN, planted), { version: PIN, marker: LOCAL_MARKER }); + + const result = runLauncher(["--print-entry"], { + cwd: project, + env: { HOME: home, USERPROFILE: home, MEMWAL_MCP_RUNTIME_DIR: planted }, + }); + + assert.notEqual(result.status, 0); + assert.match(result.stderr, /inside the project tree/); + assert.doesNotMatch(result.stdout, new RegExp(LOCAL_MARKER)); +}); diff --git a/packages/mcp/test/unknown-flags.test.mjs b/packages/mcp/test/unknown-flags.test.mjs index f6c24de5b..5cf33f52a 100644 --- a/packages/mcp/test/unknown-flags.test.mjs +++ b/packages/mcp/test/unknown-flags.test.mjs @@ -68,6 +68,8 @@ test("parseArgs treats no known flag as unknown", () => { "--help", "-h", "--logout", "--login", "login", + "--approve-project", "approve-project", + "--revoke-project", "revoke-project", "--prod", "--dev", "--staging", "--local", "--relayer", "https://r.example", "--relayer-url", "https://r.example", @@ -85,6 +87,28 @@ test("parseArgs treats no known flag as unknown", () => { assert.deepEqual(parseArgs(known).unknown, []); }); +test("an unknown flag does not swallow the project-approval commands", () => { + // They are commands, not values. Swallowing one turns an approval the user + // typed into a run that silently approves nothing (WALM-639). + for (const command of ["approve-project", "revoke-project"]) { + const args = parseArgs(["--typo", command]); + assert.deepEqual(args.unknown, ["--typo"]); + assert.equal( + command === "approve-project" ? args.approveProject : args.revokeProject, + true, + `\`${command}\` was swallowed as a flag value`, + ); + } +}); + +test("--help documents how to approve project-local credentials", () => { + // The gate is only actionable if the command that lifts it is discoverable. + const help = helpText(); + assert.ok(help.includes("approve-project"), "approve-project missing from --help"); + assert.ok(help.includes("revoke-project"), "revoke-project missing from --help"); + assert.deepEqual(parseArgs(["approve-project"]).unknown, []); +}); + test("parseArgs does not mistake a flag's value for an unknown flag", () => { // `next()` consumes the value, so "MCP Client" must never be reported. const args = parseArgs(["--label", "MCP Client"]); diff --git a/packages/mcp/test/user-prompt-hook.test.mjs b/packages/mcp/test/user-prompt-hook.test.mjs index e9e0ccbee..672af8a0f 100644 --- a/packages/mcp/test/user-prompt-hook.test.mjs +++ b/packages/mcp/test/user-prompt-hook.test.mjs @@ -2,11 +2,20 @@ * UserPromptSubmit injects one full decision rubric per session, then a * one-line nudge. It must not classify remember vs recall from English * keywords — every substantive prompt in a fresh session gets the same text. + * + * WALM-642 made the save half of that rubric depend on the user's standing + * automatic-memory answer, so "the same text" is now per state: the tests below + * pin the ON variant by asking for it explicitly, and `auto-save-optin.test.mjs` + * pins what an answered-no and an unanswered install inject instead. Every run + * is pointed at an empty MEMWAL_CREDS_DIR so the developer's own + * ~/.memwal/settings.json cannot decide the result. */ import { test } from "node:test"; import assert from "node:assert/strict"; import { spawnSync } from "node:child_process"; -import { dirname, resolve } from "node:path"; +import { mkdtempSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { dirname, join, resolve } from "node:path"; import { fileURLToPath } from "node:url"; import { DECISION_RUBRIC, @@ -15,11 +24,20 @@ import { const __dirname = dirname(fileURLToPath(import.meta.url)); const HOOK = resolve(__dirname, "../plugin/scripts/on_user_prompt.mjs"); +const EMPTY_CREDS_DIR = mkdtempSync(join(tmpdir(), "memwal-hook-test-")); function runHook(prompt, sessionId = `test-${Math.random().toString(16).slice(2)}`) { const result = spawnSync(process.execPath, [HOOK], { input: JSON.stringify({ prompt, session_id: sessionId }), encoding: "utf8", + env: { + ...process.env, + MEMWAL_CREDS_DIR: EMPTY_CREDS_DIR, + // These cases are about classification, not consent: pin the + // automatic-save rubric explicitly so they keep testing the thing + // they were written for, whatever the resolver would decide. + MEMWAL_AUTO_SAVE: "1", + }, }); assert.equal(result.status, 0, result.stderr); if (!result.stdout.trim()) return ""; diff --git a/packages/python-sdk-memwal/CHANGELOG.md b/packages/python-sdk-memwal/CHANGELOG.md index 9593c4d2d..16cfbac8e 100644 --- a/packages/python-sdk-memwal/CHANGELOG.md +++ b/packages/python-sdk-memwal/CHANGELOG.md @@ -1,5 +1,11 @@ # memwal +## 0.1.11 + +### Added + +- `list_namespaces(cursor=None, limit=None)` lists the namespaces that hold memories (name, `memory_count`, `storage_used`, `updated_at`), so an agent can discover namespaces instead of guessing. Metadata only; no decryption. Paginate on `has_more`. `MemWalSync` and the mock clients have it too. + ## 0.1.10 ### Added diff --git a/packages/python-sdk-memwal/README.md b/packages/python-sdk-memwal/README.md index 1713142a9..718019e98 100644 --- a/packages/python-sdk-memwal/README.md +++ b/packages/python-sdk-memwal/README.md @@ -100,7 +100,7 @@ async def test_memory_flow(): assert "dark mode" in result.results[0].text ``` -The mock supports remember/job polling, bulk remember, recall, analyze, embed, ask, health, restore, `forget(blob_id)`, and `clear(namespace)`. For deterministic behavior, `analyze` stores its full input as one fact instead of invoking an LLM extractor. Its simple relevance score is for application tests, not production search-quality evaluation. +The mock supports remember/job polling, bulk remember, recall, analyze, embed, ask, health, restore, list_namespaces, `forget(blob_id)`, and `clear(namespace)`. For deterministic behavior, `analyze` stores its full input as one fact instead of invoking an LLM extractor. Its simple relevance score is for application tests, not production search-quality evaluation. ### Context Manager @@ -206,6 +206,7 @@ Create a new async client. | `await analyze(text, namespace?)` | Extract and store facts | | `await ask(question, limit?, namespace?)` | Ask a question answered using memories | | `await restore(namespace, limit?)` | Restore a namespace | +| `await list_namespaces(cursor?, limit?)` | List namespaces that hold memories; paginate on `has_more` | | `await health()` | Check server health | | `await remember_manual(opts)` | Store encrypted payload + pre-computed vector | | `await recall_manual(opts)` | Search with pre-computed vector | diff --git a/packages/python-sdk-memwal/memwal/__init__.py b/packages/python-sdk-memwal/memwal/__init__.py index 18c7f6e78..345808639 100644 --- a/packages/python-sdk-memwal/memwal/__init__.py +++ b/packages/python-sdk-memwal/memwal/__init__.py @@ -44,6 +44,8 @@ EmbedResult, HealthResult, MemWalConfig, + NamespacesResult, + NamespaceSummary, RecallManualHit, RecallManualOptions, RecallManualResult, @@ -114,6 +116,8 @@ "AnalyzedFact", "HealthResult", "RestoreResult", + "NamespaceSummary", + "NamespacesResult", "ScoringWeights", "RememberManualOptions", "RememberManualResult", @@ -122,4 +126,4 @@ "RecallManualResult", ] -__version__ = "0.1.10" +__version__ = "0.1.11" diff --git a/packages/python-sdk-memwal/memwal/client.py b/packages/python-sdk-memwal/memwal/client.py index e6faa2d62..29b9377f8 100644 --- a/packages/python-sdk-memwal/memwal/client.py +++ b/packages/python-sdk-memwal/memwal/client.py @@ -34,8 +34,8 @@ import time import uuid from datetime import datetime, timezone -from typing import Any, Dict, List, Optional, Sequence, Tuple, TypeVar, Union -from urllib.parse import ParseResult, urlparse +from typing import Any, Dict, List, Optional, Sequence, Set, Tuple, TypeVar, Union +from urllib.parse import ParseResult, urlencode, urlparse import httpx import nacl.signing @@ -50,6 +50,8 @@ EmbedResult, HealthResult, MemWalConfig, + NamespacesResult, + NamespaceSummary, RecallManualHit, RecallManualOptions, RecallManualResult, @@ -289,6 +291,8 @@ def __init__(self, config: MemWalConfig) -> None: self._session_build_task: Optional[asyncio.Task[str]] = None self._relayer_version_metadata: Optional[Dict[str, Any]] = None self._compatibility_lock: Optional[asyncio.Lock] = None + self._owner_address: Optional[str] = None + self._owner_task: Optional[asyncio.Task[str]] = None # Preserve a generated key across an ambiguous transport failure. A # subsequent identical call then collapses onto the accepted paid job. self._pending_remember_keys: Dict[str, str] = {} @@ -620,7 +624,7 @@ async def wait_for_remember_jobs( continue raise - still_pending: List[str] = [] + terminal_ids: Set[str] = set() for item in batch.results: if item.status == "done": results[item.job_id] = RememberBulkItemResult( @@ -629,6 +633,7 @@ async def wait_for_remember_jobs( status="done", error=None, ) + terminal_ids.add(item.job_id) elif item.status in ("failed", "not_found"): results[item.job_id] = RememberBulkItemResult( id=item.job_id, @@ -636,9 +641,8 @@ async def wait_for_remember_jobs( status="failed", error=item.error, ) - else: - still_pending.append(item.job_id) - pending = still_pending + terminal_ids.add(item.job_id) + pending = [jid for jid in pending if jid not in terminal_ids] ordered = [results[job_id] for job_id in job_ids] succeeded = sum(1 for r in ordered if r.status == "done") @@ -962,6 +966,72 @@ async def restore(self, namespace: str, limit: int = 10) -> RestoreResult: failed=data.get("failed", 0), ) + async def list_namespaces( + self, + cursor: Optional[str] = None, + limit: Optional[int] = None, + ) -> NamespacesResult: + """List the namespaces this account holds memories in. + + Recall needs a namespace to search. Without this, an agent on an + unfamiliar account has to guess names or fall back to ``"default"``. + Returns metadata only: no blob fetch, no decryption. + + Namespaces are flat and exact-match. To work with a prefix such as + ``proj/``, filter the names here and recall each one. + + Paginate on ``has_more``, NOT page length: the relayer clamps + ``limit``, so asking for more than the cap returns exactly the cap. + + Example:: + + cursor = None + while True: + page = await memwal.list_namespaces(cursor=cursor) + for ns in page.namespaces: + print(ns.name, ns.memory_count) + cursor = page.next_cursor + if not page.has_more: + break + + Args: + cursor: A previous page's ``next_cursor``, to continue a walk or + poll for namespaces changed since then. Opaque; not a + timestamp or a namespace name. + limit: Page size. The relayer defaults to 100 and clamps to 500. + + Returns: + :class:`NamespacesResult`. + """ + owner = await self._resolve_owner() + + params: Dict[str, str] = {} + if cursor is not None: + params["updated_after"] = cursor + if limit is not None: + params["limit"] = str(limit) + query = urlencode(params) + + # The query string is part of the signed path: the relayer verifies + # against `path_and_query`, not `path`. + path = f"/v1/owners/{owner}/namespaces" + (f"?{query}" if query else "") + data = await self._signed_request("GET", path, {}, include_seal_session=False) + return NamespacesResult( + namespaces=[ + NamespaceSummary( + id=ns["id"], + name=ns["name"], + memory_count=ns["memory_count"], + storage_used=ns["storage_used"], + updated_at=ns["updated_at"], + ) + for ns in data["namespaces"] + ], + next_cursor=data.get("next_cursor"), + has_more=data["has_more"], + snapshot_version=data["snapshot_version"], + ) + async def health(self) -> HealthResult: """Check server health. No authentication required. @@ -1255,6 +1325,43 @@ async def _build_seal_session(self) -> str: finally: self._session_build_task = None + async def _resolve_owner_inner(self) -> str: + # POST /api/stats authenticates with the same delegate scheme and + # returns the owner the relayer resolved from our key. Same approach + # as the TypeScript SDK's resolveOwner(). + data = await self._signed_request( + "POST", + "/api/stats", + {"namespace": self._namespace}, + include_seal_session=False, + ) + owner = data.get("owner") + if not owner: + raise MemWalError( + "Walrus Memory could not resolve this account's owner address " + "(POST /api/stats returned no owner)." + ) + self._owner_address = owner + return owner + + async def _resolve_owner(self) -> str: + """Owner address for this account, memoised for the client's life. + + The owner-scoped read routes take the address in the path, but the + client is configured with only a delegate key and account id. + """ + if self._owner_address is not None: + return self._owner_address + + if self._owner_task is not None: + return await self._owner_task + + self._owner_task = asyncio.create_task(self._resolve_owner_inner()) + try: + return await self._owner_task + finally: + self._owner_task = None + async def _signed_request( self, method: str, @@ -1668,6 +1775,14 @@ def restore(self, namespace: str, limit: int = 10) -> RestoreResult: (matches server + TypeScript SDK).""" return self._run(self._inner.restore(namespace, limit)) + def list_namespaces( + self, + cursor: Optional[str] = None, + limit: Optional[int] = None, + ) -> NamespacesResult: + """Synchronous version of :meth:`MemWal.list_namespaces`.""" + return self._run(self._inner.list_namespaces(cursor, limit)) + def health(self) -> HealthResult: """Synchronous version of :meth:`MemWal.health`.""" return self._run(self._inner.health()) diff --git a/packages/python-sdk-memwal/memwal/mock.py b/packages/python-sdk-memwal/memwal/mock.py index 06da3d105..31737c024 100644 --- a/packages/python-sdk-memwal/memwal/mock.py +++ b/packages/python-sdk-memwal/memwal/mock.py @@ -3,11 +3,13 @@ from __future__ import annotations import asyncio +import base64 +import json import math import re import unicodedata from dataclasses import dataclass -from datetime import datetime +from datetime import datetime, timedelta, timezone from typing import Any, Dict, List, Optional, Sequence, Union from .types import ( @@ -18,6 +20,8 @@ AskResult, EmbedResult, HealthResult, + NamespacesResult, + NamespaceSummary, RecallMemory, RecallParams, RecallResult, @@ -34,6 +38,8 @@ RestoreResult, ) +_MOCK_NAMESPACE_EPOCH = datetime(2026, 1, 1, tzinfo=timezone.utc) + @dataclass class MemWalMockSeed: @@ -66,6 +72,24 @@ def _distance(query_tokens: set[str], text: str) -> float: return 1.0 - matches / len(query_tokens) +def _mock_timestamp(sequence: int) -> str: + # Same shape as the TypeScript mock's Date.toISOString(), so fixed-width + # timestamps compare correctly as strings. + moment = _MOCK_NAMESPACE_EPOCH + timedelta(seconds=sequence) + return moment.strftime("%Y-%m-%dT%H:%M:%S.000Z") + + +def _encode_namespaces_cursor(payload: Dict[str, Any]) -> str: + # The relayer's cursor: URL-safe unpadded base64 of a JSON object. + raw = json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode("utf-8") + return base64.urlsafe_b64encode(raw).decode("ascii").rstrip("=") + + +def _decode_namespaces_cursor(cursor: str) -> Dict[str, Any]: + padded = cursor + "=" * (-len(cursor) % 4) + return json.loads(base64.urlsafe_b64decode(padded)) + + def _validate_text(text: str, field: str = "text") -> None: if not isinstance(text, str) or not text.strip(): raise ValueError(f"{field} cannot be empty") @@ -344,6 +368,67 @@ async def restore(self, namespace: str, limit: int = 10) -> RestoreResult: failed=0, ) + async def list_namespaces( + self, + cursor: Optional[str] = None, + limit: Optional[int] = None, + ) -> NamespacesResult: + grouped: Dict[str, List[_Memory]] = {} + for memory in self._memories: + grouped.setdefault(memory.namespace, []).append(memory) + summaries = sorted( + ( + NamespaceSummary( + id=f"mock-ns-{name}", + name=name, + memory_count=len(memories), + storage_used=sum(len(m.text.encode("utf-8")) for m in memories), + updated_at=_mock_timestamp(max(m.sequence for m in memories)), + ) + for name, memories in grouped.items() + ), + key=lambda ns: (ns.updated_at, ns.name), + ) + + # Keyset walk pinned to a snapshot, like the relayer: writes made + # mid-walk surface on the next poll, not in the current walk. + after = _decode_namespaces_cursor(cursor) if cursor is not None else None + snapshot_at = (after or {}).get("snapshot_at") or _mock_timestamp(self._sequence) + remaining = [ + ns + for ns in summaries + if ns.updated_at <= snapshot_at + and ( + after is None + or (ns.updated_at, ns.name) > (after["updated_at"], after["namespace"]) + ) + ] + page = remaining if limit is None else remaining[:limit] + has_more = len(remaining) > len(page) + + watermark = ( + {"updated_at": page[-1].updated_at, "namespace": page[-1].name} if page else after + ) + next_cursor = None + if watermark is not None: + next_cursor = _encode_namespaces_cursor( + { + "updated_at": watermark["updated_at"], + "namespace": watermark["namespace"], + # A finished walk drops the snapshot so the next poll + # takes a fresh one. + "snapshot_at": snapshot_at if has_more else None, + } + ) + + return NamespacesResult( + namespaces=page, + next_cursor=next_cursor, + has_more=has_more, + # Matches the live relayer's current wire-format version. + snapshot_version=2, + ) + async def health(self) -> HealthResult: return HealthResult( status="ok", @@ -582,6 +667,13 @@ def ask( def restore(self, namespace: str, limit: int = 10) -> RestoreResult: return self._run(self._inner.restore(namespace, limit)) + def list_namespaces( + self, + cursor: Optional[str] = None, + limit: Optional[int] = None, + ) -> NamespacesResult: + return self._run(self._inner.list_namespaces(cursor, limit)) + def health(self) -> HealthResult: return self._run(self._inner.health()) diff --git a/packages/python-sdk-memwal/memwal/types.py b/packages/python-sdk-memwal/memwal/types.py index 09fdc49d8..cadaed9be 100644 --- a/packages/python-sdk-memwal/memwal/types.py +++ b/packages/python-sdk-memwal/memwal/types.py @@ -228,6 +228,34 @@ class RestoreResult: failed: int = 0 +@dataclass +class NamespaceSummary: + """One namespace in a :meth:`MemWal.list_namespaces` page.""" + + id: str + name: str + memory_count: int + storage_used: int + #: ``MAX(updated_at)`` across the namespace's memories (RFC 3339), the + #: same value the relayer builds the keyset cursor from. + updated_at: str + + +@dataclass +class NamespacesResult: + """Result from list_namespaces().""" + + namespaces: List[NamespaceSummary] + #: Pass back as ``cursor`` on the next call. Set on every page, including + #: the last, so a caller that finished a walk can poll from it later. + next_cursor: Optional[str] + #: Whether to keep paginating. Do NOT infer this from page length: the + #: relayer clamps ``limit``, so asking for more than the cap returns + #: exactly the cap. + has_more: bool + snapshot_version: int + + @dataclass class AskMemory: """A memory used to answer a question.""" diff --git a/packages/python-sdk-memwal/pyproject.toml b/packages/python-sdk-memwal/pyproject.toml index 549018172..f81904191 100644 --- a/packages/python-sdk-memwal/pyproject.toml +++ b/packages/python-sdk-memwal/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "memwal" -version = "0.1.10" +version = "0.1.11" description = "Python SDK for Walrus Memory — Privacy-first AI memory with Ed25519 signing" readme = "README.md" license = "MIT" diff --git a/packages/python-sdk-memwal/tests/test_client.py b/packages/python-sdk-memwal/tests/test_client.py index 44adedd90..47cc3c1bc 100644 --- a/packages/python-sdk-memwal/tests/test_client.py +++ b/packages/python-sdk-memwal/tests/test_client.py @@ -10,7 +10,7 @@ import base64 import json from datetime import datetime, timedelta, timezone -from typing import Any +from typing import Any, Sequence import httpx import nacl.signing @@ -29,6 +29,9 @@ RecallParams, RememberBulkAcceptedResult, RememberBulkItem, + RememberBulkOptions, + RememberBulkStatusItem, + RememberBulkStatusResult, RememberManualOptions, ScoringWeights, ) @@ -428,6 +431,36 @@ async def test_matching_job_ids_returns_accepted( assert result.total == 2 assert result.status == "pending" + async def test_job_omitted_from_a_poll_stays_pending( + self, memwal_client: MemWal + ) -> None: + polls = 0 + + async def partial_status( + job_ids: Sequence[str], + ) -> RememberBulkStatusResult: + nonlocal polls + polls += 1 + job_id = "job-1" if polls == 1 else "job-2" + return RememberBulkStatusResult( + results=[ + RememberBulkStatusItem( + job_id=job_id, + status="done", + blob_id=f"blob-for-{job_id}", + ) + ] + ) + + memwal_client.get_remember_bulk_status = partial_status # type: ignore[method-assign] + + result = await memwal_client.wait_for_remember_jobs( + ["job-1", "job-2"], + RememberBulkOptions(poll_interval_ms=1, timeout_ms=5_000), + ) + + assert [item.status for item in result.results] == ["done", "done"] + # ============================================================ # recall() tests diff --git a/packages/python-sdk-memwal/tests/test_integration.py b/packages/python-sdk-memwal/tests/test_integration.py index 59f17f9ef..842e8a9ee 100644 --- a/packages/python-sdk-memwal/tests/test_integration.py +++ b/packages/python-sdk-memwal/tests/test_integration.py @@ -257,6 +257,30 @@ def test_remember_custom_namespace(self) -> None: assert result.namespace == _E2E_NAMESPACE_ALT +@requires_key +class TestListNamespaces: + """list_namespaces() against live server. Read-only; writes nothing.""" + + def test_walks_every_page_on_has_more(self) -> None: + # limit=1 forces a cursor into the signed query string on every page + # after the first, which is what the relayer verifies the signature over. + mw = _sync_client() + names: list[str] = [] + cursor = None + for _ in range(1000): + page = mw.list_namespaces(cursor=cursor, limit=1) + assert len(page.namespaces) <= 1 + assert page.snapshot_version >= 1 + names.extend(ns.name for ns in page.namespaces) + cursor = page.next_cursor + if not page.has_more: + break + else: + pytest.fail("namespace walk did not finish in 1000 pages") + assert len(names) == len(set(names)), "a single walk must not repeat a namespace" + print(f"\n namespaces={len(names)}") + + @requires_key class TestRecall: """recall() against live server.""" diff --git a/packages/python-sdk-memwal/tests/test_list_namespaces.py b/packages/python-sdk-memwal/tests/test_list_namespaces.py new file mode 100644 index 000000000..c0bb58fb2 --- /dev/null +++ b/packages/python-sdk-memwal/tests/test_list_namespaces.py @@ -0,0 +1,220 @@ +""" +Tests for ``MemWal.list_namespaces()`` — owner-scoped namespace discovery. + +Mirrors ``packages/sdk/test/list-namespaces.test.mjs`` so both SDKs pin the +same wire contract. +""" + +from __future__ import annotations + +import asyncio +from urllib.parse import parse_qs + +import httpx +import nacl.signing +import pytest +import respx + +from memwal import MemWal, MemWalError, MemWalSync, NamespacesResult, NamespaceSummary +from memwal.utils import build_signature_message, bytes_to_hex, sha256_hex + +_SERVER = "https://relayer.example" +_OWNER = "0xowner0000000000000000000000000000000000000000000000000000000001" +_NAMESPACES_URL = f"{_SERVER}/v1/owners/{_OWNER}/namespaces" +_KEY_HEX = bytes_to_hex(bytes(nacl.signing.SigningKey(b"\x01" * 32))) + +_STATS = {"memory_count": 0, "storage_bytes": 0, "namespace": "default", "owner": _OWNER} + +_ONE_PAGE = { + "namespaces": [ + { + "id": "ns-1", + "name": "work", + "memory_count": 12, + "storage_used": 2048, + "updated_at": "2026-08-20T10:00:00Z", + } + ], + "next_cursor": "eyJ1cGRhdGVkX2F0IjoiMjAyNi0wOC0yMFQxMDowMDowMFoifQ", + "has_more": False, + "snapshot_version": 2, +} + + +def _client() -> MemWal: + return MemWal.create(key=_KEY_HEX, account_id="0x1", server_url=_SERVER) + + +def _stub_relayer( + namespaces_body: dict = _ONE_PAGE, + stats_body: dict | None = None, +) -> tuple[respx.Route, respx.Route]: + """Stub the three calls a list_namespaces() round-trip makes: the + compatibility preflight, the owner resolution, and the read itself. + + Anything else (``/config``, Sui GraphQL for a SEAL session) is unmocked, + so respx fails the test if the client reaches for it. + """ + respx.get(f"{_SERVER}/version").mock( + return_value=httpx.Response( + 200, + json={ + "apiVersion": "1.0.0", + "relayerVersion": "1.0.0", + "minSupportedSdk": {"typescript": "0.0.4", "python": "0.1.0", "mcp": "0.0.1"}, + }, + ) + ) + stats = respx.post(f"{_SERVER}/api/stats").mock( + return_value=httpx.Response( + 200, + json=stats_body if stats_body is not None else _STATS, + ) + ) + namespaces = respx.get(_NAMESPACES_URL).mock( + return_value=httpx.Response(200, json=namespaces_body) + ) + return stats, namespaces + + +class TestListNamespaces: + @respx.mock + async def test_reads_the_owner_scoped_namespaces_path(self) -> None: + _, namespaces = _stub_relayer() + + await _client().list_namespaces() + + assert namespaces.call_count == 1 + request = namespaces.calls[0].request + assert request.method == "GET" + assert request.url.path == f"/v1/owners/{_OWNER}/namespaces" + + @respx.mock + async def test_resolves_the_owner_once_and_reuses_it(self) -> None: + stats, namespaces = _stub_relayer() + memwal = _client() + + await memwal.list_namespaces() + await memwal.list_namespaces() + + assert stats.call_count == 1, "owner resolution must be memoised across calls" + assert namespaces.call_count == 2 + + @respx.mock + async def test_concurrent_first_calls_share_one_owner_lookup(self) -> None: + stats, namespaces = _stub_relayer() + owner_response = stats.return_value + + async def slow_stats(request: httpx.Request) -> httpx.Response: + # respx otherwise answers without yielding, so the two calls + # would run back to back instead of overlapping. + await asyncio.sleep(0.01) + return owner_response + + stats.side_effect = slow_stats + memwal = _client() + + await asyncio.gather(memwal.list_namespaces(), memwal.list_namespaces()) + + assert stats.call_count == 1 + assert namespaces.call_count == 2 + + @respx.mock + async def test_forwards_cursor_as_updated_after_and_passes_limit(self) -> None: + _, namespaces = _stub_relayer() + + await _client().list_namespaces(cursor="opaque-cursor_1", limit=25) + + params = parse_qs(namespaces.calls[0].request.url.query.decode()) + assert params == {"updated_after": ["opaque-cursor_1"], "limit": ["25"]} + + @respx.mock + async def test_omits_query_params_that_were_not_supplied(self) -> None: + _, namespaces = _stub_relayer() + + await _client().list_namespaces() + + assert namespaces.calls[0].request.url.query == b"" + + @respx.mock + async def test_signature_covers_the_query_string(self) -> None: + # The relayer verifies against `path_and_query`, not `path`. + _, namespaces = _stub_relayer() + + await _client().list_namespaces(cursor="abc", limit=5) + + request = namespaces.calls[0].request + assert request.content == b"" + headers = request.headers + message = build_signature_message( + timestamp=headers["x-timestamp"], + method="GET", + path=f"/v1/owners/{_OWNER}/namespaces?updated_after=abc&limit=5", + body_sha256=sha256_hex(""), + nonce=headers["x-nonce"], + account_id=headers["x-account-id"], + ) + verify_key = nacl.signing.VerifyKey(bytes.fromhex(headers["x-public-key"])) + verify_key.verify(message.encode("utf-8"), bytes.fromhex(headers["x-signature"])) + + @respx.mock + async def test_returns_the_relayer_wire_shape(self) -> None: + _stub_relayer() + + result = await _client().list_namespaces() + + assert result == NamespacesResult( + namespaces=[ + NamespaceSummary( + id="ns-1", + name="work", + memory_count=12, + storage_used=2048, + updated_at="2026-08-20T10:00:00Z", + ) + ], + next_cursor=_ONE_PAGE["next_cursor"], + has_more=False, + snapshot_version=2, + ) + + @respx.mock + async def test_sends_no_seal_session_on_a_metadata_only_read(self) -> None: + stats, namespaces = _stub_relayer() + + await _client().list_namespaces() + + for route in (stats, namespaces): + assert "x-seal-session" not in route.calls[0].request.headers + + @respx.mock + async def test_raises_when_stats_returns_no_owner(self) -> None: + _stub_relayer(stats_body={"memory_count": 0, "storage_bytes": 0, "namespace": "default"}) + + with pytest.raises(MemWalError, match="owner"): + await _client().list_namespaces() + + @respx.mock + async def test_a_failed_owner_lookup_is_retried_on_the_next_call(self) -> None: + stats, namespaces = _stub_relayer() + stats.side_effect = [ + httpx.Response(503, text="busy"), + httpx.Response(200, json=_STATS), + ] + memwal = _client() + + with pytest.raises(MemWalError): + await memwal.list_namespaces() + await memwal.list_namespaces() + + assert stats.call_count == 2 + assert namespaces.call_count == 1 + + @respx.mock + def test_sync_wrapper(self) -> None: + _stub_relayer() + client = MemWalSync.create(key=_KEY_HEX, account_id="0x1", server_url=_SERVER) + + result = client.list_namespaces(limit=10) + + assert [ns.name for ns in result.namespaces] == ["work"] diff --git a/packages/python-sdk-memwal/tests/test_mock.py b/packages/python-sdk-memwal/tests/test_mock.py index 30ef35cba..9a3e006fc 100644 --- a/packages/python-sdk-memwal/tests/test_mock.py +++ b/packages/python-sdk-memwal/tests/test_mock.py @@ -1,6 +1,9 @@ """Offline mock client regression tests.""" +import base64 import inspect +import json +import re import pytest @@ -162,3 +165,102 @@ async def test_sync_mock_works_inside_an_existing_event_loop(): assert stored.namespace == "notebook" assert recalled.results[0].text == "called from a running loop" + + +@pytest.mark.asyncio +async def test_mock_list_namespaces_aggregates_memories_by_namespace(): + mock = MemWalMock.create( + initial_memories=[ + MemWalMockSeed(text="one", namespace="work"), + MemWalMockSeed(text="two", namespace="work"), + MemWalMockSeed(text="旅行", namespace="home"), + ] + ) + + page = await mock.list_namespaces() + by_name = {ns.name: ns for ns in page.namespaces} + + assert sorted(by_name) == ["home", "work"] + assert by_name["work"].memory_count == 2 + assert by_name["work"].storage_used == 6 + assert by_name["home"].storage_used == len("旅行".encode("utf-8")) + assert page.has_more is False + # Matches the live relayer's current wire-format version. + assert page.snapshot_version == 2 + + +@pytest.mark.asyncio +async def test_mock_list_namespaces_reports_has_more_when_limit_truncates(): + mock = MemWalMock.create( + initial_memories=[ + MemWalMockSeed(text="a", namespace="alpha"), + MemWalMockSeed(text="b", namespace="bravo"), + MemWalMockSeed(text="c", namespace="charlie"), + ] + ) + + page = await mock.list_namespaces(limit=2) + + assert len(page.namespaces) == 2 + assert page.has_more is True, "has_more is the pagination signal, not page length" + assert page.next_cursor + + +@pytest.mark.asyncio +async def test_mock_namespace_cursor_uses_the_relayer_wire_format_and_resets_after_a_walk(): + mock = MemWalMock.create( + initial_memories=[ + MemWalMockSeed(text="a", namespace="旅行"), + MemWalMockSeed(text="b", namespace="work"), + ] + ) + + first = await mock.list_namespaces(limit=1) + assert re.fullmatch(r"[A-Za-z0-9_-]+", first.next_cursor) + padded = first.next_cursor + "=" * (-len(first.next_cursor) % 4) + cursor = json.loads(base64.urlsafe_b64decode(padded)) + assert cursor["namespace"] == "旅行" + assert cursor["updated_at"] == first.namespaces[0].updated_at + assert cursor["snapshot_at"] + + last = await mock.list_namespaces(cursor=first.next_cursor) + assert [ns.name for ns in last.namespaces] == ["work"] + assert last.has_more is False + padded = last.next_cursor + "=" * (-len(last.next_cursor) % 4) + assert json.loads(base64.urlsafe_b64decode(padded))["snapshot_at"] is None + + empty = await mock.list_namespaces(cursor=last.next_cursor) + assert empty.namespaces == [] + assert empty.next_cursor == last.next_cursor + + +@pytest.mark.asyncio +async def test_mock_namespace_walk_defers_new_writes_until_the_next_poll(): + mock = MemWalMock.create( + initial_memories=[ + MemWalMockSeed(text="a", namespace="alpha"), + MemWalMockSeed(text="b", namespace="bravo"), + ] + ) + + first = await mock.list_namespaces(limit=1) + await mock.remember("new", "bravo") + last = await mock.list_namespaces(cursor=first.next_cursor) + assert last.namespaces == [] + assert last.has_more is False + + poll = await mock.list_namespaces(cursor=last.next_cursor) + assert [ns.name for ns in poll.namespaces] == ["bravo"] + assert poll.namespaces[0].memory_count == 2 + + +def test_sync_mock_list_namespaces_matches_production(): + mock = MemWalMockSync.create(namespace="sync") + mock.remember("sync memory") + + page = mock.list_namespaces(limit=10) + + assert [ns.name for ns in page.namespaces] == ["sync"] + assert inspect.signature(MemWalMockSync.list_namespaces) == inspect.signature( + MemWalSync.list_namespaces + ) diff --git a/packages/sdk/CHANGELOG.md b/packages/sdk/CHANGELOG.md index 75fa71aa5..412e7c46e 100644 --- a/packages/sdk/CHANGELOG.md +++ b/packages/sdk/CHANGELOG.md @@ -1,5 +1,16 @@ # @mysten-incubation/memwal +## 0.1.8 + +### Added + +- `recall()` sends its own deadline as `deadline_ms` (14000, a second under the 15s it aborts at, so the relayer's answer beats the abort even on a slow connect) in the request body. A relayer that reads it answers a recall about to miss that deadline with a 504 whose `serverCode` is `RECALL_TIMEOUT` and whose body names the stuck step (`auth`, `embed`, `vector_search`, `walrus_download`, `seal_decrypt`), instead of the request aborting with no cause. Relayers that predate the field ignore it. (WALM-396) + +### Fixed + +- Every relayer request now carries a deadline. `fetch` has none of its own and the SDK passed an abort signal on exactly one method (`recall`, 15s), so the accept POST, every job-status read, and the `/version` and `/config` handshake calls could stay pending for as long as the socket stayed open. A poll loop checks its budget at the *top* of each iteration, which bounds when the next request starts rather than how long one takes — so a single stalled read ran straight past `timeoutMs`, and a `memwal_remember` documented as capping at 90s was observed by an MCP client still running after 120s. Requests default to 30s, matching the relayer's own outbound client; set it with `requestTimeoutMs` on `MemWal.create`, where a non-positive or non-finite value falls back to the default rather than disabling the bound. `restore` (60s) and `analyze` (60s) carry their own, since the route self-bounds at 55s and the extractor LLM runs inline respectively. Inside the wait loops each poll is bounded by the client deadline clamped to the remaining budget, so one stalled poll can neither outlive the wait nor swallow it. An expired request raises `MemWalRequestTimeout` with `status: 504`, which the existing transient-poll handling already retries; a caller's own abort and every other transport error propagate unchanged. +- Generated idempotency keys are derived from the content (a 30-minute bucket plus namespace and text) instead of `crypto.randomUUID()`. The per-instance key map only ever deduped retries that reused one client, and callers such as the MCP sidecar build a fresh client per session — so a replay after a reconnect read as a brand-new write and the relayer minted a second paid Walrus blob for one already in flight. The bucket bounds the collapse, since `remember_jobs` rows are never pruned. Callers passing an explicit `idempotencyKey` are unaffected, and distinct text or namespaces still derive distinct keys. + ## 0.1.7 ### Added @@ -12,6 +23,7 @@ - Declare `engines.node >= 20.0.0`, matching `memwal-mcp` and `openclaw-memory-memwal`. The SDK was the only published package without a floor. (WALM-599) - Empty-body 401s now use the same AUTH_REJECTED troubleshooting message as credential 401s instead of telling callers to run `memwal_login`. Headless SDK clients do not have that MCP tool. - `account.ts` and `manual.ts` PTBs use typed `tx.pure` helpers instead of the legacy untyped moveCall argument syntax that fails under modern `@mysten/sui`. +- An explicit `sort` on `recall()`, `"relevance"` included, is now the order: the relayer ignores `scoringWeights` for that request, and weights re-rank only when `sort` is omitted. Setting both used to return neither order, so `sort: "recent"` stopped meaning newest-first once `scoringWeights` carried a recency weight. (WALM-470) ## 0.1.6 diff --git a/packages/sdk/package.json b/packages/sdk/package.json index ca57884f5..e2a6472bb 100644 --- a/packages/sdk/package.json +++ b/packages/sdk/package.json @@ -1,7 +1,7 @@ { "name": "@mysten-incubation/memwal", - "version": "0.1.7", - "description": "Walrus Memory — Privacy-first AI memory SDK with Ed25519 delegate key auth", + "version": "0.1.8", + "description": "Walrus Memory \u2014 Privacy-first AI memory SDK with Ed25519 delegate key auth", "type": "module", "engines": { "node": ">=20.0.0" diff --git a/packages/sdk/src/ai/middleware.ts b/packages/sdk/src/ai/middleware.ts index dd4a7b72a..776521c25 100644 --- a/packages/sdk/src/ai/middleware.ts +++ b/packages/sdk/src/ai/middleware.ts @@ -165,6 +165,8 @@ export function withMemWal( }, }); + wrapped.specificationVersion = model.specificationVersion; + // Lets short-lived callers await outstanding auto-save writes before // exiting, e.g. `await model.flush()` right before `process.exit()`. wrapped.flush = async (): Promise => { @@ -212,32 +214,30 @@ export function injectMemoryContext( -1 ); + const memoryMessage = { + role: "user" as const, + content: [{ type: "text" as const, text: memoryContext }], + }; + + const result = [...prompt]; if (lastUserIndex > 0) { - const result = [...prompt]; - result.splice( - lastUserIndex, - 0, - { - role: "system" as const, - content: UNTRUSTED_MEMORY_SYSTEM_INSTRUCTION, - }, - { - role: "user" as const, - content: [{ type: "text" as const, text: memoryContext }], - } - ); - return result; + result.splice(lastUserIndex, 0, memoryMessage); + } else { + result.unshift(memoryMessage); } - return [ - { + const leading = result[0] as { role?: string; content?: unknown }; + if (leading.role === "system") { + result[0] = { + ...leading, + content: `${leading.content}\n\n${UNTRUSTED_MEMORY_SYSTEM_INSTRUCTION}`, + }; + } else { + result.unshift({ role: "system" as const, content: UNTRUSTED_MEMORY_SYSTEM_INSTRUCTION, - }, - { - role: "user" as const, - content: [{ type: "text" as const, text: memoryContext }], - }, - ...prompt, - ]; + }); + } + + return result; } diff --git a/packages/sdk/src/memwal.ts b/packages/sdk/src/memwal.ts index c49c90481..a7dec05d9 100644 --- a/packages/sdk/src/memwal.ts +++ b/packages/sdk/src/memwal.ts @@ -115,6 +115,15 @@ const SEAL_SESSION_TTL_MIN = 5; // a key server that sees it as expired. const SEAL_SESSION_SAFETY_MARGIN_MS = 30_000; +/** Per-call knobs for `signedRequest`. `timeoutMs` overrides the client-wide + * deadline for one endpoint; `signal` is the caller's own cancellation and is + * honoured alongside it, never replaced by it. */ +interface SignedRequestOptions { + includeDelegateKey?: boolean; + signal?: AbortSignal; + timeoutMs?: number; +} + type RememberStatusResponse = RememberJobStatus | { error?: string }; function sleep(ms: number): Promise { @@ -128,10 +137,148 @@ function pollingDelayMs(baseMs: number, attempt: number): number { return Math.floor(capped * jitter); } +/** Window over which the same (namespace, text) resolves to one key. + * + * `pendingRememberKeys` only dedupes retries that reuse one client instance, + * and callers such as the MCP sidecar build a fresh client per session — so a + * random key let a reconnect replay mint a second paid blob. The bucket bounds + * the collapse: `remember_jobs` rows are never pruned, so an unbucketed key + * would dedupe against a job from any point in history. */ +const IDEMPOTENCY_BUCKET_MS = 30 * 60 * 1000; + +async function derivedIdempotencyKey(requestIdentity: string): Promise { + const bucket = Math.floor(Date.now() / IDEMPOTENCY_BUCKET_MS); + return `r1-${await sha256hex(`${bucket}\0${requestIdentity}`)}`; +} + +/** Deadline for a request that names no other. + * + * `fetch` imposes none, and a poll loop checks its budget only between polls — + * so one stalled read ran past `timeoutMs` entirely. 30s mirrors the relayer's + * own outbound client. */ +const DEFAULT_REQUEST_TIMEOUT_MS = 30_000; + +/** `POST /api/restore` bounds itself at 55s server-side and answers with an + * error rather than going quiet, so the client must outlast that or it would + * abandon a response already on its way. */ +const RESTORE_REQUEST_TIMEOUT_MS = 60_000; + +/** `POST /api/analyze` runs the extractor LLM inline before it accepts, so it + * is the one write that legitimately outruns the default. */ +const ANALYZE_REQUEST_TIMEOUT_MS = 60_000; + +/** `POST /api/recall` has carried its own 15s deadline since before the rest + * had any; keeping it named makes that a decision rather than an accident. */ +const RECALL_REQUEST_TIMEOUT_MS = 15_000; + +/** What `recall()` tells the relayer it will wait, so the 504 naming the stuck + * stage arrives before the abort above rather than after it. A second under + * that abort because the two clocks start at different moments: this timer + * starts before `fetch`, the relayer's when the request lands, and the margin + * the relayer keeps for its reply cannot also cover DNS, TCP and TLS. Derived + * from the timeout so the two can never drift apart. */ +const RECALL_DEADLINE_MS = RECALL_REQUEST_TIMEOUT_MS - 1_000; + +/** + * Abort signal that fires after `ms`, or when `caller` aborts — whichever is + * first. Built by hand rather than with `AbortSignal.any`, which lands too + * recently to rely on across Node, browsers and Workers alike. + * + * Returns the signal plus a `dispose` the caller must run in a `finally`, so a + * pending timer never outlives its request. + */ +function deadlineSignal( + ms: number, + caller?: AbortSignal, +): { signal: AbortSignal; timedOut: () => boolean; dispose: () => void } { + const controller = new AbortController(); + let expired = false; + + const onCallerAbort = () => controller.abort(caller?.reason); + if (caller) { + if (caller.aborted) controller.abort(caller.reason); + else caller.addEventListener("abort", onCallerAbort, { once: true }); + } + + const timer = setTimeout(() => { + expired = true; + controller.abort(); + }, ms); + // Deliberately NOT unref'd. A deadline is the one timer somebody IS + // waiting on: unref'd, it stops firing the moment nothing else holds the + // loop open, and the stalled request it was meant to bound hangs forever + // instead. `dispose()` runs in the caller's `finally`, so the timer cannot + // outlive its request either way. + + return { + signal: controller.signal, + timedOut: () => expired, + dispose: () => { + clearTimeout(timer); + caller?.removeEventListener("abort", onCallerAbort); + }, + }; +} + +/** + * The error a request deadline produces. + * + * `status: 504` is load-bearing, not decoration: `isTransientPollingStatus` + * treats it as retryable, so one stalled poll inside a wait loop is abandoned + * and retried against the remaining budget instead of failing the whole wait. + */ +function requestTimeoutError(method: string, path: string, ms: number): Error { + const err = new Error( + `Walrus Memory request timed out after ${ms}ms (${method} ${path}). The relayer ` + + `accepted the connection but did not answer in time.`, + ); + err.name = "MemWalRequestTimeout"; + (err as Error & { status?: number }).status = 504; + return err; +} + function isTransientPollingStatus(status: number): boolean { return status === 0 || status === 429 || status >= 500; } +/** + * How long a transient polling rejection asked us to wait, in ms. + * + * A 429 from the relayer's rate limiter is not an invitation to poll again in + * a second. The limiter counts our own status reads, so answering a + * `retry_after_seconds: 60` with the 10s backoff cap re-trips it on every + * attempt and the loop starves itself: the job then reads as "still uploading" + * for as long as the caller is willing to wait — including long after it has + * actually failed, which is the state that costs a user their fact. + * + * Clamped to what is left of the caller's own budget, so a 10-minute backoff + * costs at most the wait the caller already asked for. The loop then spends + * that remainder on one final read at the boundary rather than on a backoff + * curve nobody asked for. + */ +function retryAfterDelayMs(err: unknown, deadline: number): number { + const seconds = + (err as { retryAfterSeconds?: number }).retryAfterSeconds + ?? retryAfterSecondsFromBody((err as { cause?: unknown }).cause); + if (seconds === undefined || !Number.isFinite(seconds) || seconds <= 0) return 0; + return Math.max(0, Math.min(seconds * 1000, deadline - Date.now())); +} + +/** + * The relayer states the backoff in the 429 body as well as in `Retry-After`. + * A proxy that strips the header must not cost us the hint. + */ +function retryAfterSecondsFromBody(cause: unknown): number | undefined { + if (typeof cause !== "string") return undefined; + try { + const parsed = JSON.parse(cause) as { retry_after_seconds?: unknown }; + const seconds = Number(parsed?.retry_after_seconds); + return Number.isFinite(seconds) ? seconds : undefined; + } catch { + return undefined; + } +} + /** * Normalise the legacy `(text, namespace)` and new `(text, options)` * overloads of `analyze()` / `analyzeAndWait()` into a single @@ -182,9 +329,12 @@ function normalizeSuiNetworkForGrpc(network: string): string { export class MemWal { private privateKey: Uint8Array; private publicKey: Uint8Array | null = null; + private destroyed = false; private serverUrl: string; private namespace: string; private accountId: string; + /** Deadline applied to any request that does not name its own. */ + private requestTimeoutMs: number; // ENG-1697 state — all internal, never surfaced to user code. // The public API (`MemWal.create({ key, accountId })`) is unchanged. @@ -219,6 +369,12 @@ export class MemWal { // non-localhost host. this.serverUrl = normalizeServerUrl(config.serverUrl ?? "https://relayer.memory.walrus.xyz"); this.namespace = config.namespace ?? "default"; + // A non-positive or non-finite override would disable the backstop + // entirely, which is the bug this exists to prevent. + this.requestTimeoutMs = + Number.isFinite(config.requestTimeoutMs) && (config.requestTimeoutMs as number) > 0 + ? (config.requestTimeoutMs as number) + : DEFAULT_REQUEST_TIMEOUT_MS; } /** @@ -236,6 +392,7 @@ export class MemWal { * Prevents key extraction from V8 heap dumps. */ destroy(): void { + this.destroyed = true; if (this.privateKey) { this.privateKey.fill(0); } @@ -268,7 +425,7 @@ export class MemWal { const generatedKey = options.idempotencyKey === undefined; const idempotencyKey = options.idempotencyKey ?? this.pendingRememberKeys.get(requestIdentity) - ?? crypto.randomUUID(); + ?? (await derivedIdempotencyKey(requestIdentity)); if (generatedKey) this.pendingRememberKeys.set(requestIdentity, idempotencyKey); const accepted = await this.signedRequest( @@ -319,9 +476,13 @@ export class MemWal { const { pollIntervalMs = 1500, timeoutMs = 60_000 } = opts; const deadline = Date.now() + timeoutMs; let attempt = 0; + let retryAfterMs = 0; while (Date.now() < deadline) { - await sleep(pollingDelayMs(pollIntervalMs, attempt++)); + // A retry-after the server just gave us wins over our own curve; + // the backoff resumes from where it was on the next normal poll. + await sleep(retryAfterMs > 0 ? retryAfterMs : pollingDelayMs(pollIntervalMs, attempt++)); + retryAfterMs = 0; let status: RememberStatusResponse; @@ -331,10 +492,22 @@ export class MemWal { `/api/remember/${jobId}`, {}, [200, 404], + // Bound each poll by the client deadline, and never past + // what is left of the budget. Without this the loop only + // checks the deadline between polls, so one stalled read + // runs past `timeoutMs` however small it was — the whole + // reason a 90s wait was seen still going at 120s. Taking + // the min of the two matters in both directions: the + // remaining budget keeps a poll from outliving the wait, + // and the client deadline keeps ONE stalled poll from + // swallowing the entire budget, so the loop still gets to + // retry. An expired poll surfaces as a transient 504. + { timeoutMs: this.pollDeadlineMs(deadline) }, ); } catch (err) { const httpStatus = (err as { status?: number }).status ?? 0; if (isTransientPollingStatus(httpStatus)) { + retryAfterMs = retryAfterDelayMs(err, deadline); continue; } throw err; @@ -385,7 +558,7 @@ export class MemWal { const generatedKey = opts.idempotencyKey === undefined; const idempotencyKey = opts.idempotencyKey ?? this.pendingRememberKeys.get(requestIdentity) - ?? crypto.randomUUID(); + ?? (await derivedIdempotencyKey(requestIdentity)); if (generatedKey) this.pendingRememberKeys.set(requestIdentity, idempotencyKey); const accepted = await this.rememberAsync(text, resolvedNamespace, { idempotencyKey }); @@ -460,11 +633,16 @@ export class MemWal { return accepted; } - async getRememberBulkStatus(jobIds: string[]): Promise { + async getRememberBulkStatus( + jobIds: string[], + opts: { timeoutMs?: number } = {}, + ): Promise { return this.signedRequest( "POST", "/api/remember/bulk/status", { job_ids: jobIds }, + [200], + { timeoutMs: opts.timeoutMs }, ); } @@ -484,9 +662,13 @@ export class MemWal { })); const pending = new Set(jobIds); let attempt = 0; + let retryAfterMs = 0; while (pending.size > 0 && Date.now() < deadline) { - await sleep(pollingDelayMs(pollIntervalMs, attempt++)); + // A retry-after the server just gave us wins over our own curve; + // the backoff resumes from where it was on the next normal poll. + await sleep(retryAfterMs > 0 ? retryAfterMs : pollingDelayMs(pollIntervalMs, attempt++)); + retryAfterMs = 0; const pendingIds = jobIds.filter((jobId) => pending.has(jobId)); if (pendingIds.length === 0) { @@ -496,10 +678,14 @@ export class MemWal { let batchStatus: RememberBulkStatusResult; try { - batchStatus = await this.getRememberBulkStatus(pendingIds); + // Bounded the same way as the single-job poll. + batchStatus = await this.getRememberBulkStatus(pendingIds, { + timeoutMs: this.pollDeadlineMs(deadline), + }); } catch (err) { const httpStatus = (err as { status?: number }).status ?? 0; if (isTransientPollingStatus(httpStatus)) { + retryAfterMs = retryAfterDelayMs(err, deadline); continue; } throw err; @@ -666,9 +852,7 @@ export class MemWal { const limit = options.topK ?? options.limit ?? 10; const resolvedNamespace = options.namespace ?? this.namespace; - const ac = new AbortController(); - const tid = setTimeout(() => ac.abort(), 15000); - try { + { const result = await this.signedRequest("POST", "/api/recall", { query, limit, @@ -682,7 +866,10 @@ export class MemWal { // request byte-identical and the relayer applies its own // "relevance" default. sort: options.sort, - }, { signal: ac.signal }); + // How long this call waits, so the relayer can stop just + // short of it and name the step it was stuck in. + deadline_ms: RECALL_DEADLINE_MS, + }, { timeoutMs: RECALL_REQUEST_TIMEOUT_MS }); let processed = result; if (typeof options.maxDistance === "number") { @@ -716,8 +903,6 @@ export class MemWal { } return processed; - } finally { - clearTimeout(tid); } } @@ -843,7 +1028,9 @@ export class MemWal { }; const wireOccurredAt = occurredAtToWire(options.occurredAt); if (wireOccurredAt !== undefined) body.occurred_at = wireOccurredAt; - return this.signedRequest("POST", "/api/analyze", body, [200, 202]); + return this.signedRequest("POST", "/api/analyze", body, [200, 202], { + timeoutMs: ANALYZE_REQUEST_TIMEOUT_MS, + }); } /** @@ -912,10 +1099,13 @@ export class MemWal { * ``` */ async restore(namespace: string, limit: number = 10): Promise { - const result = await this.signedRequest("POST", "/api/restore", { - namespace, - limit, - }); + const result = await this.signedRequest( + "POST", + "/api/restore", + { namespace, limit }, + [200], + { timeoutMs: RESTORE_REQUEST_TIMEOUT_MS }, + ); // Relayers older than WALM-319 omit `truncated` entirely — treat // "not present" as "not known to be truncated" rather than drop // the field or require every relayer version to send it. @@ -1015,11 +1205,11 @@ export class MemWal { * Check server health. The endpoint is public and does not require request signing. */ async health(): Promise { - const res = await fetch(`${this.serverUrl}/health`); + const res = await this.fetchWithDeadline(`${this.serverUrl}/health`); if (!res.ok) { throw new Error(`Health check failed: ${res.status}`); } - return res.json() as Promise; + return res.body as HealthResult; } /** @@ -1060,20 +1250,26 @@ export class MemWal { } private async fetchCompatibilityMetadata(): Promise { - const versionRes = await fetch(`${this.serverUrl}/version`, { method: "GET" }); + const versionRes = await this.fetchWithDeadline>( + `${this.serverUrl}/version`, + { method: "GET" }, + ); let body: Partial; if (versionRes.ok) { - body = (await versionRes.json()) as Partial; + body = versionRes.body as Partial; } else if (versionRes.status === 404 || versionRes.status === 405) { - const healthRes = await fetch(`${this.serverUrl}/health`, { method: "GET" }); + const healthRes = await this.fetchWithDeadline>( + `${this.serverUrl}/health`, + { method: "GET" }, + ); if (!healthRes.ok) { throw new Error( `Walrus Memory compatibility check failed: GET /version returned ` + `${versionRes.status}, and GET /health returned ${healthRes.status}`, ); } - body = (await healthRes.json()) as Partial; + body = healthRes.body as Partial; } else { throw new Error( `Walrus Memory compatibility check failed: GET /version returned ${versionRes.status}`, @@ -1109,11 +1305,14 @@ export class MemWal { private async fetchServerConfig(): Promise { if (this.serverConfig) return this.serverConfig; - const res = await fetch(`${this.serverUrl}/config`, { method: "GET" }); + const res = await this.fetchWithDeadline>( + `${this.serverUrl}/config`, + { method: "GET" }, + ); if (!res.ok) { throw new Error(`GET /config returned ${res.status}`); } - const body = (await res.json()) as Record; + const body = res.body as Record; if (typeof body.packageId !== "string" || !body.packageId || typeof body.network !== "string" || !body.network) { throw new Error("GET /config response missing packageId / network"); @@ -1322,12 +1521,50 @@ export class MemWal { * @param acceptedStatuses - HTTP status codes to treat as success (default [200]). * Pass [200, 202] for endpoints that return 202 Accepted. */ + /** + * `fetch` plus the JSON body read, with this client's deadline applied across both. + * + * For the handshake endpoints that skip request signing — `/health`, + * `/version`, `/config`. They are the worst place to leave unbounded: the + * compatibility probe and config fetch run before the first real call, so + * one stalled socket there hangs every method on the client, not just one. + */ + /** Deadline for one status poll: the client deadline, clamped so it can + * never outlive the wait budget it belongs to. */ + private pollDeadlineMs(deadline: number): number { + return Math.max(1, Math.min(this.requestTimeoutMs, deadline - Date.now())); + } + + private async fetchWithDeadline( + url: string, + init: RequestInit = {}, + timeoutMs?: number, + ): Promise<{ ok: boolean; status: number; body: T | undefined }> { + const ms = timeoutMs ?? this.requestTimeoutMs; + const deadline = deadlineSignal(ms); + try { + const res = await fetch(url, { ...init, signal: deadline.signal }); + return { + ok: res.ok, + status: res.status, + body: res.ok ? ((await res.json()) as T) : undefined, + }; + } catch (err) { + if (deadline.timedOut()) { + throw requestTimeoutError(init.method ?? "GET", url, ms); + } + throw err; + } finally { + deadline.dispose(); + } + } + private async signedRequest( method: string, path: string, body: object, - acceptedStatusesOrOptions: number[] | { includeDelegateKey?: boolean; signal?: AbortSignal } = [200], - requestOptions: { includeDelegateKey?: boolean; signal?: AbortSignal } = {}, + acceptedStatusesOrOptions: number[] | SignedRequestOptions = [200], + requestOptions: SignedRequestOptions = {}, ): Promise { const acceptedStatuses = Array.isArray(acceptedStatusesOrOptions) ? acceptedStatusesOrOptions @@ -1353,6 +1590,11 @@ export class MemWal { const message = `${timestamp}.${method}.${path}.${bodySha256}.${nonce}.${this.accountId}`; const msgBytes = new TextEncoder().encode(message); + if (this.destroyed) { + throw new Error( + "Walrus Memory client was destroyed; its keys are zeroed. Create a new one.", + ); + } // Sign with Ed25519 const signature = await ed.signAsync(msgBytes, this.privateKey); const publicKey = await this.getPublicKey(); @@ -1374,47 +1616,63 @@ export class MemWal { if (options.includeDelegateKey !== false) { headers["x-seal-session"] = await this.buildSealSession(); } - const res = await fetch(url, { - method, - headers, - body: method === "GET" ? undefined : bodyStr, - signal: options.signal, - }); + // Bound the request. `fetch` never times out on its own, so this is the + // only thing standing between a stalled socket and a call that hangs + // for as long as the connection stays open. + const deadlineMs = options.timeoutMs ?? this.requestTimeoutMs; + const deadline = deadlineSignal(deadlineMs, options.signal); + try { + const res = await fetch(url, { + method, + headers, + body: method === "GET" ? undefined : bodyStr, + signal: deadline.signal, + }); + + if (!acceptedStatuses.includes(res.status)) { + // LOW-26: sanitize server error bodies before surfacing to callers. + const raw = await res.text(); + const compatibilityError = compatibilityErrorFromStatus(res.status, raw); + if (compatibilityError) throw compatibilityError; + + // A stale/future-dated signature is rejected with 401 + a machine- + // readable reason header. Surface it as an actionable clock-drift + // error rather than an opaque 401 so the caller can fix node time. + const clockDriftError = clockDriftErrorFromResponse(res); + if (clockDriftError) throw clockDriftError; + + const { message, serverCode } = sanitizeServerError( + res.status, + raw, + res.headers.get("x-auth-error"), + ); + const err = new Error(message) as Error & { + status?: number; + serverCode?: string; + retryAfterSeconds?: number; + cause?: string; + }; + err.status = res.status; + if (serverCode) err.serverCode = serverCode; + const retryAfter = Number(res.headers.get("retry-after")); + if (Number.isFinite(retryAfter) && retryAfter > 0) { + err.retryAfterSeconds = retryAfter; + } + // Preserve raw body on `cause` for in-process debugging only. + err.cause = raw; + throw err; + } - if (!acceptedStatuses.includes(res.status)) { - // LOW-26: sanitize server error bodies before surfacing to callers. - const raw = await res.text(); - const compatibilityError = compatibilityErrorFromStatus(res.status, raw); - if (compatibilityError) throw compatibilityError; - - // A stale/future-dated signature is rejected with 401 + a machine- - // readable reason header. Surface it as an actionable clock-drift - // error rather than an opaque 401 so the caller can fix node time. - const clockDriftError = clockDriftErrorFromResponse(res); - if (clockDriftError) throw clockDriftError; - - const { message, serverCode } = sanitizeServerError( - res.status, - raw, - res.headers.get("x-auth-error"), - ); - const err = new Error(message) as Error & { - status?: number; - serverCode?: string; - retryAfterSeconds?: number; - cause?: string; - }; - err.status = res.status; - if (serverCode) err.serverCode = serverCode; - const retryAfter = Number(res.headers.get("retry-after")); - if (Number.isFinite(retryAfter) && retryAfter > 0) { - err.retryAfterSeconds = retryAfter; + return (await res.json()) as T; + } catch (err) { + // Translate our own expiry into something a caller can classify. + // An abort the CALLER asked for is theirs and propagates untouched. + if (deadline.timedOut() && (err as { status?: number }).status === undefined) { + throw requestTimeoutError(method, path, deadlineMs); } - // Preserve raw body on `cause` for in-process debugging only. - err.cause = raw; throw err; + } finally { + deadline.dispose(); } - - return res.json() as Promise; } } diff --git a/packages/sdk/src/types.ts b/packages/sdk/src/types.ts index e6b374c50..e253c1090 100644 --- a/packages/sdk/src/types.ts +++ b/packages/sdk/src/types.ts @@ -21,6 +21,18 @@ export interface MemWalConfig { serverUrl?: string; /** Default namespace for memory isolation (default: "default") */ namespace?: string; + /** + * Deadline for a single relayer request, in milliseconds (default: 30000). + * + * `fetch` has no timeout of its own, so without this a stalled connection + * keeps a call pending for as long as the socket stays open — which is how + * a `remember` whose poll budget was 90s could still be running after two + * minutes. Endpoints that are legitimately slower (`restore`, `analyze`) + * carry their own larger deadline and ignore this. + * + * Raise it only for a genuinely slow link; it is a backstop, not a budget. + */ + requestTimeoutMs?: number; } // ============================================================ @@ -165,6 +177,9 @@ export interface RecallOptions { * * For newest-wins, use `sort: "recent"` instead. It over-fetches * candidates server-side before ordering them by write-time. + * + * Ignored when `sort` is set: an explicit `sort`, `"relevance"` included, + * decides the order, and weights apply only when `sort` is omitted. */ scoringWeights?: ScoringWeights; /** @@ -179,6 +194,8 @@ export interface RecallOptions { * widens the candidate set; weights only re-rank the set that was already * returned, so weights alone cannot surface a record that fell outside * the window. + * + * Setting `sort` at all makes the relayer ignore `scoringWeights`. */ sort?: "relevance" | "recent"; } diff --git a/packages/sdk/src/utils.ts b/packages/sdk/src/utils.ts index 54c33eb4d..cbeb3cbea 100644 --- a/packages/sdk/src/utils.ts +++ b/packages/sdk/src/utils.ts @@ -294,7 +294,7 @@ export function normalizeServerUrl(url: string): string { /** Replace loopback URLs that leak sidecar topology into client-facing errors. */ export function redactInternalUrls(text: string): string { return text.replace( - /https?:\/\/(?:localhost|127\.0\.0\.1)(?::\d+)?[^ \t)\]>'"]*/gi, + /https?:\/\/(?:(?:[\w-]+\.)*localhost|127\.\d{1,3}\.\d{1,3}\.\d{1,3}|\[::1\]|0\.0\.0\.0)(?::\d+)?[^ \t)\]>'"]*/gi, "[internal]", ); } diff --git a/packages/sdk/test/destroy-invalidates-client.test.mjs b/packages/sdk/test/destroy-invalidates-client.test.mjs new file mode 100644 index 000000000..3408a764e --- /dev/null +++ b/packages/sdk/test/destroy-invalidates-client.test.mjs @@ -0,0 +1,43 @@ +import assert from "node:assert/strict"; +import test from "node:test"; + +import { MemWal } from "../dist/memwal.js"; + +test("a request in flight when destroy() lands is rejected, not signed with a zeroed key", async () => { + const originalFetch = globalThis.fetch; + let releaseVersion; + const versionGate = new Promise((resolve) => { + releaseVersion = resolve; + }); + + globalThis.fetch = async (url) => { + const path = new URL(url).pathname; + if (path === "/version") { + await versionGate; + return Response.json({ + apiVersion: "1.0.0", + relayerVersion: "1.0.0", + minSupportedSdk: { typescript: "0.0.4" }, + }); + } + throw new Error(`a wiped client still signed and sent ${path}`); + }; + + try { + const client = MemWal.create({ + key: new Uint8Array(32).fill(1), + accountId: "0x1", + serverUrl: "https://relayer.example", + }); + client.buildSealSession = async () => "test-session"; + + const pending = client.recall({ query: "anything" }); + await new Promise((resolve) => setImmediate(resolve)); + client.destroy(); + releaseVersion(); + + await assert.rejects(pending, /Walrus Memory client was destroyed/); + } finally { + globalThis.fetch = originalFetch; + } +}); diff --git a/packages/sdk/test/recall-deadline.test.mjs b/packages/sdk/test/recall-deadline.test.mjs new file mode 100644 index 000000000..39fdc2d04 --- /dev/null +++ b/packages/sdk/test/recall-deadline.test.mjs @@ -0,0 +1,80 @@ +import assert from "node:assert/strict"; +import test from "node:test"; + +import { MemWal } from "../dist/memwal.js"; + +// A recall that runs out of time used to abort here with nothing to say why: +// "This operation was aborted", or later "request timed out after 15000ms". +// Only the relayer knows which step was stuck, and it can only answer before +// the caller gives up if it knows when that is. So recall sends its deadline, +// and the relayer answers a 504 naming the stage just short of it (WALM-396). + +const originalFetch = globalThis.fetch; + +test.afterEach(() => { + globalThis.fetch = originalFetch; +}); + +function stubRecall(answer) { + const sent = {}; + globalThis.fetch = async (url, init = {}) => { + const path = new URL(url).pathname; + if (path === "/version") { + return Response.json({ + apiVersion: "1.0.0", + relayerVersion: "1.0.0", + minSupportedSdk: { typescript: "0.0.4" }, + }); + } + if (path === "/api/config") { + return Response.json({ packageId: "0x1", network: "testnet" }); + } + if (path === "/api/recall" && init.method === "POST") { + sent.body = JSON.parse(init.body); + return answer(); + } + throw new Error(`unexpected request ${path}`); + }; + return sent; +} + +function client() { + const c = MemWal.create({ + key: new Uint8Array(32).fill(1), + accountId: "0x1", + serverUrl: "https://relayer.example", + }); + c.buildSealSession = async () => "test-session"; + return c; +} + +test("recall tells the relayer how long it will wait", async () => { + const sent = stubRecall(() => Response.json({ results: [], total: 0 })); + + await client().recall({ query: "current task" }); + + // A second under the 15s the request is aborted at. The relayer counts its + // own margin from arrival, so without this gap a slow connect would leave + // the 504 landing after the abort it is meant to beat. + assert.equal(sent.body.deadline_ms, 14_000); +}); + +test("a relayer recall timeout reaches the caller with its code and stage", async () => { + const body = { + error: "Recall timed out after 14001ms during walrus_download", + message: "Recall timed out after 14001ms during walrus_download", + code: "RECALL_TIMEOUT", + stage: "walrus_download", + elapsed_ms: 14001, + }; + stubRecall(() => Response.json(body, { status: 504 })); + + await assert.rejects(client().recall({ query: "current task" }), (err) => { + assert.equal(err.status, 504); + assert.equal(err.serverCode, "RECALL_TIMEOUT"); + assert.match(err.message, /walrus_download/); + // The MCP sidecar reads the stage off the raw body. + assert.equal(JSON.parse(err.cause).stage, "walrus_download"); + return true; + }); +}); diff --git a/packages/sdk/test/remember-derived-key.test.mjs b/packages/sdk/test/remember-derived-key.test.mjs new file mode 100644 index 000000000..41aaa5eda --- /dev/null +++ b/packages/sdk/test/remember-derived-key.test.mjs @@ -0,0 +1,100 @@ +/** + * Idempotency keys for `remember`, derived from content rather than random. + * + * Poll cadence is deliberately NOT tested here — WALM-623 (#902) owns the + * backoff, including the immediate first attempt, and pins it in + * test/polling-delay.test.mjs. Asserting it from two places would leave one + * copy silently wrong the next time the cap moves. + */ +import assert from "node:assert/strict"; +import test from "node:test"; + +import { MemWal } from "../dist/memwal.js"; + +const originalFetch = globalThis.fetch; + +test.afterEach(() => { + globalThis.fetch = originalFetch; +}); + +/** Stub relayer. `onStatus(pollIndex)` decides what each status poll returns. */ +function stubRelayer({ onStatus, posted = [] }) { + let polls = 0; + globalThis.fetch = async (url, init = {}) => { + const path = new URL(url).pathname; + if (path === "/version") { + return Response.json({ + apiVersion: "1.0.0", + relayerVersion: "1.0.0", + minSupportedSdk: { typescript: "0.0.4" }, + }); + } + if (path === "/api/config") { + return Response.json({ packageId: "0x1", network: "testnet" }); + } + if (path === "/api/remember" && init.method === "POST") { + posted.push({ body: JSON.parse(init.body), at: Date.now() }); + return Response.json({ job_id: "job-1", status: "pending" }, { status: 202 }); + } + if (path === "/api/remember/job-1") { + return Response.json(onStatus(polls++)); + } + throw new Error(`unexpected request ${path}`); + }; + return { posted, pollCount: () => polls }; +} + +function newClient() { + const client = MemWal.create({ + key: new Uint8Array(32).fill(1), + accountId: "0x1", + serverUrl: "https://relayer.example", + }); + client.buildSealSession = async () => "test-session"; + return client; +} + +const DONE = { + job_id: "job-1", + status: "done", + blob_id: "blob-1", + owner: "0x1", + namespace: "default", +}; + +test("the same fact reuses one idempotency key across client instances", async () => { + const posted = []; + stubRelayer({ onStatus: () => DONE, posted }); + + // Two instances = what the MCP sidecar builds when the stdio bridge + // reconnects a dropped SSE stream. A per-instance random key made the + // replay look like a new write, so the relayer minted a second paid + // Walrus blob for a write already in flight. + await newClient().rememberAndWait("a durable fact"); + await newClient().rememberAndWait("a durable fact"); + + assert.equal(posted.length, 2); + assert.equal(posted[0].body.idempotency_key, posted[1].body.idempotency_key); +}); + +test("different text and different namespaces get different keys", async () => { + const posted = []; + stubRelayer({ onStatus: () => DONE, posted }); + + const client = newClient(); + await client.rememberAndWait("fact one"); + await client.rememberAndWait("fact two"); + await client.rememberAndWait("fact one", "other-namespace"); + + const keys = posted.map((p) => p.body.idempotency_key); + assert.equal(new Set(keys).size, 3, "distinct writes must not collapse onto one job"); +}); + +test("an explicit idempotency key still wins", async () => { + const posted = []; + stubRelayer({ onStatus: () => DONE, posted }); + + await newClient().rememberAndWait("a fact", undefined, { idempotencyKey: "caller-owned" }); + + assert.equal(posted[0].body.idempotency_key, "caller-owned"); +}); diff --git a/packages/sdk/test/remember-retry-after.test.mjs b/packages/sdk/test/remember-retry-after.test.mjs new file mode 100644 index 000000000..8cc1342b7 --- /dev/null +++ b/packages/sdk/test/remember-retry-after.test.mjs @@ -0,0 +1,149 @@ +import assert from "node:assert/strict"; +import test from "node:test"; + +import { MemWal } from "../dist/memwal.js"; + +const originalFetch = globalThis.fetch; + +test.afterEach(() => { + globalThis.fetch = originalFetch; +}); + +/** Stub a job whose first status read is rate-limited, and record when each + * read arrived so the test can measure the gap the SDK actually waited. */ +function rateLimitedJob({ header, body }) { + const polledAt = []; + globalThis.fetch = async (url, init = {}) => { + const path = new URL(url).pathname; + if (path === "/version") { + return Response.json({ + apiVersion: "1.0.0", + relayerVersion: "1.0.0", + minSupportedSdk: { typescript: "0.0.4" }, + }); + } + if (path === "/api/config") { + return Response.json({ packageId: "0x1", network: "testnet" }); + } + if (path === "/api/remember" && init.method === "POST") { + return Response.json({ job_id: "limited-job", status: "pending" }, { status: 202 }); + } + if (path === "/api/remember/limited-job") { + polledAt.push(Date.now()); + if (polledAt.length === 1) { + return new Response(JSON.stringify(body), { + status: 429, + headers: { + "content-type": "application/json", + ...(header ? { "retry-after": header } : {}), + }, + }); + } + return Response.json({ + job_id: "limited-job", + status: "done", + blob_id: "blob-1", + owner: "0x1", + namespace: "default", + }); + } + throw new Error(`unexpected request ${path}`); + }; + return polledAt; +} + +function client() { + const memwal = MemWal.create({ + key: new Uint8Array(32).fill(1), + accountId: "0x1", + serverUrl: "https://relayer.example", + }); + memwal.buildSealSession = async () => "test-session"; + return memwal; +} + +// The rate limiter counts our own status reads, so re-polling on the 1.5s→10s +// curve after a `retry_after_seconds: 60` re-trips it every attempt and the +// loop starves itself — the job reads as "still uploading" for as long as the +// caller waits, including long after it has failed. `pollIntervalMs: 1` makes +// the unfixed behaviour ~1ms, so the wait being ≥ the stated backoff is +// unambiguous. +test("a rate-limited status poll waits the retry-after from the body", async () => { + const polledAt = rateLimitedJob({ + body: { error: "Rate limit exceeded", retry_after_seconds: 0.5 }, + }); + + const result = await client().rememberAndWait("fact", undefined, { + pollIntervalMs: 1, + timeoutMs: 10_000, + }); + + assert.equal(result.blob_id, "blob-1"); + assert.equal(polledAt.length, 2); + assert.ok( + polledAt[1] - polledAt[0] >= 450, + `expected to honour the 500ms backoff, waited ${polledAt[1] - polledAt[0]}ms`, + ); +}); + +test("a rate-limited status poll waits the Retry-After header", async () => { + const polledAt = rateLimitedJob({ + header: "1", + body: { error: "Rate limit exceeded" }, + }); + + const result = await client().rememberAndWait("fact", undefined, { + pollIntervalMs: 1, + timeoutMs: 10_000, + }); + + assert.equal(result.blob_id, "blob-1"); + assert.ok( + polledAt[1] - polledAt[0] >= 950, + `expected to honour the 1s backoff, waited ${polledAt[1] - polledAt[0]}ms`, + ); +}); + +test("a retry-after longer than the remaining budget is clamped to it", async () => { + // Honouring a 10-minute backoff must not turn a 1s wait into a 10-minute + // one. The budget still buys a final read at its own boundary — that read + // is free and may be the answer — but nothing beyond it. + const polledAt = []; + globalThis.fetch = async (url, init = {}) => { + const path = new URL(url).pathname; + if (path === "/version") { + return Response.json({ + apiVersion: "1.0.0", + relayerVersion: "1.0.0", + minSupportedSdk: { typescript: "0.0.4" }, + }); + } + if (path === "/api/config") { + return Response.json({ packageId: "0x1", network: "testnet" }); + } + if (path === "/api/remember" && init.method === "POST") { + return Response.json({ job_id: "slow-job", status: "pending" }, { status: 202 }); + } + if (path === "/api/remember/slow-job") { + polledAt.push(Date.now()); + if (polledAt.length === 1) { + return new Response( + JSON.stringify({ error: "Rate limit exceeded", retry_after_seconds: 600 }), + { status: 429, headers: { "content-type": "application/json" } }, + ); + } + return Response.json({ job_id: "slow-job", status: "pending" }); + } + throw new Error(`unexpected request ${path}`); + }; + + const startedAt = Date.now(); + await assert.rejects( + client().rememberAndWait("fact", undefined, { pollIntervalMs: 1, timeoutMs: 800 }), + /timed out/, + ); + assert.ok( + Date.now() - startedAt < 3_000, + `a 600s retry-after must be clamped to the caller's own timeout, took ${Date.now() - startedAt}ms`, + ); +}); diff --git a/packages/sdk/test/request-timeout.test.mjs b/packages/sdk/test/request-timeout.test.mjs new file mode 100644 index 000000000..2220368ce --- /dev/null +++ b/packages/sdk/test/request-timeout.test.mjs @@ -0,0 +1,226 @@ +import assert from "node:assert/strict"; +import test from "node:test"; + +import { MemWal } from "../dist/memwal.js"; + +const originalFetch = globalThis.fetch; + +test.afterEach(() => { + globalThis.fetch = originalFetch; +}); + +/** + * `fetch` has no timeout of its own, and the SDK used to pass a signal on + * exactly one method (`recall`). Everything else — the accept POST, every job + * status read, and the `/version` and `/config` handshake calls that run before + * any of them — could stay pending for as long as the socket stayed open. + * + * A poll loop only checks its budget between polls, so that was not merely + * untidy: one stalled read ran straight past `timeoutMs`, which is how a + * `memwal_remember` documented as capping at 90s was seen still running after + * 120s by an MCP client. + */ + +/** Stands in for a stalled socket: never answers, but honours abort the way a + * real `fetch` does — which is also what proves the signal reaches it. */ +function hangUntilAborted(init = {}) { + return new Promise((_, reject) => { + const signal = init.signal; + if (!signal) return; // no signal reaching fetch => hangs forever => test times out + if (signal.aborted) return reject(abortError()); + signal.addEventListener("abort", () => reject(abortError()), { once: true }); + }); +} + +function abortError() { + const err = new Error("This operation was aborted"); + err.name = "AbortError"; + return err; +} + +/** Stub relayer whose handshake always answers; `onApi` decides the rest. */ +function stubRelayer(onApi) { + globalThis.fetch = async (url, init = {}) => { + const path = new URL(url).pathname; + if (path === "/version") { + return Response.json({ + apiVersion: "1.0.0", + relayerVersion: "1.0.0", + minSupportedSdk: { typescript: "0.0.4" }, + }); + } + if (path === "/config") { + return Response.json({ packageId: "0x1", network: "testnet" }); + } + return onApi(path, init, url); + }; +} + +function clientWith(extra = {}) { + const client = MemWal.create({ + key: new Uint8Array(32).fill(1), + accountId: "0x1", + serverUrl: "https://relayer.example", + ...extra, + }); + client.buildSealSession = async () => "test-session"; + return client; +} + +test("a stalled accept is bounded instead of hanging forever", async () => { + stubRelayer((path, init) => + path === "/api/remember" ? hangUntilAborted(init) : Promise.reject(new Error(path)), + ); + + const started = Date.now(); + await assert.rejects( + clientWith({ requestTimeoutMs: 120 }).rememberAsync("a durable fact"), + (err) => { + assert.equal(err.name, "MemWalRequestTimeout"); + // 504 is load-bearing: `isTransientPollingStatus` treats it as + // retryable, so a stalled poll inside a wait loop is retried + // against the remaining budget rather than failing the whole wait. + assert.equal(err.status, 504); + assert.match(err.message, /POST \/api\/remember/); + return true; + }, + ); + assert.ok(Date.now() - started < 2_000, "should give up at the deadline, not hang"); +}); + +test("a stalled handshake is bounded too", async () => { + // `/version` runs before any memory call, so leaving it unbounded hangs + // every method on the client rather than one request. + globalThis.fetch = async (url, init = {}) => hangUntilAborted(init); + + await assert.rejects( + clientWith({ requestTimeoutMs: 120 }).rememberAsync("a durable fact"), + (err) => { + assert.equal(err.name, "MemWalRequestTimeout"); + return true; + }, + ); +}); + +test("waitForRememberJob stays inside its budget when every poll stalls", async () => { + stubRelayer((path, init) => + path.startsWith("/api/remember/") ? hangUntilAborted(init) : Promise.reject(new Error(path)), + ); + + const started = Date.now(); + await assert.rejects( + clientWith({ requestTimeoutMs: 10_000 }).waitForRememberJob("job-1", { + timeoutMs: 400, + pollIntervalMs: 50, + }), + /timed out/, + ); + // The budget is the bound now. Previously the loop only checked it between + // polls, so a single stalled read outlived it by however long the socket + // stayed open — here that would have been the 10s client deadline. + assert.ok( + Date.now() - started < 3_000, + `wait overran its budget: ${Date.now() - started}ms`, + ); +}); + +test("an expired poll is retried rather than failing the whole wait", async () => { + let call = 0; + stubRelayer((path, init) => { + if (!path.startsWith("/api/remember/")) return Promise.reject(new Error(path)); + // First poll stalls; the retry answers. + if (call++ === 0) return hangUntilAborted(init); + return Response.json({ + job_id: "job-1", + status: "done", + blob_id: "blob-1", + owner: "0x1", + namespace: "default", + }); + }); + + const result = await clientWith({ requestTimeoutMs: 100 }).waitForRememberJob("job-1", { + timeoutMs: 5_000, + pollIntervalMs: 20, + }); + + assert.equal(result.blob_id, "blob-1"); + assert.ok(call >= 2, "the stalled poll should have been retried, not fatal"); +}); + +test("a transport failure that is not a deadline keeps its own identity", async () => { + // The deadline path must not swallow real network errors — an operator + // debugging DNS or TLS needs the original, not "timed out". + stubRelayer(() => Promise.reject(new TypeError("fetch failed"))); + + await assert.rejects( + clientWith({ requestTimeoutMs: 5_000 }).rememberAsync("a durable fact"), + (err) => { + assert.notEqual(err.name, "MemWalRequestTimeout"); + assert.match(err.message, /fetch failed/); + return true; + }, + ); +}); + +test("an unusable requestTimeoutMs falls back to the default rather than disabling the bound", async () => { + // 0 / negative / NaN would mean "no deadline", which is the bug this + // exists to prevent — a typo must not silently restore it. + for (const bad of [0, -1, Number.NaN, undefined]) { + const client = clientWith({ requestTimeoutMs: bad }); + assert.equal(client.requestTimeoutMs, 30_000, `bad value ${bad} disabled the bound`); + } + assert.equal(clientWith({ requestTimeoutMs: 1234 }).requestTimeoutMs, 1234); +}); + +test("a body that stalls after the headers is bounded too", { timeout: 10_000 }, async () => { + stubRelayer((path, init) => { + if (path !== "/api/remember") return Promise.reject(new Error(path)); + return Promise.resolve({ + ok: true, + status: 202, + headers: new Headers(), + json: () => hangUntilAborted(init), + text: () => hangUntilAborted(init), + }); + }); + + const started = Date.now(); + await assert.rejects( + clientWith({ requestTimeoutMs: 120 }).rememberAsync("a durable fact"), + (err) => { + assert.equal(err.name, "MemWalRequestTimeout"); + assert.equal(err.status, 504); + assert.match(err.message, /POST \/api\/remember/); + return true; + }, + ); + assert.ok( + Date.now() - started < 2_000, + "the deadline should still be armed while the body is read", + ); +}); + +test("a classified error survives a deadline that expires during its body read", async () => { + const raw = JSON.stringify({ error: "Rate limit exceeded", retry_after_seconds: 60 }); + stubRelayer((path) => { + if (path !== "/api/remember") return Promise.reject(new Error(path)); + return Promise.resolve({ + ok: false, + status: 429, + headers: new Headers({ "content-type": "application/json", "retry-after": "60" }), + text: () => new Promise((resolve) => setTimeout(() => resolve(raw), 300)), + json: () => new Promise((resolve) => setTimeout(() => resolve(JSON.parse(raw)), 300)), + }); + }); + + await assert.rejects( + clientWith({ requestTimeoutMs: 120 }).rememberAsync("a durable fact"), + (err) => { + assert.equal(err.status, 429); + assert.equal(err.retryAfterSeconds, 60); + assert.equal(err.cause, raw); + return true; + }, + ); +}); diff --git a/packages/sdk/test/sanitize-server-error.test.mjs b/packages/sdk/test/sanitize-server-error.test.mjs index 623afdad4..341d058f3 100644 --- a/packages/sdk/test/sanitize-server-error.test.mjs +++ b/packages/sdk/test/sanitize-server-error.test.mjs @@ -74,3 +74,12 @@ test("localhost sidecar URLs are stripped from error text", () => { assert.doesNotMatch(message, /localhost:9000/); assert.match(message, /\[internal\]/); }); + +test("IPv6 loopback sidecar URLs are stripped from error text", () => { + const { message } = sanitizeServerError( + 500, + "Sidecar seal/encrypt request failed: error sending request for url (http://[::1]:9000/seal/encrypt)", + ); + assert.doesNotMatch(message, /::1/); + assert.match(message, /\[internal\]/); +}); diff --git a/packages/sdk/test/untrusted-memory.test.mjs b/packages/sdk/test/untrusted-memory.test.mjs index c01305bc9..0884657cc 100644 --- a/packages/sdk/test/untrusted-memory.test.mjs +++ b/packages/sdk/test/untrusted-memory.test.mjs @@ -50,3 +50,26 @@ test("recalled bytes are never inserted into a system message", () => { assert.equal(userMessages.length, 2); assert.equal(userMessages[0].content[0].text, attack); }); + +test("the trust policy stays ahead of the first non-system message", () => { + const memory = "recalled fact"; + const prompt = [ + { role: "system", content: "be helpful" }, + { role: "user", content: [{ type: "text", text: "first" }] }, + { role: "assistant", content: [{ type: "text", text: "reply" }] }, + { role: "user", content: [{ type: "text", text: "second" }] }, + ]; + const enriched = injectMemoryContext(prompt, memory); + + const firstNonSystem = enriched.findIndex((message) => message.role !== "system"); + assert.equal( + enriched.slice(firstNonSystem).some((message) => message.role === "system"), + false + ); + assert.ok( + enriched.some( + (message) => + message.role === "user" && message.content[0]?.text === memory + ) + ); +}); diff --git a/packages/sdk/test/withMemWal-flush.test.mjs b/packages/sdk/test/withMemWal-flush.test.mjs index 4449a8764..114817f00 100644 --- a/packages/sdk/test/withMemWal-flush.test.mjs +++ b/packages/sdk/test/withMemWal-flush.test.mjs @@ -1,6 +1,8 @@ import assert from "node:assert/strict"; import test from "node:test"; +import { generateText } from "ai"; + import { MemWal } from "../dist/memwal.js"; import { withMemWal } from "../dist/ai/middleware.js"; @@ -62,3 +64,25 @@ test("flush() awaits the fire-and-forget auto-save analyze() call", async () => MemWal.prototype.analyze = originalAnalyze; } }); + +test("a wrapped v2 model still reports its finish reason", async () => { + const originalRecall = MemWal.prototype.recall; + const originalAnalyze = MemWal.prototype.analyze; + + MemWal.prototype.recall = async () => ({ results: [] }); + MemWal.prototype.analyze = async () => { }; + + try { + const model = withMemWal(fakeLanguageModel(), { + key: TEST_KEY, + accountId: TEST_ACCOUNT_ID, + }); + + const result = await generateText({ model, prompt: "hi" }); + + assert.equal(result.finishReason, "stop"); + } finally { + MemWal.prototype.recall = originalRecall; + MemWal.prototype.analyze = originalAnalyze; + } +}); diff --git a/scripts/check-mcp-pin.mjs b/scripts/check-mcp-pin.mjs new file mode 100755 index 000000000..a2974ca67 --- /dev/null +++ b/scripts/check-mcp-pin.mjs @@ -0,0 +1,163 @@ +#!/usr/bin/env node + +/** + * The version the plugin launcher installs must exist on npm. + * + * `pinnedVersion()` in packages/mcp/plugin/scripts/lib/mcp-launch.mjs reads + * `mcpPackageVersion` out of plugin.json and installs exactly that. npm has no + * such version, the install fails, and the launch fails with it — every fresh + * plugin install, on every editor, until someone edits the manifest. That is not + * hypothetical: on 2026-09-17 the launchers pinned `@0.0.14` while npm carried + * only 0.0.13 and 0.0.14-dev.0, because plugin.json had been bumped for a release + * that had not shipped. + * + * A unit test cannot catch this — whether a version exists is a fact about npm, + * not about the repository — so it lives here. + * + * Two findings, deliberately unequal: + * + * ERROR the pin is not published. The launcher cannot install it. Blocking. + * One exception: on a ref headed for `main`, a pin equal to + * packages/mcp/package.json's version is the release being cut, and + * release-mcp.yml publishes exactly that from this tree moments later. + * Erroring there would redden every release. The exception is scoped to + * `main` on purpose: `dev` pinning an unpublished release version is + * precisely the 2026-09-17 bug, and must stay an error. + * + * WARNING the pin is published but is not what the dist-tag for this branch + * points at. Not blocking, because it cannot be: pushing to `dev` + * publishes a NEW `dev.N` and never commits that number back, so the + * manifest is one behind by construction for most of its life. Failing + * on that would redden CI after every merge. It is still worth saying — + * 0.0.14-dev.0 outlived the relayer change that made its own + * instructions wrong, and nothing anywhere said so. + */ + +import fs from "node:fs"; +import path from "node:path"; + +const PACKAGE = "@mysten-incubation/memwal-mcp"; +const MANIFEST = "packages/mcp/plugin/plugin.json"; +const PACKAGE_JSON = "packages/mcp/package.json"; +const REGISTRY = "https://registry.npmjs.org"; + +/** Which dist-tag this ref is supposed to be tracking. */ +function expectedTag() { + const override = process.argv + .find((a) => a.startsWith("--tag=")) + ?.slice("--tag=".length); + if (override) return override; + + // On a pull request GITHUB_REF_NAME is "/merge", so the branch that + // matters is the one being merged INTO. + const branch = process.env.GITHUB_BASE_REF || process.env.GITHUB_REF_NAME || "dev"; + if (branch === "main") return "latest"; + if (branch === "staging") return "rc"; + return "dev"; +} + +function annotate(level, message) { + if (process.env.GITHUB_ACTIONS) console.log(`::${level}::${message}`); + console.log(`${level.toUpperCase()}: ${message}`); +} + +function summarise(lines) { + const file = process.env.GITHUB_STEP_SUMMARY; + if (file) fs.appendFileSync(file, lines.join("\n") + "\n"); +} + +async function main() { + const manifestPath = path.join(process.cwd(), MANIFEST); + const manifest = JSON.parse(fs.readFileSync(manifestPath, "utf8")); + // Mirror pinnedVersion(): mcpPackageVersion wins, version is the fallback. + const pin = manifest.mcpPackageVersion ?? manifest.version; + if (typeof pin !== "string" || pin.trim() === "") { + annotate("error", `${MANIFEST} has no usable "mcpPackageVersion" or "version"`); + process.exit(1); + } + + let packument; + try { + const response = await fetch(`${REGISTRY}/${PACKAGE}`, { + headers: { accept: "application/json" }, + }); + if (!response.ok) throw new Error(`registry answered HTTP ${response.status}`); + packument = await response.json(); + } catch (error) { + // Failing open would make a green check mean nothing, so this is fatal — + // but say plainly that it is the registry, not the pin, so nobody starts + // editing plugin.json to chase a network blip. + annotate("error", `could not read ${PACKAGE} from npm (${error.message}). Retry the job.`); + process.exit(1); + } + + const versions = Object.keys(packument.versions ?? {}); + const tags = packument["dist-tags"] ?? {}; + const tag = expectedTag(); + const tagged = tags[tag]; + + const summary = [ + "### MCP launcher pin", + "", + `| | |`, + `|---|---|`, + `| \`mcpPackageVersion\` | \`${pin}\` |`, + `| \`version\` | \`${manifest.version}\` |`, + `| npm \`${tag}\` | \`${tagged ?? "(none)"}\` |`, + "", + ]; + + if (!versions.includes(pin)) { + // The release being cut is allowed to pin ahead of npm: release-mcp.yml + // publishes packages/mcp/package.json's version from this very tree a + // moment after the push. Only on a ref headed for `main`, though — + // `dev` pinning an unpublished release version is the bug this exists + // to catch. + const releasing = + tag === "latest" + && pin === JSON.parse(fs.readFileSync(path.join(process.cwd(), PACKAGE_JSON), "utf8")).version; + if (releasing) { + annotate( + "warning", + `${MANIFEST} pins ${pin}, which npm does not carry yet. That is the release ` + + `this ref publishes, so it is expected here — but nothing has verified the ` + + `publish succeeded. Check the release job before announcing it.`, + ); + summary.push(`Pinned to \`${pin}\`, the release this ref publishes. Not on npm yet.`); + summarise(summary); + return; + } + annotate( + "error", + `${MANIFEST} pins ${PACKAGE}@${pin}, which is not published. ` + + `The trusted launcher installs that exact version, so every plugin install ` + + `would fail. Published dist-tags: ` + + Object.entries(tags) + .map(([name, value]) => `${name}=${value}`) + .join(", "), + ); + summary.push(`**Not published.** The launcher cannot install \`${pin}\`.`); + summarise(summary); + process.exit(1); + } + + if (tagged && pin !== tagged) { + annotate( + "warning", + `${MANIFEST} pins ${pin}, but npm "${tag}" is ${tagged}. The pin is published, ` + + `so nothing breaks — but installs keep serving ${pin}. Bump ` + + `mcpPackageVersion if that is not deliberate.`, + ); + summary.push( + `Published, but behind the \`${tag}\` tag. Installs serve \`${pin}\`, not \`${tagged}\`.`, + ); + summarise(summary); + return; + } + + console.log(`OK: ${PACKAGE}@${pin} is published and matches the "${tag}" dist-tag.`); + summary.push(`Published, and matches the \`${tag}\` tag.`); + summarise(summary); +} + +await main(); diff --git a/scripts/verify-manual-sdk-release.mjs b/scripts/verify-manual-sdk-release.mjs index 314758341..e18db4897 100644 --- a/scripts/verify-manual-sdk-release.mjs +++ b/scripts/verify-manual-sdk-release.mjs @@ -1,17 +1,17 @@ #!/usr/bin/env node -import { readFileSync } from "node:fs"; +import { existsSync, readFileSync } from "node:fs"; const releases = [ { name: "TypeScript SDK", - version: "0.1.7", + version: "0.1.8", manifests: [["packages/sdk/package.json", "version"]], changelogs: ["packages/sdk/CHANGELOG.md", "docs/sdk/changelog.mdx"], }, { name: "Python SDK", - version: "0.1.10", + version: "0.1.11", manifests: [ ["packages/python-sdk-memwal/pyproject.toml", "toml-version"], ["packages/python-sdk-memwal/memwal/__init__.py", "python-version"], @@ -23,7 +23,7 @@ const releases = [ }, { name: "MCP package", - version: "0.0.13", + version: "0.0.14", manifests: [ ["packages/mcp/package.json", "version"], [".claude-plugin/marketplace.json", "plugin-version"], @@ -63,37 +63,96 @@ for (const release of releases) { console.log(`${release.name} ${release.version}: manifests and changelogs synchronized`); } +// The plugin no longer launches the server through `npx @`: npx resolves +// the name against the project the MCP client is started in, so a package planted +// there could answer to the pinned spec (WALM-640). Every launch site must run the +// plugin's launcher, which installs the pin under ~/.memwal/runtime and runs that +// absolute entry point. The pin itself is plugin/plugin.json's version, already +// checked against packages/mcp/package.json above — or its `mcpPackageVersion`, +// which names the published prerelease to install while that release version is +// not yet on npm. A pin the registry cannot serve fails the launch outright, so the +// override may only name a prerelease of the release it stands in for, and it has +// to go when that release is published. const mcpVersion = JSON.parse(readFileSync("packages/mcp/package.json", "utf8")).version; -const expectedPluginArgs = ["-y", `@mysten-incubation/memwal-mcp@${mcpVersion}`]; -for (const pluginPath of [ - "packages/mcp/plugin/.mcp.json", - "packages/mcp/plugin/.cursor-mcp.json", - "packages/mcp/plugin/.codex-mcp.json", +const pinOverride = JSON.parse( + readFileSync("packages/mcp/plugin/plugin.json", "utf8"), +).mcpPackageVersion; +if (pinOverride !== undefined && !pinOverride.startsWith(`${mcpVersion}-`)) { + throw new Error( + `packages/mcp/plugin/plugin.json: "mcpPackageVersion" is ${pinOverride}, which is not ` + + `a prerelease of ${mcpVersion}; drop it once ${mcpVersion} is on npm`, + ); +} +const LAUNCHER = "scripts/launch_mcp.mjs"; +for (const [pluginPath, rootPlaceholder] of [ + ["packages/mcp/plugin/.mcp.json", "${CLAUDE_PLUGIN_ROOT}"], + ["packages/mcp/plugin/.cursor-mcp.json", "${CURSOR_PLUGIN_ROOT}"], + ["packages/mcp/plugin/.codex-mcp.json", "${PLUGIN_ROOT}"], ]) { - const actual = JSON.parse(readFileSync(pluginPath, "utf8")).mcpServers.memwal.args; - if (JSON.stringify(actual) !== JSON.stringify(expectedPluginArgs)) { + const server = JSON.parse(readFileSync(pluginPath, "utf8")).mcpServers.memwal; + const expected = { command: "node", args: [`${rootPlaceholder}/${LAUNCHER}`] }; + if ( + server.command !== expected.command || + JSON.stringify(server.args) !== JSON.stringify(expected.args) + ) { throw new Error( - `${pluginPath}: expected ${JSON.stringify(expectedPluginArgs)}, received ${JSON.stringify(actual)}`, + `${pluginPath}: expected ${JSON.stringify(expected)}, received ${JSON.stringify({ command: server.command, args: server.args })}`, ); } } const installerPath = "packages/mcp/plugin/scripts/install_codex_hooks.mjs"; const installer = readFileSync(installerPath, "utf8"); -const expectedPin = expectedPluginArgs[1]; -if (installer.includes('["-y", "@mysten-incubation/memwal-mcp"]')) { +if (/command\s*=\s*\\?"npx/.test(installer)) { throw new Error( - `${installerPath}: expected ${JSON.stringify(expectedPluginArgs)}, received ${JSON.stringify(["-y", "@mysten-incubation/memwal-mcp"])}`, + `${installerPath}: registers the MCP server through npx; it must register the ` + + `absolute path to ${LAUNCHER} (WALM-640)`, ); } -if ( - !installer.includes(expectedPin) && - !installer.includes("@mysten-incubation/memwal-mcp@${") -) { +if (!installer.includes("launch_mcp.mjs")) { + throw new Error(`${installerPath}: does not register ${LAUNCHER}`); +} +const launcherPath = `packages/mcp/plugin/${LAUNCHER}`; +if (!existsSync(launcherPath)) { + throw new Error(`${launcherPath}: missing, but every launch site points at it`); +} + +// Registering the launcher for *new* installations is only half of it: every user who +// ran the installer before WALM-640 has `command = "npx"` in ~/.codex/config.toml, and +// an installer that skips an existing block leaves them on the vulnerable resolution +// for ever. Exercise the migration rather than grepping for its absence. +const codexConfigPath = "packages/mcp/plugin/scripts/lib/codex-config.mjs"; +if (!existsSync(codexConfigPath)) { + throw new Error(`${codexConfigPath}: missing, but ${installerPath} migrates through it`); +} +const { planMcpRegistration } = await import(`../${codexConfigPath}`); +const legacyConfig = [ + "[features]", + "codex_hooks = true", + "", + "[mcp_servers.memwal]", + 'command = "npx"', + `args = ["-y", "@mysten-incubation/memwal-mcp@${mcpVersion}"]`, + 'env = { MEMWAL_NAMESPACE = "work" }', + "", + "[mcp_servers.other]", + 'command = "other"', + "", +].join("\n"); +const migrated = planMcpRegistration(legacyConfig, "/abs/plugin/scripts/launch_mcp.mjs"); +if (migrated.action !== "migrated") { throw new Error( - `${installerPath}: expected ${JSON.stringify(expectedPluginArgs)}, received missing version pin`, + `${codexConfigPath}: an existing npx [mcp_servers.memwal] block must be migrated, ` + + `received action "${migrated.action}" (WALM-640)`, ); } -console.log(`MCP package ${mcpVersion}: plugin npx args pin ${expectedPin}`); +if (/command\s*=\s*"npx"/.test(migrated.content) || !migrated.content.includes(LAUNCHER)) { + throw new Error(`${codexConfigPath}: migration did not replace npx with ${LAUNCHER}`); +} +if (!migrated.content.includes("MEMWAL_NAMESPACE") || !migrated.content.includes("[mcp_servers.other]")) { + throw new Error(`${codexConfigPath}: migration dropped keys it does not own`); +} +console.log(`MCP package ${mcpVersion}: every launch site runs ${LAUNCHER} by absolute path`); +console.log(`MCP package ${mcpVersion}: an existing npx Codex registration is migrated, not skipped`); function readVersion(content, kind) { if (kind === "version") return JSON.parse(content).version; diff --git a/services/server/migrations/021_drop_failed_write_report_ack.sql b/services/server/migrations/021_drop_failed_write_report_ack.sql new file mode 100644 index 000000000..3bd5aca46 --- /dev/null +++ b/services/server/migrations/021_drop_failed_write_report_ack.sql @@ -0,0 +1,16 @@ +-- Drop the failed-write-report ack column that a PR #921 preview deploy +-- created on relayer.dev and then left behind. +-- +-- History: 021_failed_write_report_ack.sql (feature #918) added +-- remember_jobs.failure_reported_at so a failed write is reported to the +-- user once. origin/dev shipped the SQL file and the queries but never +-- wired the file into VectorDb::new(), so a fresh schema never grew the +-- column. PR #921 briefly wired it (d71efdbb), deployed that build to +-- relayer.dev.memwal.ai, then reverted the feature (da73a486) without a +-- down-migration. The revert is the right code fix; this file is the +-- leftover schema. Idempotent: DROP IF EXISTS is a no-op on databases +-- that never ran the ADD (CI, local, any env that only ever ran origin/dev). +-- +-- Index first: it is a partial index ON this column. +DROP INDEX IF EXISTS remember_jobs_unreported_failures_idx; +ALTER TABLE remember_jobs DROP COLUMN IF EXISTS failure_reported_at; diff --git a/services/server/migrations/022_remember_jobs_recent_outcomes.sql b/services/server/migrations/022_remember_jobs_recent_outcomes.sql new file mode 100644 index 000000000..12ec240ee --- /dev/null +++ b/services/server/migrations/022_remember_jobs_recent_outcomes.sql @@ -0,0 +1,16 @@ +-- Partial index for GET /health's recent_write_outcomes probe. +-- +-- The probe is `status IN ('failed','done','uploaded') AND updated_at >= $1` +-- with no owner. Existing indexes are (owner), (status), (owner, status, +-- updated_at DESC), and (owner, idempotency_key) — none of those can serve +-- a window scan on updated_at alone, and `done` is almost the whole table +-- so idx_remember_jobs_status is a near-seq-scan. remember_jobs is never +-- pruned. Without this index the 1s statement timeout always fires on a +-- large table and the probe fails open, so writes=degraded silently stops +-- working at the scale where it matters. +-- +-- Own file: CREATE INDEX CONCURRENTLY cannot run inside a transaction +-- (sqlx::raw_sql wraps a file as one), same reason as 016/018. +CREATE INDEX CONCURRENTLY IF NOT EXISTS remember_jobs_recent_outcomes_idx + ON remember_jobs (updated_at DESC) + WHERE status IN ('failed', 'done', 'uploaded'); diff --git a/services/server/scripts/__tests__/sidecar-query-helpers.test.ts b/services/server/scripts/__tests__/sidecar-query-helpers.test.ts index 1391c8b75..b9fa162e1 100644 --- a/services/server/scripts/__tests__/sidecar-query-helpers.test.ts +++ b/services/server/scripts/__tests__/sidecar-query-helpers.test.ts @@ -27,6 +27,7 @@ import { WALRUS_PACKAGE_ID, } from "../sidecar/config.js"; import { + addressBalanceExpiration, assertAddressBalanceRegisterTransaction, assertSponsoredRegisterTransactionKind, createdBlobObjectIdFromTransaction, @@ -1024,3 +1025,118 @@ test("normal registration checkpoints the exact Blob from transaction effects", }) ); }); + +test("a rejected register expiration names the bound that broke the guard", () => { + // The guard rejects on three conditions joined by `||`. When every upload + // job on testnet hit this (2026-09-17) the logged sentence named the + // invariant but not the value, so the relayer logs could not say which + // condition fired. The message has to carry the shape it saw. + const signer = new Ed25519Keypair(); + const transaction = new Transaction(); + transaction.setSender(signer.toSuiAddress()); + transaction.setGasOwner(signer.toSuiAddress()); + transaction.setGasBudget(1_000n); + transaction.setGasPrice(1n); + transaction.setGasPayment([]); + transaction.setExpiration({ + ValidDuring: { + minEpoch: "1", + maxEpoch: "2", + minTimestamp: "5", + maxTimestamp: null, + chain: "69WiPg3DAQiwdxfncX6wYQ2siKwAe6L9BZthQea3JNMD", + nonce: 1, + }, + }); + + const resolved = TransactionDataBuilder.restore(transaction.getData() as never); + assert.throws( + () => assertAddressBalanceRegisterTransaction(resolved), + (error: Error) => { + assert.match(error.message, /must use a ValidDuring address-balance expiration/); + assert.match(error.message, /expiration=ValidDuring/); + assert.match(error.message, /minTimestamp="5"/); + assert.match(error.message, /maxTimestamp=null/); + return true; + }, + ); +}); + + +/** Build a register the way `prepareRegisterTransaction` does: gas from the + * address balance, a WAL withdrawal, and the expiration the production path now + * applies. Every other register fixture in this file hand-sets `ValidDuring`, + * which is exactly why the direct-signed path could reject every write in + * production while the suite stayed green. */ +async function directSignedRegister(): Promise { + const signer = new Ed25519Keypair(); + const transaction = new Transaction(); + transaction.setSender(signer.toSuiAddress()); + transaction.setGasOwner(signer.toSuiAddress()); + transaction.setGasBudget(1_000n); + transaction.setGasPrice(1n); + transaction.setGasPayment([]); + const walType = `0x${"2".repeat(64)}::wal::WAL`; + const withdrawal = transaction.withdrawal({ amount: 1n, type: walType }); + const wal = transaction.moveCall({ + target: "0x2::coin::redeem_funds", + typeArguments: [walType], + arguments: [withdrawal], + }); + transaction.moveCall({ + target: "0x2::coin::destroy_zero", + typeArguments: [walType], + arguments: [wal], + }); + transaction.setExpiration( + addressBalanceExpiration(7n, await transaction.build({ onlyTransactionKind: true })), + ); + return TransactionDataBuilder.fromBytes(await transaction.build()); +} + +test("the expiration prepareRegisterTransaction applies is the one the guard demands", async () => { + // The positive half of the dev outage: `flow.register()` hands back a + // transaction with no expiration, and the guard one line later requires a + // ValidDuring window, so every direct-signed register failed its own + // assertion and no blob was certified. A register carrying what + // addressBalanceExpiration builds passes, and reports the epoch the journal + // then uses as its expiry guard. + assert.equal(assertAddressBalanceRegisterTransaction(await directSignedRegister()), 7n); +}); + +test("a register whose expiration is not ValidDuring is rejected", async () => { + // The `$kind` arm of the guard, which the timestamp case above does not + // reach. Built with the real expiration first: a withdrawal cannot be + // resolved offline without one, which is its own evidence that the window + // is not optional here. + const data = await directSignedRegister(); + data.expiration = { $kind: "Epoch", Epoch: "7" } as typeof data.expiration; + assert.throws( + () => assertAddressBalanceRegisterTransaction(data), + /must use a ValidDuring address-balance expiration/, + ); +}); + +test("the address-balance nonce is derived from the transaction, not drawn at random", () => { + const kind = new Uint8Array([1, 2, 3]); + const otherKind = new Uint8Array([1, 2, 4]); + + // Re-preparing the same register must rebuild byte-identical, or the + // journal stops being idempotent and a replay pays for a second blob. + assert.equal( + addressBalanceExpiration(9n, kind).ValidDuring.nonce, + addressBalanceExpiration(9n, kind).ValidDuring.nonce, + ); + // Two different registers must still reserve under different nonces. + assert.notEqual( + addressBalanceExpiration(9n, kind).ValidDuring.nonce, + addressBalanceExpiration(9n, otherKind).ValidDuring.nonce, + ); + + const { minEpoch, maxEpoch, minTimestamp, maxTimestamp } = + addressBalanceExpiration(9n, kind).ValidDuring; + assert.equal(minEpoch, "9"); + assert.equal(maxEpoch, "9"); + assert.equal(minTimestamp, null); + assert.equal(maxTimestamp, null); +}); diff --git a/services/server/scripts/__tests__/sidecar-slot-lifecycle.test.ts b/services/server/scripts/__tests__/sidecar-slot-lifecycle.test.ts new file mode 100644 index 000000000..5e56e43f5 --- /dev/null +++ b/services/server/scripts/__tests__/sidecar-slot-lifecycle.test.ts @@ -0,0 +1,49 @@ +import test from "node:test"; +import assert from "node:assert/strict"; + +process.env.WALRUS_UPLOAD_MAX_CONCURRENCY = "1"; +process.env.WALRUS_UPLOAD_PER_WALLET_CONCURRENCY = "1"; +process.env.WALRUS_UPLOAD_ACQUIRE_TIMEOUT_MS = "1000"; + +const { acquireWalrusUploadSlots, getUploadCounts, walrusUploadLimitSnapshot } = + await import("../sidecar/concurrency.js"); + +test("global acquire timeout releases the wallet reservation and queue count", async () => { + const release = await acquireWalrusUploadSlots(0, "held-global"); + try { + await assert.rejects( + acquireWalrusUploadSlots(1, "queued-global"), + /global upload slot/ + ); + assert.deepEqual(getUploadCounts(), { active: 1, queued: 0 }); + assert.deepEqual(walrusUploadLimitSnapshot(1).wallet, { + capacity: 1, + available: 1, + queued: 0, + }); + } finally { + release(); + } + const releaseNext = await acquireWalrusUploadSlots(1, "after-timeout"); + releaseNext(); + assert.deepEqual(getUploadCounts(), { active: 0, queued: 0 }); +}); + +test("releasing an old holder twice cannot free a successor's slot", async () => { + const releaseFirst = await acquireWalrusUploadSlots(0, "first"); + const second = acquireWalrusUploadSlots(0, "second"); + releaseFirst(); + const releaseSecond = await second; + try { + releaseFirst(); + assert.deepEqual(getUploadCounts(), { active: 1, queued: 0 }); + assert.deepEqual(walrusUploadLimitSnapshot(0).wallet, { + capacity: 1, + available: 0, + queued: 0, + }); + } finally { + releaseSecond(); + } + assert.deepEqual(getUploadCounts(), { active: 0, queued: 0 }); +}); diff --git a/services/server/scripts/__tests__/sidecar-upload-metrics.test.ts b/services/server/scripts/__tests__/sidecar-upload-metrics.test.ts new file mode 100644 index 000000000..2351c56a3 --- /dev/null +++ b/services/server/scripts/__tests__/sidecar-upload-metrics.test.ts @@ -0,0 +1,69 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import type { Server } from "node:http"; + +const { Ed25519Keypair } = await import("@mysten/sui/keypairs/ed25519"); +process.env.SERVER_SUI_PRIVATE_KEYS = new Ed25519Keypair().getSecretKey(); +process.env.SIDECAR_AUTH_TOKEN = "upload-metrics-test-token"; +process.env.SUI_NETWORK = "testnet"; +process.env.SUI_GRPC_URL = "https://upload-metrics.testnet.example"; +process.env.WALRUS_PACKAGE_ID = `0x${"a".repeat(64)}`; +process.env.WALRUS_UPLOAD_MAX_CONCURRENCY = "4"; +process.env.WALRUS_UPLOAD_PER_WALLET_CONCURRENCY = "1"; +process.env.WALRUS_UPLOAD_ACQUIRE_TIMEOUT_MS = "30000"; + +const { getWalrusClient, suiClient } = await import("../sidecar/clients.js"); +let dependencyCalls = 0; +(suiClient.ledgerService as any).getServiceInfo = async () => { + dependencyCalls += 1; + throw new Error("upload metrics must not call Sui"); +}; +(getWalrusClient() as any).getBlobType = async () => { + dependencyCalls += 1; + throw new Error("upload metrics must not call Walrus"); +}; +const { acquireWalrusUploadSlots } = await import("../sidecar/concurrency.js"); +const { createSidecarApp } = await import("../sidecar/app.js"); + +async function listen(mode: "full" | "writer"): Promise<{ server: Server; baseUrl: string }> { + return new Promise((resolve) => { + const server = createSidecarApp(mode).listen(0, "127.0.0.1", () => { + const address = server.address(); + assert.ok(address && typeof address !== "string"); + resolve({ server, baseUrl: `http://127.0.0.1:${address.port}` }); + }); + }); +} + +async function close(server: Server): Promise { + await new Promise((resolve, reject) => { + server.close((error) => (error ? reject(error) : resolve())); + }); +} + +for (const mode of ["full", "writer"] as const) { + test(`${mode} mode: /metrics/uploads reports live queue counters without a token`, async () => { + const { server, baseUrl } = await listen(mode); + // Wallet 0 has one slot: the first upload runs, the second waits in the queue. + const releaseRunning = await acquireWalrusUploadSlots(0, "running-upload"); + const queuedUpload = acquireWalrusUploadSlots(0, "queued-upload"); + try { + const res = await fetch(`${baseUrl}/metrics/uploads`); + assert.equal(res.status, 200); + assert.deepEqual(await res.json(), { + activeWalrusUploads: 1, + queuedWalrusUploads: 1, + walrusUploadLimits: { + globalCapacity: 4, + perWalletCapacity: 1, + acquireTimeoutMs: 30000, + }, + }); + assert.equal(dependencyCalls, 0); + } finally { + releaseRunning(); + (await queuedUpload)(); + await close(server); + } + }); +} diff --git a/services/server/scripts/mcp/__tests__/analyze-fast-return.test.ts b/services/server/scripts/mcp/__tests__/analyze-fast-return.test.ts new file mode 100644 index 000000000..5b8619377 --- /dev/null +++ b/services/server/scripts/mcp/__tests__/analyze-fast-return.test.ts @@ -0,0 +1,171 @@ +/** + * `memwal_analyze` returns once the facts are extracted and queued. + * + * It was the last tool still blocking to terminal after the two remember tools + * moved to a bounded wait, which left it the slowest in the set by a wide + * margin — 37.0s measured against dev in the same session where + * `memwal_remember` came back in 0.2s. The wait has the same shape as bulk's + * (N Walrus writes, one upload per wallet), so there was no reason for the + * answer to be shaped differently. + * + * What must stay true, and is what these tests pin: extraction is still waited + * for, because the facts are the part an agent can act on; the reply never + * reads as saved when the writes are still in flight; and every job_id comes + * back paired with the fact it carries, so a later partial failure is + * actionable. + */ +// A small non-zero wait so the bounded-wait branch is reachable at all. The +// shipped default is now 0 — the tools return at accept — so leaving it unset +// would skip the wait entirely and every assertion below about a partly landed +// batch would have nothing to observe. +// Set before the dynamic import, because the budget is read once at load. +process.env.MEMWAL_MCP_REMEMBER_WAIT_MS = "500"; + +import assert from "node:assert/strict"; +import test, { type TestContext } from "node:test"; +import { Client } from "@modelcontextprotocol/sdk/client/index.js"; +import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js"; +import type { MemWalSession } from "../auth.js"; + +// Dynamic, not static: ESM hoists static imports above the assignment above, +// so the module would read the real default before the line that changes it. +const { createMcpServer } = await import("../server.js"); + +const FACTS = ["User drinks oat milk", "User lives in Ho Chi Minh City"]; + +function sessionWith( + opts: { states?: Array<"done" | "failed" | "timeout">; facts?: string[] } = {}, + calls: string[] = [], +): MemWalSession { + const facts = opts.facts ?? FACTS; + const jobIds = facts.map((_, i) => `analyze-job-${i + 1}`); + return { + oauthScope: "memwal:read memwal:write", + namespace: "default", + memwal: { + async analyze(text: string) { + calls.push(`analyze:${text.slice(0, 20)}`); + return { + job_ids: jobIds, + facts: facts.map((t) => ({ text: t })), + fact_count: facts.length, + status: "accepted", + owner: "0xowner", + }; + }, + async analyzeAndWait() { + calls.push("analyzeAndWait"); + throw new Error("analyze must not block to terminal any more"); + }, + async waitForRememberJobs(ids: string[]) { + calls.push(`waitJobs:${ids.join(",")}`); + const states = opts.states ?? facts.map(() => "done" as const); + return { + results: states.map((status, i) => ({ + id: jobIds[i], + blob_id: status === "done" ? `blob-${i + 1}` : "", + status, + namespace: "default", + })), + total: states.length, + succeeded: states.filter((s) => s === "done").length, + // Deliberately total-minus-succeeded, because that is what + // the real `waitForRememberJobs` returns — a `timeout` + // counts here. A stub that filtered on "failed" instead + // would model a value the SDK never produces, and the + // still-uploading-counted-as-failed case below could not + // fail no matter what the tool printed. + failed: states.length - states.filter((s) => s === "done").length, + }; + }, + }, + } as unknown as MemWalSession; +} + +async function callAnalyze( + session: MemWalSession, + t: TestContext, + text = "a long note about the user", +): Promise { + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair(); + const server = createMcpServer(session); + const client = new Client({ name: "analyze-test", version: "1.0.0" }); + t.after(async () => { + await client.close(); + await server.close(); + }); + await server.connect(serverTransport); + await client.connect(clientTransport); + + const result = await client.callTool({ + name: "memwal_analyze", + arguments: { text }, + }); + return (result as { content: Array<{ text: string }> }).content + .map((c) => c.text) + .join("\n"); +} + +test("analyze never blocks to terminal", async (t) => { + const calls: string[] = []; + await callAnalyze(sessionWith({}, calls), t); + assert.ok( + !calls.includes("analyzeAndWait"), + `took the blocking path: ${calls.join(", ")}`, + ); + assert.ok(calls.some((c) => c.startsWith("analyze:")), calls.join(", ")); +}); + +test("the extracted facts come back even though the writes have not landed", async (t) => { + // Extraction is the half an agent can use straight away. Handing back only + // job_ids would make the tool useless until a second call. + const text = await callAnalyze(sessionWith(), t); + for (const fact of FACTS) assert.match(text, new RegExp(fact)); +}); + +test("an in-flight analyze does not read as saved", async (t) => { + const text = await callAnalyze(sessionWith({ states: ["timeout", "timeout"] }), t); + assert.match(text, /NOT SAVED YET|ACCEPTED, NOT YET SAVED/); + assert.match(text, /memwal_remember_status/); + assert.doesNotMatch(text, /^Saved to Walrus Memory/m); +}); + +test("every job_id is paired with the fact it carries", async (t) => { + // "one of these failed" is only actionable if the agent can tell which. + const text = await callAnalyze(sessionWith({ states: ["timeout", "timeout"] }), t); + assert.match(text, /analyze-job-1/); + assert.match(text, /analyze-job-2/); + assert.match(text, /analyze-job-1 — User drinks oat milk/); +}); + +test("a partly landed batch reports both halves", async (t) => { + const text = await callAnalyze(sessionWith({ states: ["done", "timeout"] }), t); + assert.match(text, /blob-1/, "the landed write shows its blob_id"); + assert.match(text, /analyze-job-2/, "the straggler shows its job_id"); + assert.match(text, /memwal_remember_status/, "and how to settle it"); +}); + +test("a still-uploading write is not counted or labelled as failed", async (t) => { + // The straggler block already tells the agent this job is on its way and + // must not be re-sent. Printing `failed=1` next to it contradicts that, + // and an agent that believes the count re-sends — a duplicate paid Walrus + // write queued behind the original. + const text = await callAnalyze(sessionWith({ states: ["done", "timeout"] }), t); + assert.doesNotMatch(text, /failed=/, "a timeout is in flight, not failed"); + assert.match(text, /1 still uploading/, "it is counted as in flight instead"); + assert.doesNotMatch(text, /\[timeout\]/, "and not labelled with the raw status"); + assert.match(text, /still uploading, job_id=analyze-job-2/); +}); + +test("a genuinely failed write is still counted as failed", async (t) => { + const text = await callAnalyze(sessionWith({ states: ["done", "failed"] }), t); + assert.match(text, /failed=1/, "a terminal failure must still be reported"); + assert.doesNotMatch(text, /still uploading/); +}); + +test("text with nothing worth saving says so instead of handing back an empty batch", async (t) => { + const text = await callAnalyze(sessionWith({ facts: [] }), t); + assert.match(text, /Extracted 0 facts/); + assert.match(text, /nothing was saved/); + assert.doesNotMatch(text, /memwal_remember_status/); +}); diff --git a/services/server/scripts/mcp/__tests__/live-e2e.mjs b/services/server/scripts/mcp/__tests__/live-e2e.mjs new file mode 100644 index 000000000..9420d4ecd --- /dev/null +++ b/services/server/scripts/mcp/__tests__/live-e2e.mjs @@ -0,0 +1,169 @@ +/** + * End-to-end exercise of the PR's MCP tools against a LIVE relayer. + * + * The unit tests drive these tools through a mocked session. That pins the + * wording and the branching, but it cannot answer the question that actually + * matters before merge: does `memwal_remember_bulk` hand back job_ids the + * real `memwal_remember_status` can then resolve into real blob_ids? + * + * So this wires the same in-process MCP server the unit tests use to a real + * MemWal session, and walks the flow a user's agent would walk. It is not a + * unit test — it writes paid blobs — which is why it lives here as a script + * rather than under `node --test`. + * + * NODE_USE_ENV_PROXY=1 MEMWAL_CREDS_DIR= \ + * node --import tsx mcp/__tests__/live-e2e.mjs + */ +import { readFileSync } from "node:fs"; +import assert from "node:assert/strict"; +import { Client } from "@modelcontextprotocol/sdk/client/index.js"; +import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js"; +import { MemWal } from "@mysten-incubation/memwal"; +import { createMcpServer } from "../server.js"; + +const credsDir = process.env.MEMWAL_CREDS_DIR; +if (!credsDir) throw new Error("set MEMWAL_CREDS_DIR"); +const creds = JSON.parse(readFileSync(`${credsDir}/credentials.json`, "utf8")); +const NS = "memwal-probe"; + +const memwal = new MemWal({ + key: creds.delegatePrivateKey, + accountId: creds.accountId, + serverUrl: creds.relayerUrl, + namespace: NS, +}); + +const session = { + accountId: creds.accountId, + delegateKeyHex: creds.delegatePrivateKey, + delegatePubKeyHex: creds.delegatePublicKeyHex, + namespace: NS, + memwal, + relayerUrl: creds.relayerUrl, + authMethod: "delegate-key", + oauthScope: "memwal:read memwal:write", + agentClient: "other", +}; + +const server = createMcpServer(session); +const client = new Client({ name: "live-e2e", version: "1.0.0" }, { capabilities: {} }); +const [clientT, serverT] = InMemoryTransport.createLinkedPair(); +await Promise.all([client.connect(clientT), server.connect(serverT)]); + +const ms = () => performance.now(); +const fmt = (n) => (n < 1000 ? `${Math.round(n)}ms` : `${(n / 1000).toFixed(1)}s`); +const textOf = (r) => r.content?.[0]?.text ?? ""; + +async function call(name, args) { + const t0 = ms(); + const r = await client.callTool({ name, arguments: args }); + return { ms: ms() - t0, text: textOf(r), isError: r.isError === true }; +} + +function show(label, r, extra = "") { + console.log( + ` ${label.padEnd(30)} ${fmt(r.ms).padStart(8)} ${r.isError ? "ERR " : "ok "} ${extra}`, + ); +} + +const stamp = new Date().toISOString(); +const failures = []; +function check(name, fn) { + try { + fn(); + console.log(` ✔ ${name}`); + } catch (e) { + failures.push(name); + console.log(` ✖ ${name}\n ${e.message.split("\n")[0]}`); + } +} + +console.log(`\nlive e2e — relayer ${creds.relayerUrl} ns=${NS}\n`); + +// ── single write ──────────────────────────────────────────────── +console.log("single remember → status"); +const single = await call("memwal_remember", { + text: `live-e2e single ${stamp}`, + namespace: NS, +}); +show("memwal_remember", single); +const singleJob = /job_id=([0-9a-f-]+)/.exec(single.text)?.[1]; + +check("remember returns fast", () => assert.ok(single.ms < 5000, `${fmt(single.ms)}`)); +check("remember hands back a job_id", () => assert.ok(singleJob, single.text.slice(0, 120))); +check("remember does not claim it is saved", () => + assert.doesNotMatch(single.text, /Saved to Walrus Memory/)); + +const statusNoWait = await call("memwal_remember_status", { job_id: singleJob, waitMs: 0 }); +show("status waitMs=0", statusNoWait); +check("waitMs=0 answers immediately", () => + assert.ok(statusNoWait.ms < 3000, `${fmt(statusNoWait.ms)}`)); + +// ── bulk write → job_ids → batch status ───────────────────────── +console.log("\nbulk remember → status with job_ids"); +const bulk = await call("memwal_remember_bulk", { + namespace: NS, + facts: [1, 2, 3].map((i) => `live-e2e bulk ${stamp} item ${i}`), +}); +show("memwal_remember_bulk", bulk); +const bulkJobs = [...bulk.text.matchAll(/job_id=([0-9a-f-]+)/g)].map((m) => m[1]); + +check("bulk returns fast", () => assert.ok(bulk.ms < 10_000, `${fmt(bulk.ms)}`)); +check("bulk hands back one job_id per fact", () => + assert.equal(bulkJobs.length, 3, `got ${bulkJobs.length}: ${bulk.text.slice(0, 200)}`)); +check("bulk pairs each job_id with its fact", () => + assert.match(bulk.text, /item 1/)); + +// THE question: does job_ids actually resolve to data? +// 45s is the tool's ceiling — deliberately under the MCP client's own 60s +// request deadline, so the tool answers rather than the client giving up. +const batch = await call("memwal_remember_status", { job_ids: bulkJobs, waitMs: 45_000 }); +show("status job_ids (45s budget)", batch); +check("job_ids returns a line per job", () => { + // Guard the loop: with no ids collected it would pass by doing nothing, + // which is exactly the case this check exists to catch. + assert.ok(bulkJobs.length > 0, "no job_ids to resolve"); + for (const id of bulkJobs) assert.match(batch.text, new RegExp(id.slice(0, 8))); +}); +check("job_ids accounts for every job, saved or not", () => { + // Not "must have blob_ids": measured p50 is ~34s and p90 ~65s, so a 45s + // budget legitimately expires with writes still in flight. What must hold + // is that every job comes back with a definite state and the report never + // reads as success when nothing landed. + assert.match(batch.text, /\d+\/\d+ saved/, batch.text.slice(0, 200)); + const blobs = [...batch.text.matchAll(/blob_id=([A-Za-z0-9_-]{20,})/g)]; + const saved = Number(/(\d+)\/\d+ saved/.exec(batch.text)?.[1] ?? "0"); + assert.equal(blobs.length, saved, "a blob_id for each job reported saved, and no more"); + if (saved === 0) assert.match(batch.text, /still uploading/); +}); + +// ── edge cases ────────────────────────────────────────────────── +console.log("\nedge cases"); +const both = await call("memwal_remember_status", { job_id: "a", job_ids: ["b"] }); +show("job_id + job_ids together", both); +check("rejects both ids at once", () => assert.ok(both.isError || /not both/i.test(both.text))); + +const neither = await call("memwal_remember_status", {}); +show("neither id", neither); +check("rejects an empty call", () => assert.ok(neither.isError || /Pass job_id/i.test(neither.text))); + +const unknown = await call("memwal_remember_status", { + job_id: "00000000-0000-0000-0000-000000000000", + waitMs: 0, +}); +show("unknown job_id", unknown); +check("an unknown job is an error, not a silent ok", () => + assert.ok(unknown.isError || /not found/i.test(unknown.text))); + +// ── recall, including the failed-write report ─────────────────── +console.log("\nrecall"); +const recall = await call("memwal_recall", { query: "live-e2e", limit: 5 }); +show("memwal_recall", recall, `${recall.text.split("\n").length} lines`); +check("recall answers under the SDK's 15s abort", () => + assert.ok(recall.ms < 15_000, `${fmt(recall.ms)}`)); +check("recall never invents a failure report", () => { + if (/FAILED/.test(recall.text)) assert.match(recall.text, /NOT stored/); +}); + +console.log(`\n${failures.length === 0 ? "all checks passed" : `${failures.length} FAILED: ${failures.join(", ")}`}\n`); +process.exit(failures.length === 0 ? 0 : 1); diff --git a/services/server/scripts/mcp/__tests__/remember-bulk-fast-return.test.ts b/services/server/scripts/mcp/__tests__/remember-bulk-fast-return.test.ts new file mode 100644 index 000000000..f94be7de7 --- /dev/null +++ b/services/server/scripts/mcp/__tests__/remember-bulk-fast-return.test.ts @@ -0,0 +1,222 @@ +// Same opt-in as remember-fast-return: the wait is the default, this file +// covers the accept-and-continue path behind the knob. +process.env.MEMWAL_MCP_REMEMBER_WAIT_MS = "0"; + +import assert from "node:assert/strict"; +import test, { type TestContext } from "node:test"; +import { Client } from "@modelcontextprotocol/sdk/client/index.js"; +import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js"; +import type { MemWalSession } from "../auth.js"; + +const { createMcpServer } = await import("../server.js"); + +/** + * `memwal_remember_bulk` used to block until every job in the batch reached a + * terminal state, on a fixed 120s budget. That made it the slowest tool while + * the server instructions steer an agent to it for any multi-fact turn, and a + * batch is N separate Walrus writes contending for one upload slot per wallet + * — so a five-fact batch could burn the whole budget and return nothing but + * timeouts. + * + * It now mirrors `memwal_remember`: accept, then hand back the job_ids. These + * tests pin the two things that keeps honest — an accepted batch must never + * read as saved, and every job_id must come back paired with its fact so a + * later partial failure is actionable. + */ + +interface BulkBehaviour { + /** Per-job terminal state, in input order. */ + states: Array<"done" | "failed" | "timeout">; +} + +function sessionWith(b: BulkBehaviour, calls: string[] = []): MemWalSession { + const jobIds = b.states.map((_, i) => `job-${i + 1}`); + return { + oauthScope: "memwal:read memwal:write", + namespace: "default", + memwal: { + async rememberBulkAsync(items: Array<{ text: string }>) { + calls.push(`bulkAsync:${items.map((i) => i.text).join("|")}`); + return { job_ids: jobIds, total: jobIds.length, status: "accepted" }; + }, + async rememberBulkAndWait() { + calls.push("bulkAndWait"); + throw new Error("bulk must not block to terminal any more"); + }, + async waitForRememberJobs(ids: string[]) { + calls.push(`waitJobs:${ids.join(",")}`); + return { + results: b.states.map((status, i) => ({ + id: jobIds[i], + blob_id: status === "done" ? `blob-${i + 1}` : "", + status, + namespace: "default", + error: status === "failed" ? "walrus upload rejected" : undefined, + })), + total: b.states.length, + succeeded: b.states.filter((s) => s === "done").length, + failed: b.states.filter((s) => s !== "done").length, + }; + }, + async getRememberBulkStatus(ids: string[]) { + calls.push(`bulkStatus:${ids.join(",")}`); + return { + results: ids.map((id, i) => ({ + job_id: id, + status: b.states[i] === "timeout" ? "running" : b.states[i], + blob_id: b.states[i] === "done" ? `blob-${i + 1}` : undefined, + error: b.states[i] === "failed" ? "walrus upload rejected" : undefined, + })), + }; + }, + }, + } as unknown as MemWalSession; +} + +async function clientFor(session: MemWalSession, t: TestContext): Promise { + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair(); + const server = createMcpServer(session); + const client = new Client({ name: "remember-bulk-fast-return-test", version: "1.0.0" }); + t.after(async () => { + await client.close(); + await server.close(); + }); + await server.connect(serverTransport); + await client.connect(clientTransport); + return client; +} + +function textOf(result: unknown): string { + return (result as { content: Array<{ text: string }> }).content + .map((c) => c.text) + .join("\n"); +} + +test("memwal_remember_bulk returns at accept and never reads as saved", async (t) => { + const calls: string[] = []; + const client = await clientFor( + sessionWith({ states: ["done", "done"] }, calls), + t, + ); + + const text = textOf( + await client.callTool({ + name: "memwal_remember_bulk", + arguments: { facts: ["likes espresso", "works in Hanoi"] }, + }), + ); + + // Accepted, not stored — and it must not have blocked on the batch. + assert.match(text, /ACCEPTED, NOT YET SAVED/); + assert.ok(calls.some((c) => c.startsWith("bulkAsync:")), "should accept via rememberBulkAsync"); + assert.ok(!calls.includes("bulkAndWait"), "must not block to terminal"); + assert.ok( + !calls.some((c) => c.startsWith("waitJobs:")), + "a zero budget must not enter the wait loop", + ); + // No blob_id may appear — that is the token an agent reads as "stored". + assert.ok(!/blob_id=/.test(text), `accepted batch leaked a blob_id: ${text}`); +}); + +test("every job_id comes back paired with its fact", async (t) => { + const client = await clientFor(sessionWith({ states: ["done", "done", "done"] }), t); + + const text = textOf( + await client.callTool({ + name: "memwal_remember_bulk", + arguments: { facts: ["alpha fact", "beta fact", "gamma fact"] }, + }), + ); + + // "one of these failed" is only actionable if the agent can tell which. + assert.match(text, /job_id=job-1 — alpha fact/); + assert.match(text, /job_id=job-2 — beta fact/); + assert.match(text, /job_id=job-3 — gamma fact/); +}); + +test("the accepted batch tells the agent how to settle it and not to re-send", async (t) => { + const client = await clientFor(sessionWith({ states: ["done"] }), t); + + const text = textOf( + await client.callTool({ + name: "memwal_remember_bulk", + arguments: { facts: ["a durable fact"] }, + }), + ); + + assert.match(text, /memwal_remember_status/); + assert.match(text, /job_ids=/); + // Re-sending queues a second paid copy behind the first. + assert.match(text, /[Dd]o not re-send/); +}); + +test("memwal_remember_status settles a whole batch in one call", async (t) => { + const calls: string[] = []; + const client = await clientFor( + sessionWith({ states: ["done", "failed", "timeout"] }, calls), + t, + ); + + const text = textOf( + await client.callTool({ + name: "memwal_remember_status", + arguments: { job_ids: ["job-1", "job-2", "job-3"], waitMs: 1000 }, + }), + ); + + // A mixed batch must report every outcome rather than throwing on the + // first failure — otherwise the blob_ids that DID land are lost. + assert.match(text, /1\/3 saved/); + assert.match(text, /blob_id=blob-1/); + assert.match(text, /NOT STORED/); + assert.match(text, /still uploading/); + // And it must name which ids still need chasing, and which need re-sending. + assert.match(text, /job_ids=\[job-3\]/); + assert.match(text, /must be sent again/); + assert.ok(calls.some((c) => c.startsWith("waitJobs:")), "should use the batch wait"); +}); + +test("memwal_remember_status with waitMs=0 reads a batch without waiting", async (t) => { + const calls: string[] = []; + const client = await clientFor(sessionWith({ states: ["done", "timeout"] }, calls), t); + + const text = textOf( + await client.callTool({ + name: "memwal_remember_status", + arguments: { job_ids: ["job-1", "job-2"], waitMs: 0 }, + }), + ); + + assert.ok( + calls.some((c) => c.startsWith("bulkStatus:")), + "a zero budget must be a single batched read", + ); + assert.ok(!calls.some((c) => c.startsWith("waitJobs:")), "must not enter the wait loop"); + assert.match(text, /1\/2 saved/); +}); + +test("memwal_remember_status rejects job_id and job_ids together", async (t) => { + const client = await clientFor(sessionWith({ states: ["done"] }), t); + + const result = await client.callTool({ + name: "memwal_remember_status", + arguments: { job_id: "job-1", job_ids: ["job-2"] }, + }); + + // They would describe different writes; guessing which one the caller + // meant would report the wrong fact's fate. + assert.equal((result as { isError?: boolean }).isError, true); + assert.match(textOf(result), /not both/); +}); + +test("memwal_remember_status requires one of job_id or job_ids", async (t) => { + const client = await clientFor(sessionWith({ states: ["done"] }), t); + + const result = await client.callTool({ + name: "memwal_remember_status", + arguments: { waitMs: 0 }, + }); + + assert.equal((result as { isError?: boolean }).isError, true); + assert.match(textOf(result), /job_id.*job_ids/s); +}); diff --git a/services/server/scripts/mcp/__tests__/remember-deadline.test.ts b/services/server/scripts/mcp/__tests__/remember-deadline.test.ts new file mode 100644 index 000000000..bf1a0ceb4 --- /dev/null +++ b/services/server/scripts/mcp/__tests__/remember-deadline.test.ts @@ -0,0 +1,423 @@ +// Bound the SDK calls that have no deadline of their own. Set before the +// module under test is imported — ACCEPT_DEADLINE_MS is read once at load. +process.env.MEMWAL_MCP_ACCEPT_DEADLINE_MS = "150"; +// This file is about the ACCEPT leg, so skip the terminal wait that is now the +// default — otherwise every case here would also drive the poll loop. +process.env.MEMWAL_MCP_REMEMBER_WAIT_MS = "0"; + +import assert from "node:assert/strict"; +import test, { type TestContext } from "node:test"; +import { Client } from "@modelcontextprotocol/sdk/client/index.js"; +import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js"; +import type { MemWalSession } from "../auth.js"; + +// Dynamic, not static: ESM hoists every static import above the assignment +// above, so the module would read the real 15s default before the line that +// shortens it ever runs — and each tool test would then take 15 seconds. +const { createMcpServer } = await import("../server.js"); +const { ACCEPT_DEADLINE_MS, withDeadline, MAX_ABSORBED_COOLDOWN_MS, DEFAULT_ACCEPT_DEADLINE_MS } = + await import("../tools/remember-wait.js"); + +/** + * The SDK's `signedRequest` aborts a request only when the caller passes a + * signal, and of the memory methods only `recall()` does. `rememberAsync`, + * `rememberBulkAsync` and every job-status poll call it with none, so `fetch` + * runs with no deadline. `timeoutMs` is checked at the top of the poll loop, so + * it bounds when the next poll STARTS, not how long one takes — which is how a + * tool documented as capping at 90s was observed still running past 120s. + * + * Returning at accept does not fix that by itself: the accept POST is one of + * the unbounded calls. These tests pin that every entry point is bounded, and + * that the resulting error never reads as a saved write. + */ + +/** A promise that never settles — what a stalled socket looks like from here. */ +function hangs(): Promise { + return new Promise(() => {}); +} + +function sessionThatHangs(): MemWalSession { + return { + oauthScope: "memwal:read memwal:write", + namespace: "default", + memwal: { + rememberAsync: hangs, + rememberBulkAsync: hangs, + getRememberStatus: hangs, + getRememberBulkStatus: hangs, + waitForRememberJob: hangs, + waitForRememberJobs: hangs, + }, + } as unknown as MemWalSession; +} + +async function clientFor(session: MemWalSession, t: TestContext): Promise { + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair(); + const server = createMcpServer(session); + const client = new Client({ name: "remember-deadline-test", version: "1.0.0" }); + t.after(async () => { + await client.close(); + await server.close(); + }); + await server.connect(serverTransport); + await client.connect(clientTransport); + return client; +} + +function textOf(result: unknown): string { + return (result as { content: Array<{ text: string }> }).content + .map((c) => c.text) + .join("\n"); +} + +test("the accept deadline is read from the environment and validated", () => { + assert.equal(ACCEPT_DEADLINE_MS, 150); +}); + +test("withDeadline passes a value through untouched when work finishes first", async () => { + assert.equal(await withDeadline(Promise.resolve("ok"), 1_000, "nope"), "ok"); +}); + +test("withDeadline rejects with a named error once the deadline passes", async () => { + await assert.rejects( + withDeadline(hangs(), 20, "relayer went quiet"), + (err: Error) => { + // `wrapTool` routes on the name, so it has to be distinct from a + // job failure — this is "we do not know", not "it failed". + assert.equal(err.name, "MemWalRelayerUnresponsive"); + assert.match(err.message, /relayer went quiet/); + return true; + }, + ); +}); + +test("withDeadline does not reject once the work has already resolved", async () => { + // A leftover timer firing after resolution would reject a promise nobody + // is racing any more, surfacing as an unhandled rejection. + const value = await withDeadline(Promise.resolve(1), 10, "nope"); + assert.equal(value, 1); + await new Promise((r) => setTimeout(r, 40)); +}); + +test("memwal_remember cannot hang forever on a stalled accept", async (t) => { + const client = await clientFor(sessionThatHangs(), t); + + const result = await client.callTool({ + name: "memwal_remember", + arguments: { text: "a durable fact" }, + }); + + assert.equal((result as { isError?: boolean }).isError, true); + const text = textOf(result); + // Must not read as stored, and must say a retry is safe rather than + // leaving the agent to guess (a blind retry would risk a second paid copy). + assert.ok(!/^Saved to Walrus Memory/m.test(text), `read as saved: ${text}`); + assert.match(text, /did not accept/); + assert.match(text, /[Rr]etry/); +}); + +test("memwal_remember_bulk cannot hang forever on a stalled accept", async (t) => { + const client = await clientFor(sessionThatHangs(), t); + + const result = await client.callTool({ + name: "memwal_remember_bulk", + arguments: { facts: ["one", "two"] }, + }); + + assert.equal((result as { isError?: boolean }).isError, true); + const text = textOf(result); + assert.match(text, /did not accept/); + // /api/remember/bulk carries NO idempotency key — the handler mints a + // fresh uuid per item — so the single path's "retrying is safe" line is a + // lie here, and an expensive one: withDeadline does not cancel the request, + // so the relayer has usually accepted by the time this fires. + assert.match(text, /Do NOT retry blindly/); + assert.match(text, /second time at full cost|SECOND time at full cost/i); + assert.doesNotMatch( + text, + /Retrying in this session is safe/, + "bulk must never claim idempotency it does not have", + ); +}); + +test("memwal_remember's accept timeout still says a retry is safe", async (t) => { + // The single path DOES carry a content-derived idempotency key, so the + // opposite advice is correct there — and worth pinning, because collapsing + // both messages into one is exactly how the bulk bug happened. + const client = await clientFor(sessionThatHangs(), t); + const result = await client.callTool({ + name: "memwal_remember", + arguments: { text: "a durable fact" }, + }); + assert.match(textOf(result), /Retrying is safe/); +}); + +test("memwal_remember_status cannot hang forever on a stalled read", async (t) => { + const client = await clientFor(sessionThatHangs(), t); + + const result = await client.callTool({ + name: "memwal_remember_status", + arguments: { job_id: "job-1", waitMs: 0 }, + }); + + assert.equal((result as { isError?: boolean }).isError, true); + assert.match(textOf(result), /did not accept/); +}); + +test("a stalled batch status read is bounded too", async (t) => { + const client = await clientFor(sessionThatHangs(), t); + + const result = await client.callTool({ + name: "memwal_remember_status", + arguments: { job_ids: ["job-1", "job-2"], waitMs: 0 }, + }); + + assert.equal((result as { isError?: boolean }).isError, true); + assert.match(textOf(result), /did not accept/); +}); + +/** + * The status tool must answer before the MCP client gives up on it. + * + * `@modelcontextprotocol/sdk` times a request out at + * `DEFAULT_REQUEST_TIMEOUT_MSEC` (60s) unless the caller overrides it. A tool + * whose advertised maximum equals that deadline loses every race it enters: + * the client reports `MCP error -32001: Request timed out` and the agent + * cannot tell a slow write from a broken tool. Caught live against the + * production relayer with `waitMs: 60000` on a three-job batch. + */ +test("the status wait ceiling stays under the MCP client's own deadline", async (t: TestContext) => { + const { DEFAULT_REQUEST_TIMEOUT_MSEC } = await import( + "@modelcontextprotocol/sdk/shared/protocol.js" + ); + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair(); + const server = createMcpServer({ + oauthScope: "memwal:read memwal:write", + } as MemWalSession); + const client = new Client({ name: "status-deadline-test", version: "1.0.0" }); + t.after(async () => { + await client.close(); + await server.close(); + }); + await server.connect(serverTransport); + await client.connect(clientTransport); + + const { tools } = await client.listTools(); + const status = tools.find((tool) => tool.name === "memwal_remember_status"); + assert.ok(status, "memwal_remember_status is registered"); + + const max = ( + status.inputSchema as { properties?: { waitMs?: { maximum?: number } } } + ).properties?.waitMs?.maximum; + assert.equal(typeof max, "number", "waitMs advertises a maximum"); + assert.ok( + max < DEFAULT_REQUEST_TIMEOUT_MSEC, + `waitMs max ${max}ms must stay under the client deadline ${DEFAULT_REQUEST_TIMEOUT_MSEC}ms`, + ); + // Headroom for the round trip, not just a strict inequality. + assert.ok( + DEFAULT_REQUEST_TIMEOUT_MSEC - max >= 10_000, + `only ${DEFAULT_REQUEST_TIMEOUT_MSEC - max}ms of headroom before the client gives up`, + ); +}); + +/** + * Rate limiting is the quietest way this system loses a memory: the relayer + * answers 429 with `retry_after_seconds`, and before this the tool surfaced it + * as a bare `Tool error` while the fact was simply never written. Observed + * live against production during benchmarking, four calls in a row. + */ + +function rejectsWith(status, serverCode, retryAfterSeconds, thenValue) { + let n = 0; + return async () => { + if (n++ === 0) { + const e = new Error(`stub ${status}`); + Object.assign(e, { status, serverCode, retryAfterSeconds }); + throw e; + } + return thenValue; + }; +} + +function sessionRejecting(fn) { + return { + oauthScope: "memwal:read memwal:write", + namespace: "default", + memwal: { rememberAsync: fn, rememberBulkAsync: fn }, + } as unknown as MemWalSession; +} + +test("a short cooldown is absorbed and the write still lands", async (t) => { + // 503 AUTH_UPSTREAM_UNAVAILABLE advises ~5s; inside the absorb budget, so + // the agent should never see it. + // + // A sub-second cooldown here only because this file shortens + // ACCEPT_DEADLINE_MS to 150ms: the retry sleeps INSIDE the accept deadline, + // so the wait has to fit within it. In production that is 8s of absorb + // inside a 15s deadline — see the invariant pinned below. + const client = await clientFor( + sessionRejecting(rejectsWith(503, "AUTH_UPSTREAM_UNAVAILABLE", 0.05, { job_id: "job-1", status: "pending" })), + t, + ); + const res = await client.callTool({ name: "memwal_remember", arguments: { text: "a fact" } }); + assert.notEqual((res as { isError?: boolean }).isError, true); + assert.match(textOf(res), /job_id=job-1/); +}); + +test("a 60s rate-limit cooldown is reported, not slept through", async (t) => { + // Sleeping 60s inside a tool call is the hang this work exists to remove, + // and the MCP client would time out first. + const client = await clientFor( + sessionRejecting(rejectsWith(429, "Rate limit exceeded", 60, { job_id: "nope", status: "pending" })), + t, + ); + const started = Date.now(); + const res = await client.callTool({ name: "memwal_remember", arguments: { text: "a fact" } }); + assert.ok(Date.now() - started < 5_000, "must not sit on a 60s cooldown"); + assert.equal((res as { isError?: boolean }).isError, true); + const text = textOf(res); + // The agent must not tell the user this is being saved. + assert.match(text, /NOT\s+SAVED/i); + assert.match(text, /60s|about 60/); + // And it should say how to stop burning the budget. + assert.match(text, /memwal_remember_bulk/); +}); + +test("a rate-limited batch says the facts were not saved", async (t) => { + const client = await clientFor( + sessionRejecting(rejectsWith(429, "Rate limit exceeded", 60, { job_ids: [], total: 0, status: "x" })), + t, + ); + const res = await client.callTool({ name: "memwal_remember_bulk", arguments: { facts: ["a", "b"] } }); + assert.equal((res as { isError?: boolean }).isError, true); + assert.match(textOf(res), /NOT\s+SAVED/i); +}); + +/** + * The relayer runs three limiter layers with different windows — + * `delegate_key` and `account_burst` per minute, `account_sustained` per HOUR + * (services/server/src/rate_limit.rs). The message used to assert "the limit is + * per delegate key and resets in about Ns; retry after that" for all of them. + * + * Both halves were wrong on the hourly layer, and an agent that believed them + * waited out the advised window and retried into the same denial. Measured on + * dev 2026-09-17: a 429 carrying `account_sustained` / `1000 + * weighted-requests/hour` / `retry_after_seconds: 300`, retried after 300s and + * again after 420s, denied both times. + */ +function rejectsWithRateLimitBody(layer: string, limit: string, retryAfterSeconds: number) { + return async () => { + const e = new Error( + `Walrus Memory server error (429): ` + + JSON.stringify({ + error: "Rate limit exceeded", + layer, + limit, + retry_after_seconds: retryAfterSeconds, + }), + ); + Object.assign(e, { status: 429, serverCode: "Rate limit exceeded", retryAfterSeconds }); + throw e; + }; +} + +test("a rate-limit message names the layer the relayer actually denied on", async (t) => { + const client = await clientFor( + sessionRejecting( + rejectsWithRateLimitBody("account_sustained", "1000 weighted-requests/hour", 300), + ), + t, + ); + const res = await client.callTool({ name: "memwal_remember", arguments: { text: "a fact" } }); + assert.equal((res as { isError?: boolean }).isError, true); + const text = textOf(res); + + assert.match(text, /NOT\s+SAVED/i); + assert.match(text, /account_sustained/); + assert.match(text, /1000 weighted-requests\/hour/); + // It is the account's budget, not the delegate key's. + assert.match(text, /per account/); + assert.doesNotMatch(text, /per delegate key/); + // And the advised number must not be sold as a reset. + assert.match(text, /not a reset/i); + assert.doesNotMatch(text, /resets in/i); +}); + +test("a per-minute layer keeps the plain retry advice", async (t) => { + const client = await clientFor( + sessionRejecting(rejectsWithRateLimitBody("delegate_key", "60 weighted-requests/min", 60)), + t, + ); + const res = await client.callTool({ name: "memwal_remember", arguments: { text: "a fact" } }); + const text = textOf(res); + assert.match(text, /delegate_key/); + assert.match(text, /per delegate key/); + assert.match(text, /Retry after ~60s/); + // The hourly caveat belongs only to the hourly layer. + assert.doesNotMatch(text, /not a reset/i); +}); + +test("a 429 with no parsable body still reports the limit", async (t) => { + const client = await clientFor( + sessionRejecting(rejectsWith(429, "Rate limit exceeded", 60, { job_id: "nope", status: "pending" })), + t, + ); + const res = await client.callTool({ name: "memwal_remember", arguments: { text: "a fact" } }); + const text = textOf(res); + assert.match(text, /NOT\s+SAVED/i); + assert.match(text, /Retry after ~60s/); + // No body means no layer to name — it must not invent one. + assert.doesNotMatch(text, /Limit hit:/); +}); + +/** + * `memwal_analyze`'s extraction leg runs the extractor LLM inline, so it is the + * one leg long enough to race the MCP client's own ceiling. Budgeted level with + * that ceiling it lost the race, and `analyze` carries no idempotency key, so + * the caller's retry re-extracted and re-stored every fact. + */ +test("the analyze extraction deadline fits under the client ceiling", async () => { + const { DEFAULT_REQUEST_TIMEOUT_MSEC } = await import( + "@modelcontextprotocol/sdk/shared/protocol.js" + ); + const { ANALYZE_EXTRACTION_DEADLINE_MS, MCP_CLIENT_DEFAULT_TIMEOUT_MS } = await import( + "../tools/remember-wait.js" + ); + + // The constant we derive from must be the one the SDK actually applies. + assert.equal(MCP_CLIENT_DEFAULT_TIMEOUT_MS, DEFAULT_REQUEST_TIMEOUT_MSEC); + assert.ok( + ANALYZE_EXTRACTION_DEADLINE_MS < DEFAULT_REQUEST_TIMEOUT_MSEC, + `extraction ${ANALYZE_EXTRACTION_DEADLINE_MS}ms must finish before the client gives up ` + + `at ${DEFAULT_REQUEST_TIMEOUT_MSEC}ms`, + ); + assert.ok( + DEFAULT_REQUEST_TIMEOUT_MSEC - ANALYZE_EXTRACTION_DEADLINE_MS >= 10_000, + `only ${DEFAULT_REQUEST_TIMEOUT_MSEC - ANALYZE_EXTRACTION_DEADLINE_MS}ms of headroom`, + ); +}); + +test("a non-retryable error is not retried", async (t) => { + // A 500 could have been thrown after a write started; retrying bulk there + // would store every fact twice. + let calls = 0; + const client = await clientFor( + sessionRejecting(async () => { calls++; const e = new Error("boom"); Object.assign(e, { status: 500 }); throw e; }), + t, + ); + const res = await client.callTool({ name: "memwal_remember_bulk", arguments: { facts: ["a"] } }); + assert.equal((res as { isError?: boolean }).isError, true); + assert.equal(calls, 1, "a 500 must not be retried"); +}); + +test("the absorb budget fits inside the accept deadline", () => { + // withAcceptDeadline wraps withRelayerRetry, so a cooldown we choose to sit + // on is spent against the accept deadline. If the absorb budget ever grew + // past it, every absorbed retry would be cut off mid-wait and surface as + // "did not accept" instead of succeeding. + assert.ok( + MAX_ABSORBED_COOLDOWN_MS < DEFAULT_ACCEPT_DEADLINE_MS, + `absorb ${MAX_ABSORBED_COOLDOWN_MS}ms must stay under the ${DEFAULT_ACCEPT_DEADLINE_MS}ms accept deadline`, + ); +}); diff --git a/services/server/scripts/mcp/__tests__/remember-fast-return.test.ts b/services/server/scripts/mcp/__tests__/remember-fast-return.test.ts new file mode 100644 index 000000000..e54b2e881 --- /dev/null +++ b/services/server/scripts/mcp/__tests__/remember-fast-return.test.ts @@ -0,0 +1,203 @@ +// Pin the wait at 0 so this file reads the product default even if the +// process already had another MEMWAL_MCP_REMEMBER_WAIT_MS. Default is return +// at accept; 90000 restores block-until-done. +process.env.MEMWAL_MCP_REMEMBER_WAIT_MS = "0"; + +import assert from "node:assert/strict"; +import test, { type TestContext } from "node:test"; +import { Client } from "@modelcontextprotocol/sdk/client/index.js"; +import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js"; +import type { MemWalSession } from "../auth.js"; + +// Dynamic: static imports hoist above the assignment above, so the module +// would read the real default before the line that overrides it ever runs. +const { createMcpServer } = await import("../server.js"); +const { REMEMBER_WAIT_MS, parseWaitBudget, MAX_REMEMBER_WAIT_MS } = + await import("../tools/remember-wait.js"); + +/** + * Default wait is 0: `memwal_remember` returns at accept with a job_id. + * These tests pin what keeps that honest: an accepted result must never read + * as saved or carry a blob_id, and `memwal_remember_status` — the only thing + * that can observe a job failing after acceptance — must report that failure + * as an error rather than as a write still in flight. + */ + +interface FakeJob { + /** Calls to waitForRememberJob before the job reports done. */ + pendingPolls: number; + blobId?: string; + /** When set, the job fails with this message instead of completing. */ + failWith?: string; +} + +function sessionWith(job: FakeJob, calls: string[] = []): MemWalSession { + let polls = 0; + return { + oauthScope: "memwal:read memwal:write", + memwal: { + async rememberAsync(text: string, namespace?: string) { + calls.push(`rememberAsync:${text}:${namespace ?? ""}`); + return { job_id: "job-1", status: "running" }; + }, + async waitForRememberJob(jobId: string) { + calls.push(`wait:${jobId}`); + if (job.failWith) { + throw Object.assign( + new Error(`remember job failed: ${job.failWith}`), + { status: 500, jobId } + ); + } + if (polls++ < job.pendingPolls) { + throw Object.assign( + new Error(`remember job timed out (job_id=${jobId})`), + { status: 504, jobId } + ); + } + return { + id: jobId, + job_id: jobId, + blob_id: job.blobId ?? "blob-abc", + owner: "0xowner", + namespace: "default", + }; + }, + async getRememberStatus(jobId: string) { + calls.push(`status:${jobId}`); + if (job.failWith) { + return { job_id: jobId, status: "failed", error: job.failWith }; + } + return polls++ < job.pendingPolls + ? { job_id: jobId, status: "running" } + : { + job_id: jobId, + status: "done", + blob_id: job.blobId ?? "blob-abc", + namespace: "default", + }; + }, + }, + } as unknown as MemWalSession; +} + +async function clientFor(session: MemWalSession, t: TestContext): Promise { + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair(); + const server = createMcpServer(session); + const client = new Client({ name: "remember-fast-return-test", version: "1.0.0" }); + t.after(async () => { + await client.close(); + await server.close(); + }); + await server.connect(serverTransport); + await client.connect(clientTransport); + return client; +} + +function textOf(result: unknown): string { + return (result as { content: Array<{ text: string }> }).content + .map((c) => c.text) + .join("\n"); +} + +test("the default returns at accept, and blocking is opt-in", () => { + // Assignment above is so this module reads 0 even if the process already + // had another value. The product default is 0. + assert.equal(parseWaitBudget(undefined), 0); + assert.equal(parseWaitBudget(""), 0); + assert.equal(REMEMBER_WAIT_MS, 0); + assert.equal(parseWaitBudget("0"), 0); + // An operator who wants the old behaviour still has it. + assert.equal(parseWaitBudget("90000"), MAX_REMEMBER_WAIT_MS); +}); + +test("a typo'd budget falls back to the default instead of picking one nobody asked for", () => { + // Number("10s") is NaN, and every NaN comparison is false — an unvalidated + // parse would sail past a range check. + assert.equal(parseWaitBudget("10s"), 0); + assert.equal(parseWaitBudget("abc"), 0); + assert.equal(parseWaitBudget("-1"), 0); +}); + +test("a budget past the ceiling is clamped, not honoured", () => { + assert.equal(parseWaitBudget("90000"), 90_000); + assert.equal(parseWaitBudget("600000"), 90_000); +}); + +test("memwal_remember returns at accept and never reads as saved", async (t) => { + const calls: string[] = []; + const client = await clientFor(sessionWith({ pendingPolls: 0 }, calls), t); + const result = await client.callTool({ + name: "memwal_remember", + arguments: { text: "a durable fact" }, + }); + + // Accepted is the expected path, not a failure — the job is a durable row + // the relayer drives, and keeps going after the tool returns. + assert.equal(result.isError, undefined); + + const text = textOf(result); + assert.match(text, /ACCEPTED, NOT YET SAVED/); + assert.match(text, /job_id=job-1/); + assert.match(text, /memwal_remember_status/); + // The whole point of the wording: an agent must not be able to read this + // as a completed write. + assert.doesNotMatch(text, /Saved to Walrus Memory/); + assert.ok(!text.includes("blob_id="), "an accepted result must not carry a blob_id"); + + // A zero budget must not poll at all — the job was accepted, and waiting + // zero milliseconds for it is not a thing worth a round trip. + assert.deepEqual(calls, ["rememberAsync:a durable fact:"]); +}); + +test("the accepted message does not claim a duration it never waited", async (t) => { + const client = await clientFor(sessionWith({ pendingPolls: 0 }), t); + const result = await client.callTool({ + name: "memwal_remember", + arguments: { text: "a fact" }, + }); + assert.doesNotMatch(textOf(result), /after 0\.0s/); +}); + +test("memwal_remember_status reports the blob_id once the job lands", async (t) => { + const client = await clientFor(sessionWith({ pendingPolls: 0, blobId: "blob-late" }), t); + const result = await client.callTool({ + name: "memwal_remember_status", + arguments: { job_id: "job-1" }, + }); + + assert.equal(result.isError, undefined); + const text = textOf(result); + assert.match(text, /Saved to Walrus Memory/); + assert.match(text, /blob_id=blob-late/); +}); + +test("memwal_remember_status with waitMs=0 reads state without waiting", async (t) => { + const calls: string[] = []; + const client = await clientFor(sessionWith({ pendingPolls: 5 }, calls), t); + const result = await client.callTool({ + name: "memwal_remember_status", + arguments: { job_id: "job-1", waitMs: 0 }, + }); + + assert.equal(result.isError, undefined); + assert.match(textOf(result), /STILL UPLOADING/); + // A zero budget must be a single GET — waitForRememberJob sleeps before + // its first poll, so routing it there would report "still running" + // without ever asking the relayer. + assert.deepEqual(calls, ["status:job-1"]); +}); + +test("memwal_remember_status surfaces a failed job as an error", async (t) => { + const client = await clientFor( + sessionWith({ pendingPolls: 0, failWith: "walrus upload rejected" }, []), + t + ); + const result = await client.callTool({ + name: "memwal_remember_status", + arguments: { job_id: "job-1", waitMs: 0 }, + }); + + assert.equal(result.isError, true); + assert.match(textOf(result), /Walrus Memory job failed/); + assert.match(textOf(result), /walrus upload rejected/); +}); diff --git a/services/server/scripts/mcp/__tests__/remember-wait-hang.test.ts b/services/server/scripts/mcp/__tests__/remember-wait-hang.test.ts new file mode 100644 index 000000000..82ca4ccb1 --- /dev/null +++ b/services/server/scripts/mcp/__tests__/remember-wait-hang.test.ts @@ -0,0 +1,160 @@ +// The wait path only exists when there IS a wait, so this file keeps the +// default budget rather than the accept-and-continue knob the deadline tests +// use. Without that, `memwal_remember` returns at accept and never reaches the +// code under test — which is exactly how the first version of this test came +// to pass against the bug it was written for. +process.env.MEMWAL_MCP_REMEMBER_WAIT_MS = "5000"; + +import assert from "node:assert/strict"; +import test, { type TestContext } from "node:test"; +import { Client } from "@modelcontextprotocol/sdk/client/index.js"; +import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js"; +import type { MemWalSession } from "../auth.js"; + +const { createMcpServer } = await import("../server.js"); +const { REMEMBER_WAIT_MS } = await import("../tools/remember-wait.js"); + +function textOf(result: unknown): string { + return (result as { content: Array<{ text: string }> }).content + .map((c) => c.text) + .join("\n"); +} + +async function clientFor(session: MemWalSession, t: TestContext): Promise { + const [ct, st] = InMemoryTransport.createLinkedPair(); + const server = createMcpServer(session); + const client = new Client({ name: "wait-hang", version: "1.0.0" }); + t.after(async () => { await client.close(); await server.close(); }); + await server.connect(st); + await client.connect(ct); + return client; +} + +test("this file actually exercises the wait path", () => { + // Guards the mistake above: if the budget is ever 0 here, every test below + // passes without running the code it targets. + assert.notEqual(REMEMBER_WAIT_MS, 0); +}); + +test("a relayer that goes quiet mid-wait still hands back the job_id", async (t) => { + // The write was accepted — it is a row in remember_jobs and still running. + // Our wait deadline firing means the relayer stopped answering US. The + // job_id is the only way to settle it, and carrying it is the entire + // reason the pending result exists; the deadline error used to fall + // through to `throw` and discard it. + const client = await clientFor({ + oauthScope: "memwal:read memwal:write", + namespace: "default", + memwal: { + async rememberAsync() { + return { job_id: "job-live", status: "pending" }; + }, + async waitForRememberJob() { + // What withWaitDeadline raises once the relayer goes silent. + const err = new Error("Walrus Memory stopped responding while waiting"); + err.name = "MemWalRelayerUnresponsive"; + throw err; + }, + }, + } as unknown as MemWalSession, t); + + const res = await client.callTool({ + name: "memwal_remember", + arguments: { text: "a durable fact" }, + }); + const text = textOf(res); + + assert.match(text, /job_id=job-live/, `job_id was dropped: ${text}`); + assert.doesNotMatch(text, /^Saved to Walrus Memory/m, "must not read as stored"); +}); + +test("a job that genuinely failed is still an error, not a pending result", async (t) => { + // The counterpart: a 500 means the write is dead, so it must NOT be + // laundered into "still uploading" by the same branch. + const client = await clientFor({ + oauthScope: "memwal:read memwal:write", + namespace: "default", + memwal: { + async rememberAsync() { + return { job_id: "job-dead", status: "pending" }; + }, + async waitForRememberJob() { + const err = new Error("remember job failed: walrus upload rejected"); + Object.assign(err, { status: 500 }); + throw err; + }, + }, + } as unknown as MemWalSession, t); + + const res = await client.callTool({ + name: "memwal_remember", + arguments: { text: "a durable fact" }, + }); + assert.equal((res as { isError?: boolean }).isError, true); + assert.match(textOf(res), /failed/i); +}); + +test("a relayer that goes quiet mid-wait still hands back every bulk job_id", async (t) => { + // Same hole as the single-write case, on the path where it costs most: a + // batch loses N job_ids at once, and `waitForRememberJobs` reporting + // stragglers per item does not help — the throw comes from the deadline + // wrapper around it, before any per-item result exists. + const client = await clientFor({ + oauthScope: "memwal:read memwal:write", + namespace: "default", + memwal: { + async rememberBulkAsync() { + return { job_ids: ["bulk-1", "bulk-2"], total: 2, status: "accepted" }; + }, + async waitForRememberJobs() { + const err = new Error("Walrus Memory stopped responding while waiting"); + err.name = "MemWalRelayerUnresponsive"; + throw err; + }, + }, + } as unknown as MemWalSession, t); + + const text = textOf( + await client.callTool({ + name: "memwal_remember_bulk", + arguments: { facts: ["first fact", "second fact"] }, + }), + ); + assert.match(text, /bulk-1/, `job_ids were dropped: ${text}`); + assert.match(text, /bulk-2/, `job_ids were dropped: ${text}`); + assert.doesNotMatch(text, /^Saved \d+\/\d+/m, "must not read as stored"); +}); + +test("a relayer that goes quiet mid-wait keeps analyze's job_ids AND its facts", async (t) => { + // Analyze pays an LLM extraction before the writes are queued. Throwing + // away the wait discarded both the job_ids and that extraction, so the + // caller could neither settle the running writes nor recover the facts + // without paying for them again. + const client = await clientFor({ + oauthScope: "memwal:read memwal:write", + namespace: "default", + memwal: { + async analyze() { + return { + job_ids: ["an-1", "an-2"], + facts: [{ text: "drinks oat milk" }, { text: "ships on Fridays" }], + status: "accepted", + }; + }, + async waitForRememberJobs() { + const err = new Error("Walrus Memory stopped responding while waiting"); + err.name = "MemWalRelayerUnresponsive"; + throw err; + }, + }, + } as unknown as MemWalSession, t); + + const text = textOf( + await client.callTool({ + name: "memwal_analyze", + arguments: { text: "a transcript worth extracting from" }, + }), + ); + assert.match(text, /an-1/, `job_ids were dropped: ${text}`); + assert.match(text, /drinks oat milk/, `extracted facts were dropped: ${text}`); +}); diff --git a/services/server/scripts/mcp/__tests__/secret-redaction.test.ts b/services/server/scripts/mcp/__tests__/secret-redaction.test.ts new file mode 100644 index 000000000..3d86d30e5 --- /dev/null +++ b/services/server/scripts/mcp/__tests__/secret-redaction.test.ts @@ -0,0 +1,631 @@ +/** + * Unit coverage for the credential redactor (WALM-642). + * + * The write-path tests next door prove the three tools call this. These prove + * what it does, and — just as important — what it leaves alone: Walrus storage + * is append-only, so a false negative is permanent, but a false positive + * silently destroys the fact the user asked to keep. Both directions are pinned + * here. + */ +import assert from "node:assert/strict"; +import test from "node:test"; + +import { + sanitizeFact, + sanitizeFactBatch, + sanitizePassage, + redactionNotice, + refusalNotice, + droppedSpanNotice, + type RedactionKind, +} from "../tools/redaction.js"; + +/** Assert a secret is gone, the surrounding fact survived, and the kind is named. */ +function assertRedacted( + input: string, + secret: string, + kind: RedactionKind, + keeps: string[], +): string { + const out = sanitizeFact(input); + assert.equal(out.refusal, undefined, `unexpectedly refused: ${input}`); + assert.ok(out.changed, `nothing was redacted in: ${input}`); + assert.ok( + !out.text.includes(secret), + `the secret survived redaction (kind=${kind})`, + ); + assert.ok(out.kinds.includes(kind), `expected kind ${kind}, got ${out.kinds}`); + for (const keep of keeps) { + assert.ok(out.text.includes(keep), `lost "${keep}" from: ${input}`); + } + return out.text; +} + +// ── the shapes that must never reach storage ──────────────────────────────── + +test("a connection string keeps its host and loses its credentials", () => { + // The WALM-642 repro, almost verbatim: a preference stated next to a URL + // carrying a password. + const text = assertRedacted( + "I prefer dark mode, and the staging db is postgres://admin:hunter2@db.internal:5432/app", + "hunter2", + "url-credentials", + ["I prefer dark mode", "db.internal:5432/app", "postgres://"], + ); + assert.ok(!text.includes("admin:hunter2")); +}); + +test("vendor-prefixed API keys are recognised without any context", () => { + const cases: Array<[string, string]> = [ + ["sk-ant-api03-abcdefghijklmnopqrstuvwxyz0123456789", "sk-ant"], + ["sk-abcdefghijklmnopqrstuvwxyz0123", "openai"], + ["ghp_abcdefghijklmnopqrstuvwxyz0123456789", "github"], + ["github_pat_11ABCDEFG0abcdefghijklmnop", "github fine-grained"], + ["AKIAIOSFODNN7EXAMPLE", "aws"], + ["xoxb-1234567890-abcdefghij", "slack"], + ["glpat-abcdefghij0123456789", "gitlab"], + ]; + for (const [secret, label] of cases) { + assertRedacted( + `My deploy notes: the CI runner uses ${secret} for pushes`, + secret, + "vendor-api-key", + ["My deploy notes", "CI runner"], + ); + assert.ok(label); + } +}); + +test("a PEM private key is removed whole, terminated or not", () => { + const body = "MIIEowIBAAKCAQEAx7Vk9mJ0ZwQ3\nabcdefghijklmnopqrstuvwxyz0123456789\n"; + const closed = + `Deploy key for the box:\n-----BEGIN RSA PRIVATE KEY-----\n${body}-----END RSA PRIVATE KEY-----\nIt lives in 1Password.`; + const out = assertRedacted(closed, body.trim(), "private-key-block", ["Deploy key"]); + assert.ok(!out.includes("BEGIN RSA PRIVATE KEY")); + + // A paste that was cut off has no END line. The body must still go. + const truncated = `Deploy key for the box:\n-----BEGIN OPENSSH PRIVATE KEY-----\n${body}`; + const cut = sanitizeFact(truncated); + assert.ok(!cut.text.includes("MIIEowIBAAKCAQEAx7Vk9mJ0ZwQ3")); + assert.ok(cut.kinds.includes("private-key-block")); +}); + +test("a JWT is removed", () => { + const jwt = + "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.dBjftJeZ4CVPmB92K27uhbUJU1p1r_wW1gFWFOEjXk"; + assertRedacted( + `Our session tokens look like ${jwt} and expire hourly`, + jwt, + "jwt", + ["Our session tokens", "expire hourly"], + ); +}); + +test("credential assignments lose the value and keep the key name", () => { + for (const [line, secret] of [ + ["password=hunter2", "hunter2"], + ["api_key: abc123def456", "abc123def456"], + ['token = "t0ps3cr3t-value"', "t0ps3cr3t-value"], + ["client_secret:swordfish99", "swordfish99"], + ] as Array<[string, string]>) { + const out = assertRedacted( + `My local override file has ${line} and I never commit it`, + secret, + "credential-assignment", + ["local override file", "never commit it"], + ); + assert.ok(/redacted:credential-assignment/.test(out)); + } +}); + +test("authorization and cookie headers are removed", () => { + assertRedacted( + "To call the API: Authorization: Bearer abc123xyz789 — then GET /v1/me", + "abc123xyz789", + "auth-header", + ["To call the API", "/v1/me"], + ); + assertRedacted( + "The dashboard needs Cookie: session=9f8e7d6c5b4a3 to load my profile", + "9f8e7d6c5b4a3", + "auth-header", + ["The dashboard needs", "to load my profile"], + ); +}); + +test("a labelled seed phrase is removed", () => { + const words = + "abandon ability able about above absent absorb abstract absurd abuse access accident"; + assertRedacted( + `My wallet recovery phrase is ${words} and the wallet is on Sui mainnet`, + words, + "seed-phrase", + ["My wallet", "Sui mainnet"], + ); +}); + +test("MemWal's own delegate private key is removed, in every shape it arrives in", () => { + // The 64-hex Ed25519 seed from ~/.memwal/credentials.json. auth.ts marks it + // "NEVER log this": whoever holds it can read and write the user's memories + // until the delegate is revoked. It is pure lowercase hex, so the entropy + // rule deliberately does not see it — the LABEL is what catches it. + const SEED = "4f3c2b1a9e8d7c6b5a4f3e2d1c0b9a8f7e6d5c4b3a2f1e0d9c8b7a6f5e4d3c2b"; + + // A pasted line from the file, the file itself, prose, and the label after + // the value. Each keeps a fact around it so the refusal path is not what is + // being measured here. + for (const text of [ + `Notes from setup: delegatePrivateKey ${SEED} was written on this laptop`, + `Notes from setup: "delegatePrivateKey": "${SEED}" is in the file`, + `Notes from setup: my delegate private key is ${SEED} on this laptop`, + `Notes from setup: ${SEED} is my private key for this laptop`, + `Notes from setup: secret_key = ${SEED} on this laptop`, + `Notes from setup: the signing key 0x${SEED} lives on this laptop`, + ]) { + const out = sanitizeFact(text); + assert.equal(out.refusal, undefined, `unexpectedly refused: ${text}`); + assert.ok(!out.text.includes(SEED), `the delegate key survived: ${text}`); + assert.ok(out.text.includes("Notes from setup"), `lost the fact: ${text}`); + } +}); + +test("an UNLABELLED hex run is still left alone — that is what the label gate buys", () => { + // The regression this design protects. Same 64 hex characters as the test + // above; the only difference is that nothing calls them a key. + const HEX = "4f3c2b1a9e8d7c6b5a4f3e2d1c0b9a8f7e6d5c4b3a2f1e0d9c8b7a6f5e4d3c2b"; + for (const fact of [ + `The release digest is ${HEX} and I verified it`, + `Pin the deployment to 0x${HEX}`, + // "credentials.json" is a phrase MemWal prose uses constantly, and it + // sits within a window of the SHA here. It is excluded from the label + // list precisely so this sentence keeps its commit id. + "My creds live in ~/.memwal/credentials.json and the fix landed in 4f2b8c1e9d7a3f5b6c0e2d4a8b1f3c5e7d9a0b2c", + // Documented in auth.ts as "Safe to display" — the public half must not + // be swept up with the private one. + `My delegatePublicKeyHex is ${HEX}`, + ]) { + const out = sanitizeFact(fact); + assert.equal(out.text, fact, `redacted an unlabelled hex run: ${fact}`); + assert.equal(out.changed, false); + } +}); + +test("a seed phrase is caught however the label is spelled", () => { + const WORDS = + "abandon ability able about above absent absorb abstract absurd abuse access accident"; + for (const text of [ + `Wallet notes: my recovery phrase is ${WORDS} for the mainnet wallet`, + `Wallet notes: seed_phrase: ${WORDS} for the mainnet wallet`, + `Wallet notes: "mnemonic": "${WORDS}" for the mainnet wallet`, + `Wallet notes: seedPhrase=${WORDS} for the mainnet wallet`, + ]) { + const out = sanitizeFact(text); + assert.equal(out.refusal, undefined, `unexpectedly refused: ${text}`); + assert.ok(!out.text.includes(WORDS), `the mnemonic survived: ${text}`); + assert.ok(out.text.includes("Wallet notes"), `lost the fact: ${text}`); + } + + // Still true, and still documented: a bare word run with no label at all is + // indistinguishable from a sentence, so it is left to the model rules. + const bare = sanitizeFact(`I wrote down ${WORDS} yesterday`); + assert.equal(bare.changed, false); +}); + +test("a long mixed-case base64 blob is removed", () => { + const blob = + "QWxhZGRpbjpvcGVuIHNlc2FtZQBcdefGHIjklMNOpqrSTUvwxYZ0123456789abcDEF0123"; + assertRedacted( + `The signing material is ${blob} which I keep in the vault`, + blob, + "high-entropy-secret", + ["The signing material", "keep in the vault"], + ); +}); + +// ── the shapes that must survive untouched ────────────────────────────────── + +test("a plain preference passes through byte-for-byte", () => { + // The single most important assertion in this file: the ordinary case must + // be indistinguishable from having no redactor at all. + for (const fact of [ + "I always use pnpm, and TypeScript strict mode on every project.", + "Tui luôn dùng pnpm và order cafe là matcha oat latte.", + "Deploy to staging on Thursdays, never on Friday afternoons.", + "My password manager is 1Password and I rotate keys every quarter.", + "The API key for that service is stored in Vault, not in the repo.", + ]) { + const out = sanitizeFact(fact); + assert.equal(out.text, fact, `changed a clean fact: ${fact}`); + assert.equal(out.changed, false); + assert.equal(out.count, 0); + assert.deepEqual(out.kinds, []); + assert.equal(out.refusal, undefined); + } +}); + +test("the identifiers MemWal itself stores are not mistaken for secrets", () => { + // A generic entropy rule would eat every one of these, which is why there + // isn't one. See the trade-off note at the top of redaction.ts. + for (const fact of [ + "My account id is 0x7f3a9c2e5b8d1f4a6c9e2b5d8f1a4c7e0b3d6f9a2c5e8b1d4f7a0c3e6b9d2f5a", + "The blob landed as blob_id=Xj9vKq2mP7nR4tW8yB1cE5gH0dF3sA6uZ2xN8qL4kM7", + "Pin the build to commit 4f2b8c1e9d7a3f5b6c0e2d4a8b1f3c5e7d9a0b2c", + "The sha256 of the release tarball is e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", + "My Sui package id is 0xe80f2feec1c139616a86c9f71210152e2a7ca552b20841f2e192f99f75864437", + "The migration artifact hashes to 9a8b7c6d5e4f3a2b1c0d9e8f7a6b5c4d3e2f1a0b9c8d7e6f5a4b3c2d1e0f9a8b", + ]) { + const out = sanitizeFact(fact); + assert.equal(out.text, fact, `redacted a legitimate identifier: ${fact}`); + assert.equal(out.changed, false); + } +}); + +// ── refusals ──────────────────────────────────────────────────────────────── + +test("an explicit do-not-save is honoured", () => { + for (const text of [ + "My bank PIN is 4821 — don't save this.", + "Do not remember this: I'm interviewing elsewhere.", + "Off the record, I'm leaving the team in March.", + "Please don't store that anywhere.", + ]) { + const out = sanitizeFact(text); + assert.equal(out.refusal, "no-save-directive", `not refused: ${text}`); + assert.equal(out.text, ""); + } +}); + +test("an ordinary preference about not saving things is NOT a do-not-save", () => { + // The demonstrative ("this"/"that"/"it") is what separates the two, and + // without it the directive check would eat real preferences. + for (const text of [ + "I don't save screenshots to the Desktop, they go to ~/Pictures.", + "Never store build artifacts in the repo — use the cache.", + "Don't keep logs longer than 30 days on staging.", + ]) { + const out = sanitizeFact(text); + assert.equal(out.refusal, undefined, `wrongly refused: ${text}`); + assert.equal(out.text, text); + } +}); + +test("pasted third-party content is not saved as a fact about the user", () => { + const fenced = "```\nERROR 500 from vendor api\n at handler (index.js:42)\n```"; + assert.equal(sanitizeFact(fenced).refusal, "pasted-content"); + + const quotedBlock = "> their PM said the deadline slips\n> and the scope is unchanged"; + assert.equal(sanitizeFact(quotedBlock).refusal, "pasted-content"); + + const longQuote = `"${"the vendor's release note says the same thing again and again. ".repeat(5)}"`; + assert.ok(longQuote.length >= 200); + assert.equal(sanitizeFact(longQuote).refusal, "pasted-content"); +}); + +test("a short quoted fact is still saved — an agent quotes the user routinely", () => { + const quoted = '"I always use pnpm"'; + const out = sanitizeFact(quoted); + assert.equal(out.refusal, undefined); + assert.equal(out.text, quoted); +}); + +test("a text that is nothing but a secret is refused, not stored as a placeholder", () => { + const out = sanitizeFact("sk-abcdefghijklmnopqrstuvwxyz0123"); + assert.equal(out.refusal, "credential-only"); + assert.equal(out.text, ""); + assert.ok(out.kinds.includes("vendor-api-key")); +}); + +// ── what the caller is told ───────────────────────────────────────────────── + +test("the notices name the kind and never the value", () => { + const secret = "hunter2"; + const out = sanitizeFact(`db is postgres://admin:${secret}@db.internal/app for staging`); + const notice = redactionNotice(out.kinds, out.count); + assert.match(notice, /url-credentials/); + assert.match(notice, /do not re-send/i); + assert.ok(!notice.includes(secret), "the notice must not echo the secret"); + + const refusal = refusalNotice("credential-only"); + assert.match(refusal, /NOT SAVED/); + assert.match(refusal, /Do not retry this text/); +}); + +test("nothing removed means nothing said", () => { + assert.equal(redactionNotice([], 0), ""); +}); + +/* ─────────────────────────────────────────────────────────────────────────── + * Review findings on the WALM-642 branch. Each of these leaked, verbatim, + * through `sanitizeFact` before the fix beside it. + * ------------------------------------------------------------------------ */ + +// ── finding 3: the credential-assignment gate missed whole spellings ──────── + +test("SCREAMING_SNAKE and quoted-JSON credential names are assignments too", () => { + // `_` is a word character, so `\b` never fired between `_` and the + // keyword, and `_` is not `[a-z]` so the camelCase lookbehind did not + // either. That left the spelling credentials actually arrive in — a pasted + // env file or shell export — completely unguarded. + for (const [text, secret] of [ + ["POSTGRES_PASSWORD=hunter2", "hunter2"], + [ + "AWS_SECRET_ACCESS_KEY=wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + ], + ["DB_PASSWORD=hunter2SuperSecret", "hunter2SuperSecret"], + ["X_AUTH_TOKEN: abcd1234", "abcd1234"], + ["SESSION_SECRET=abcd1234efgh5678", "abcd1234efgh5678"], + ["my_api_key=abcd1234efgh5678", "abcd1234efgh5678"], + ] as Array<[string, string]>) { + const out = sanitizeFact(text); + assert.ok(out.changed, `not redacted at all: ${text}`); + assert.ok(!out.text.includes(secret), `the secret survived: ${text}`); + } + + // A JSON object puts a closing quote between the key and the colon, which + // the old separator group demanded come immediately after the keyword. + const json = sanitizeFact( + 'Save my config: {"username": "alice", "password": "hunter2-prod-9xQ"} for staging', + ); + assert.equal(json.refusal, undefined); + assert.ok(!json.text.includes("hunter2-prod-9xQ"), "the JSON password survived"); + assert.ok(json.text.includes("alice"), "the username is not a credential"); + assert.ok(json.text.includes("for staging"), "the fact was lost with the password"); +}); + +test("widening the gate did not widen it onto ordinary prose", () => { + // The lookbehind now admits `_`, `-` and digits. These are the sentences + // that must not start matching because of it. + for (const fact of [ + "My password manager is 1Password and I rotate keys every quarter.", + "The API key for that service is stored in Vault, not in the repo.", + "My creds live in ~/.memwal/credentials.json and the fix landed in 4f2b8c1e9a7d6f5c4b3a29180716253443219876", + "We renamed the secret_store module to vault_client last sprint.", + "Token bucket rate limiting is what the relayer uses.", + ]) { + const out = sanitizeFact(fact); + assert.equal(out.text, fact, `changed a clean fact: ${fact}`); + assert.equal(out.changed, false); + } +}); + +// ── finding 5: the label rule only ever looked at hex ─────────────────────── + +test("a labelled secret that is not hex is key material too", () => { + // 40 characters, base64-ish, three label words in front of it — and it came + // back completely unchanged, because `hasAdjacentCredentialLabel` was + // consulted only for HEX_RUN and the entropy rule demands 64+. + const AWS_SECRET = "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"; + assertRedacted( + `My AWS secret access key is ${AWS_SECRET} for prod`, + AWS_SECRET, + "labelled-key-material", + ["My AWS secret access key", "for prod"], + ); + + // The pair, which is the shape it arrives in: the non-secret id was being + // redacted by the vendor rule while the secret half survived. + const pair = sanitizeFact( + `AWS creds for staging: AKIAIOSFODNN7EXAMPLE and the secret key is ${AWS_SECRET}`, + ); + assert.ok(!pair.text.includes(AWS_SECRET), "the secret half of the pair survived"); + assert.ok(!pair.text.includes("AKIAIOSFODNN7EXAMPLE")); +}); + +test("the label is still the only discriminator — bare runs pass", () => { + // The whole design: widening the SHAPE the rule can see must not weaken the + // gate in front of it, or every identifier this product exists to remember + // starts disappearing. + for (const fact of [ + "The blob landed as blob_id=Xj9vKq2mP7nR4tW8yB1cE5gH0dF3sA6uZ2xN8qL4kM7", + "Pin the build to commit 4f2b8c1e9d7a3f5b6c0e2d4a8b1f3c5e7d9a0b2c", + "My Sui package id is 0xe80f2feec1c139616a86c9f71210152e2a7ca552b20841f2e192f99f75864437", + "My delegatePublicKeyHex is 4f3c2b1a9e8d7c6b5a4f3e2d1c0b9a8f7e6d5c4b3a2f1e0d9c8b7a6f5e4d3c2b and it is safe to share", + "The build id is 20260918T0930Z-linux-arm64-release-candidate", + ]) { + const out = sanitizeFact(fact); + assert.equal(out.text, fact, `redacted a legitimate identifier: ${fact}`); + assert.equal(out.changed, false); + } +}); + +// ── findings 6 + 7: the URL rule was quadratic, and mangled ordinary URLs ─── + +test("an ordinary URL with a query string is not userinfo", () => { + // The password group excluded `/` and whitespace but not `?` or `=`, so + // this matched with `app.example.com` as the user and `8443?owner=alice` as + // the password: host and port destroyed, `corp.com` promoted to hostname, + // and a notice claiming a credential had been removed when there was none. + const fact = + "Our dashboard is at https://app.example.com:8443?owner=alice@corp.com and we deploy Fridays"; + const out = sanitizeFact(fact); + assert.equal(out.text, fact, "an ordinary URL was mangled"); + assert.equal(out.changed, false); + assert.equal(out.count, 0, "a redaction was reported where none happened"); + + // ...while a real connection string is untouched by the narrowing. + const real = sanitizeFact("staging is postgres://admin:hunter2@db.internal:5432/app"); + assert.ok(!real.text.includes("hunter2")); + assert.ok(real.text.includes("db.internal:5432/app"), "the host was lost"); +}); + +test("a long passage is screened in linear time", () => { + // `URL_USERINFO`'s scheme repeat was unbounded, so it was tried and + // abandoned at every start offset: 30 KB took 317 ms, 60 KB 1254 ms and + // 120 KB 4814 ms. `memwal_analyze` takes a whole transcript, the sidecar is + // single-threaded and a tools/call times out at 60 s, so a long paste was a + // stall for every other caller too. + // + // The budget is deliberately loose — this is a guard against a quadratic + // pattern coming back, not a benchmark, and CI machines are noisy. The + // shape is what matters: 4x the input must not be 16x the time. + const worst = "a".repeat(120 * 1024); + const started = Date.now(); + sanitizeFact(worst); + const elapsed = Date.now() - started; + assert.ok( + elapsed < 1000, + `120 KB took ${elapsed} ms — a redaction pattern has gone superlinear again`, + ); + + const small = "a".repeat(30 * 1024); + const t0 = Date.now(); + sanitizeFact(small); + const smallMs = Math.max(Date.now() - t0, 1); + assert.ok( + elapsed / smallMs < 12, + `120 KB/30 KB ratio was ${(elapsed / smallMs).toFixed(1)}x — that is quadratic, not linear`, + ); +}); + +// ── finding 4: a secret split across bulk entries ────────────────────────── + +test("a credential split across two bulk entries is still caught", () => { + // The reporter's case. Every label-gated rule searches a window inside ONE + // string, so putting the label in one entry and the value in the next was + // a way past all of them — including the rule that exists specifically for + // MemWal's own delegate private key. + const SEED = "4f3c2b1a9e8d7c6b5a4f3e2d1c0b9a8f7e6d5c4b3a2f1e0d9c8b7a6f5e4d3c2b"; + const split = sanitizeFactBatch([ + "my delegate private key for the mainnet account", + SEED, + ]); + assert.ok( + !split.map((r) => r.text).join("\n").includes(SEED), + "the delegate private key survived being split across entries", + ); + assert.equal(split[1].refusal, "credential-only", "a bare key must be dropped, not stored"); + assert.equal(split[0].refusal, undefined, "the label entry is a fact and should survive"); + + // The same value with filler words around it, and with the label AFTER it. + for (const facts of [ + ["my delegate private key for the mainnet account", `the value to use is ${SEED} for that account`], + [SEED, "that is my delegate private key"], + ["I set up a second laptop", "my delegate private key is below", SEED], + ]) { + const out = sanitizeFactBatch(facts); + assert.ok( + !out.map((r) => r.text).join("\n").includes(SEED), + `the key survived: ${JSON.stringify(facts)}`, + ); + } +}); + +test("the batch screen does not fire without a label anywhere in it", () => { + // Cross-entry awareness widens what counts as adjacent, which is exactly + // the kind of change that starts eating identifiers. The gate is still the + // label: a batch with none must be byte-for-byte untouched. + const facts = [ + "The release digest is 4f3c2b1a9e8d7c6b5a4f3e2d1c0b9a8f7e6d5c4b3a2f1e0d9c8b7a6f5e4d3c2b", + "My Sui package id is 0xe80f2feec1c139616a86c9f71210152e2a7ca552b20841f2e192f99f75864437", + "The blob landed as blob_id=Xj9vKq2mP7nR4tW8yB1cE5gH0dF3sA6uZ2xN8qL4kM7", + "I always use pnpm", + ]; + const out = sanitizeFactBatch(facts); + assert.deepEqual(out.map((r) => r.text), facts); + assert.deepEqual(out.map((r) => r.changed), [false, false, false, false]); +}); + +test("per-entry refusal granularity survives the batch screen", () => { + // The property that already worked: one bad entry is dropped and the rest + // of the batch still lands. + const out = sanitizeFactBatch([ + "I always use pnpm", + "sk-abcdefghijklmnopqrstuvwxyz0123", + "Please don't save this one", + "Deploy on Thursdays", + ]); + assert.equal(out[0].refusal, undefined); + assert.equal(out[0].text, "I always use pnpm"); + assert.equal(out[1].refusal, "credential-only"); + assert.equal(out[2].refusal, "no-save-directive"); + assert.equal(out[3].text, "Deploy on Thursdays"); +}); + +// ── finding 8: one stray line discarded a whole transcript ───────────────── + +const TRANSCRIPT = [ + "user: I always use pnpm for every project", + "assistant: noted", + "user: my staging box is app.example.com:8443", + "assistant: got it", + "user: my bank PIN is 4821 - don't save this part", + "assistant: understood", + "user: I deploy on Thursdays and never on Fridays", + "assistant: makes sense", +].join("\n"); + +test("one do-not-save line drops its span, not the whole transcript", () => { + const out = sanitizePassage(TRANSCRIPT); + assert.equal(out.refusal, undefined, "the whole passage was discarded again"); + assert.ok(!out.text.includes("4821"), "the thing they asked not to save was kept"); + assert.ok(!out.text.includes("don't save this part")); + // Everything else survived, on both sides of the dropped span. + assert.ok(out.text.includes("I always use pnpm for every project")); + assert.ok(out.text.includes("app.example.com:8443")); + assert.ok(out.text.includes("I deploy on Thursdays")); + // And the caller is told what went, by position and reason — never by text. + assert.ok(out.dropped.length > 0); + assert.ok(out.dropped.every((d) => d.reason === "no-save-directive")); + const notice = droppedSpanNotice(out.dropped, out.segments); + assert.match(notice, /span\(s\) were dropped/); + assert.ok(!notice.includes("4821"), "the notice echoed the dropped content"); +}); + +test("a directive takes its own paragraph with it", () => { + // Scoping per line would save the PIN and drop only the sentence asking + // not to, which is the one outcome worse than dropping too much. + const out = sanitizePassage( + "I always use pnpm\n\nMy bank PIN is 4821\ndon't save this\n\nI deploy on Thursdays", + ); + assert.equal(out.refusal, undefined); + assert.ok(!out.text.includes("4821"), "the directive's neighbour was saved anyway"); + assert.ok(out.text.includes("I always use pnpm")); + assert.ok(out.text.includes("I deploy on Thursdays")); +}); + +test("a fenced transcript is analysed, a fenced code paste is still refused", () => { + // The tool documents a fenced transcript as its canonical input, and + // refused exactly that as pasted content. + const fenced = sanitizePassage("```\n" + TRANSCRIPT + "\n```"); + assert.equal(fenced.refusal, undefined, "a fenced transcript was discarded"); + assert.ok(fenced.text.includes("I always use pnpm for every project")); + assert.ok(!fenced.text.includes("4821")); + + // A short fence, or one carrying a language tag, is a paste and stays one. + assert.equal( + sanitizePassage("```\nERROR 500 from the vendor API\n at handler (index.js:42)\n```") + .refusal, + "pasted-content", + ); + assert.equal( + sanitizePassage('```json\n{\n "a": 1,\n "b": 2,\n "c": 3,\n "d": 4\n}\n```').refusal, + "pasted-content", + ); + // ...and a fence INSIDE a passage is dropped without taking the rest. + const mixed = sanitizePassage( + "I always use pnpm\n\n```\nERROR 500\n```\n\nI deploy on Thursdays", + ); + assert.equal(mixed.refusal, undefined); + assert.ok(mixed.text.includes("I always use pnpm")); + assert.ok(mixed.text.includes("I deploy on Thursdays")); + assert.ok(!mixed.text.includes("ERROR 500")); + assert.ok(mixed.dropped.some((d) => d.reason === "pasted-content")); +}); + +test("a passage with nothing left is still refused outright", () => { + // Scoping the refusal must not turn "save nothing" into "save something". + const out = sanitizePassage("My bank PIN is 4821 — don't save this."); + assert.equal(out.refusal, "no-save-directive"); + assert.equal(out.text, ""); + assert.ok(!out.text.includes("4821")); +}); + +test("a clean passage is forwarded byte-for-byte", () => { + const clean = "I always use pnpm, TypeScript strict mode, and deploy on Thursdays."; + const out = sanitizePassage(clean); + assert.equal(out.text, clean); + assert.equal(out.changed, false); + assert.deepEqual(out.dropped, []); + assert.equal(droppedSpanNotice(out.dropped, out.segments), ""); +}); diff --git a/services/server/scripts/mcp/__tests__/status-inflight-blob-id.test.ts b/services/server/scripts/mcp/__tests__/status-inflight-blob-id.test.ts new file mode 100644 index 000000000..0100bff5a --- /dev/null +++ b/services/server/scripts/mcp/__tests__/status-inflight-blob-id.test.ts @@ -0,0 +1,163 @@ +import assert from "node:assert/strict"; +import test, { type TestContext } from "node:test"; +import { Client } from "@modelcontextprotocol/sdk/client/index.js"; +import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js"; +import type { MemWalSession } from "../auth.js"; + +const { createMcpServer } = await import("../server.js"); + +/** + * An unfinished write must not be reported with a blob_id. + * + * `persist_uploaded_state` (services/server/src/jobs.rs) writes `blob_id` and + * sets status `uploaded` as soon as the blob is on Walrus — but the job is not + * done: SetMetadataAndTransfer has not handed the blob object to the user and + * insert_vector has not indexed it, so the fact is not recallable and a later + * failure can still lose it. `uploaded` is therefore counted as in flight. + * + * The render did not follow: it printed the blob_id of any row that had one, + * so a real reply on dev 2026-09-21 read + * + * 0/6 saved, 6 still uploading. + * 1. [still uploading] job_id=77e30990-… blob_id=M4jEdPGmnmGTPQMQ… + * + * which contradicts this server's own instruction that "only a blob_id in the + * tool reply means the fact is already stored". An agent that believes the + * blob_id tells the user a fact is saved while it is still in flight. + */ + +function sessionWith(opts: { + /** What the SDK's internal poll loop reports when its budget expires. */ + waitStates?: Array<"done" | "failed" | "timeout">; + /** What a direct batch read reports. */ + probeRows: Array<{ status: string; blob_id?: string }>; +}): MemWalSession { + const { waitStates, probeRows } = opts; + const idFor = (i: number) => `job-${i + 1}`; + return { + oauthScope: "memwal:read memwal:write", + namespace: "default", + memwal: { + async waitForRememberJobs(ids: string[]) { + const states = waitStates ?? ids.map(() => "timeout" as const); + return { + results: states.map((status, i) => ({ + id: idFor(i), + blob_id: status === "done" ? `blob-${i + 1}` : "", + status, + namespace: "default", + error: + status === "timeout" + ? "polling timed out after 10000ms" + : undefined, + })), + total: states.length, + succeeded: states.filter((s) => s === "done").length, + failed: states.filter((s) => s !== "done").length, + }; + }, + async getRememberBulkStatus(ids: string[]) { + return { + results: ids.map((id, i) => ({ + job_id: id, + status: probeRows[i]?.status ?? "running", + blob_id: probeRows[i]?.blob_id, + error: undefined, + })), + }; + }, + }, + } as unknown as MemWalSession; +} + +async function clientFor(session: MemWalSession, t: TestContext): Promise { + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair(); + const server = createMcpServer(session); + const client = new Client({ name: "status-inflight-blob-id-test", version: "1.0.0" }); + t.after(async () => { + await client.close(); + await server.close(); + }); + await server.connect(serverTransport); + await client.connect(clientTransport); + return client; +} + +function textOf(result: unknown): string { + return (result as { content: Array<{ text: string }> }).content + .map((c) => c.text) + .join("\n"); +} + +test("a waited batch whose probe finds `uploaded` shows no blob_id", async (t) => { + // Nothing moved inside the budget, so settleBatch confirms with one direct + // read — which is where an `uploaded` row (real blob_id, unfinished job) + // enters the report. + const client = await clientFor( + sessionWith({ + probeRows: [ + { status: "uploaded", blob_id: "blob-uploaded-1" }, + { status: "running" }, + ], + }), + t, + ); + + const text = textOf( + await client.callTool({ + name: "memwal_remember_status", + arguments: { job_ids: ["job-1", "job-2"] }, + }), + ); + + assert.match(text, /0\/2 saved, 2 still uploading/); + assert.ok( + !text.includes("blob-uploaded-1"), + `an in-flight row must not carry a blob_id:\n${text}`, + ); +}); + +test("a zero-budget read of an `uploaded` job shows no blob_id", async (t) => { + // The other path into the same render: waitMs=0 reads the relayer directly, + // so `uploaded` arrives with its blob_id and no probe is involved. + const client = await clientFor( + sessionWith({ probeRows: [{ status: "uploaded", blob_id: "blob-uploaded-1" }] }), + t, + ); + + const text = textOf( + await client.callTool({ + name: "memwal_remember_status", + arguments: { job_ids: ["job-1"], waitMs: 0 }, + }), + ); + + assert.match(text, /0\/1 saved, 1 still uploading/); + assert.ok( + !text.includes("blob-uploaded-1"), + `an in-flight row must not carry a blob_id:\n${text}`, + ); +}); + +test("a settled batch still shows every blob_id", async (t) => { + // The guard must not cost a caller the blob_ids it is entitled to: `done` + // is exactly the state the instruction text points at. + const client = await clientFor( + sessionWith({ + waitStates: ["done", "done"], + probeRows: [{ status: "done" }, { status: "done" }], + }), + t, + ); + + const text = textOf( + await client.callTool({ + name: "memwal_remember_status", + arguments: { job_ids: ["job-1", "job-2"] }, + }), + ); + + assert.match(text, /2\/2 saved/); + assert.ok(text.includes("blob-1"), text); + assert.ok(text.includes("blob-2"), text); +}); diff --git a/services/server/scripts/mcp/__tests__/status-rate-limit.test.ts b/services/server/scripts/mcp/__tests__/status-rate-limit.test.ts new file mode 100644 index 000000000..de7216e3d --- /dev/null +++ b/services/server/scripts/mcp/__tests__/status-rate-limit.test.ts @@ -0,0 +1,302 @@ +import assert from "node:assert/strict"; +import test, { type TestContext } from "node:test"; +import { Client } from "@modelcontextprotocol/sdk/client/index.js"; +import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js"; +import type { MemWalSession } from "../auth.js"; + +const { createMcpServer } = await import("../server.js"); + +/** + * A rate-limited poll must not read as "still uploading". + * + * `waitForRememberJobs` runs its own poll loop and swallows each poll's error, + * stamping every row `timeout` when the budget runs out. So a batch whose + * polls were all REFUSED with 429 came back looking exactly like a batch that + * was genuinely still uploading — and the agent, told it was progressing, + * polled again on a budget it had already spent. + * + * Observed on dev 2026-09-17: ~10 minutes of "0/N saved, N still uploading" + * while every `/api/remember/bulk/status` was being denied. + */ + +interface Behaviour { + /** What the SDK's internal poll loop reports when its budget expires. */ + waitStates: Array<"done" | "failed" | "timeout">; + /** How the confirming direct read behaves. */ + probe: "rate-limited" | "still-running" | "done" | "throws" | "partial"; +} + +function rateLimit429(): Error { + const e = new Error( + 'Walrus Memory server error (429): {"error":"Rate limit exceeded",' + + '"layer":"account_sustained","limit":"1000 weighted-requests/hour",' + + '"retry_after_seconds":300}', + ); + (e as Error & { status?: number; retryAfterSeconds?: number }).status = 429; + (e as Error & { retryAfterSeconds?: number }).retryAfterSeconds = 300; + return e; +} + +function sessionWith(b: Behaviour, calls: string[] = []): MemWalSession { + const jobIds = b.waitStates.map((_, i) => `job-${i + 1}`); + return { + oauthScope: "memwal:read memwal:write", + namespace: "default", + memwal: { + async waitForRememberJobs(ids: string[]) { + calls.push(`waitJobs:${ids.join(",")}`); + return { + results: b.waitStates.map((status, i) => ({ + id: jobIds[i], + blob_id: status === "done" ? `blob-${i + 1}` : "", + status, + namespace: "default", + error: + status === "timeout" + ? "polling timed out after 30000ms" + : status === "failed" + ? "walrus upload rejected" + : undefined, + })), + total: b.waitStates.length, + succeeded: b.waitStates.filter((s) => s === "done").length, + failed: b.waitStates.filter((s) => s !== "done").length, + }; + }, + async getRememberBulkStatus(ids: string[]) { + calls.push(`bulkStatus:${ids.join(",")}`); + if (b.probe === "rate-limited") throw rateLimit429(); + if (b.probe === "throws") throw new Error("relayer unreachable"); + // The relayer is not obliged to echo back one row per id it was + // asked about: a job it cannot find is simply absent from + // `results`. Answer for the first id only. + const answered = b.probe === "partial" ? ids.slice(0, 1) : ids; + return { + results: answered.map((id, i) => ({ + job_id: id, + status: b.probe === "done" ? "done" : "running", + blob_id: b.probe === "done" ? `blob-${i + 1}` : undefined, + error: undefined, + })), + }; + }, + }, + } as unknown as MemWalSession; +} + +async function clientFor(session: MemWalSession, t: TestContext): Promise { + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair(); + const server = createMcpServer(session); + const client = new Client({ name: "status-rate-limit-test", version: "1.0.0" }); + t.after(async () => { + await client.close(); + await server.close(); + }); + await server.connect(serverTransport); + await client.connect(clientTransport); + return client; +} + +function textOf(result: unknown): string { + return (result as { content: Array<{ text: string }> }).content + .map((c) => c.text) + .join("\n"); +} + +test("a batch whose polls were all rate-limited reports the limit, not progress", async (t) => { + const calls: string[] = []; + const client = await clientFor( + sessionWith({ waitStates: ["timeout", "timeout"], probe: "rate-limited" }, calls), + t, + ); + + const result = await client.callTool({ + name: "memwal_remember_status", + arguments: { job_ids: ["job-1", "job-2"], waitMs: 30000 }, + }); + const text = textOf(result); + + // The whole point: the words that sent the agent back to poll must be gone. + assert.ok( + !/still uploading/i.test(text), + `a refused poll still read as progress: ${text}`, + ); + assert.match(text, /rate limit/i); + // And it must name the layer it actually hit — the hourly one is not the + // per-minute one, and waiting out the wrong window converges on nothing. + assert.match(text, /account_sustained/); + assert.ok( + calls.some((c) => c.startsWith("bulkStatus:")), + "nothing moving must be confirmed with one direct read", + ); +}); + +test("a batch that is genuinely still uploading is left alone", async (t) => { + const calls: string[] = []; + const client = await clientFor( + sessionWith({ waitStates: ["timeout", "timeout"], probe: "still-running" }, calls), + t, + ); + + const text = textOf( + await client.callTool({ + name: "memwal_remember_status", + arguments: { job_ids: ["job-1", "job-2"], waitMs: 30000 }, + }), + ); + + assert.match(text, /still uploading/i); + assert.ok(!/rate limit/i.test(text), `invented a rate limit: ${text}`); +}); + +test("the probe is skipped when any row already settled", async (t) => { + const calls: string[] = []; + const client = await clientFor( + sessionWith({ waitStates: ["done", "timeout"], probe: "rate-limited" }, calls), + t, + ); + + const text = textOf( + await client.callTool({ + name: "memwal_remember_status", + arguments: { job_ids: ["job-1", "job-2"], waitMs: 30000 }, + }), + ); + + // A settled row proves the polls were getting through, so spending another + // request to confirm would be the opposite of the point. + assert.ok( + !calls.some((c) => c.startsWith("bulkStatus:")), + "probed despite evidence the polls were working", + ); + assert.match(text, /1\/2 saved/); +}); + +test("a probe that fails for any other reason does not invent an outcome", async (t) => { + const client = await clientFor( + sessionWith({ waitStates: ["timeout", "timeout"], probe: "throws" }), + t, + ); + + const text = textOf( + await client.callTool({ + name: "memwal_remember_status", + arguments: { job_ids: ["job-1", "job-2"], waitMs: 30000 }, + }), + ); + + // Falls back to what the wait reported rather than reading a failed probe + // as a failed write. + assert.match(text, /still uploading/i); + assert.ok(!/NOT stored/.test(text), `a failed probe was read as a failed write: ${text}`); +}); + +/** + * Single-job wait path: the confirming probe used to fall through + * `failed` / `not_found` into `stillRunning`. That is the case the probe + * exists for — polls that were all 429s (or otherwise swallowed) so the + * SDK stamped timeout even though the job already died. The agent was then + * told not to re-send, which is silent loss. + */ + +interface SingleJobBehaviour { + probe: "failed" | "not_found" | "running"; +} + +function timeout504(): Error { + const e = new Error("polling timed out after 30000ms"); + (e as Error & { status?: number }).status = 504; + return e; +} + +function sessionWithSingleJob(b: SingleJobBehaviour, calls: string[] = []): MemWalSession { + return { + oauthScope: "memwal:read memwal:write", + namespace: "default", + memwal: { + async waitForRememberJob(id: string) { + calls.push(`waitJob:${id}`); + throw timeout504(); + }, + async getRememberStatus(id: string) { + calls.push(`getStatus:${id}`); + if (b.probe === "failed") { + return { job_id: id, status: "failed", error: "walrus upload rejected" }; + } + if (b.probe === "not_found") { + return { job_id: id, status: "not_found" }; + } + return { job_id: id, status: "running" }; + }, + }, + } as unknown as MemWalSession; +} + +test("a single job whose wait times out and whose probe is failed is an error, not progress", async (t) => { + const calls: string[] = []; + const client = await clientFor(sessionWithSingleJob({ probe: "failed" }, calls), t); + + const result = await client.callTool({ + name: "memwal_remember_status", + arguments: { job_id: "job-1", waitMs: 30000 }, + }); + const text = textOf(result); + + assert.equal((result as { isError?: boolean }).isError, true); + assert.ok( + !/still uploading/i.test(text), + `a terminal failure still read as progress: ${text}`, + ); + assert.match(text, /failed/i); + assert.ok( + calls.includes("getStatus:job-1"), + "timed-out wait must be confirmed with one direct read", + ); +}); + +test("a single job whose wait times out and whose probe is not_found is an error, not progress", async (t) => { + const client = await clientFor(sessionWithSingleJob({ probe: "not_found" }), t); + + const result = await client.callTool({ + name: "memwal_remember_status", + arguments: { job_id: "job-1", waitMs: 30000 }, + }); + const text = textOf(result); + + assert.equal((result as { isError?: boolean }).isError, true); + assert.ok( + !/still uploading/i.test(text), + `a missing job still read as progress: ${text}`, + ); + assert.match(text, /not found/i); +}); + +test("a probe that answers for fewer jobs than it was asked about loses none", async (t) => { + // The confirming probe used to REPLACE the wait's rows with whatever it + // returned. A relayer that omits an id it cannot find therefore erased + // that job from the report entirely: the caller was never told it had to + // send the fact again, which is the silent loss this tool exists to stop. + const calls: string[] = []; + const client = await clientFor( + sessionWith({ waitStates: ["timeout", "timeout"], probe: "partial" }, calls), + t, + ); + + const text = textOf( + await client.callTool({ + name: "memwal_remember_status", + arguments: { job_ids: ["job-1", "job-2"], waitMs: 30000 }, + }), + ); + + assert.ok(calls.some((c) => c.startsWith("bulkStatus:")), "the probe must have run"); + assert.match(text, /job-1/, `the answered job is missing: ${text}`); + assert.match( + text, + /job-2/, + `the job the probe skipped was dropped from the report: ${text}`, + ); + // Both are still in flight, so the count must cover both, not just the + // one row the probe happened to answer for. + assert.match(text, /2 still uploading/i, `wrong in-flight count: ${text}`); +}); diff --git a/services/server/scripts/mcp/__tests__/tool-annotations.test.ts b/services/server/scripts/mcp/__tests__/tool-annotations.test.ts index ebb9bae49..86c3db199 100644 --- a/services/server/scripts/mcp/__tests__/tool-annotations.test.ts +++ b/services/server/scripts/mcp/__tests__/tool-annotations.test.ts @@ -36,6 +36,10 @@ test("tools/list publishes safe titles and behavior annotations for every remote title: "Remember Multiple Facts", annotations: { readOnlyHint: false, destructiveHint: false }, }, + memwal_remember_status: { + title: "Check a Remember Job", + annotations: { readOnlyHint: true, destructiveHint: false }, + }, memwal_analyze: { title: "Analyze and Remember", annotations: { readOnlyHint: false, destructiveHint: true }, diff --git a/services/server/scripts/mcp/__tests__/tool-failure.test.ts b/services/server/scripts/mcp/__tests__/tool-failure.test.ts new file mode 100644 index 000000000..c65e3af2d --- /dev/null +++ b/services/server/scripts/mcp/__tests__/tool-failure.test.ts @@ -0,0 +1,316 @@ +/** + * A tool call that timed out used to reach the agent as + * `Tool error: This operation was aborted` — no step, no word on whether the + * relayer was up, nothing to decide between retrying, waiting and reporting. + * `wrapTool` now names the cause, checks the relayer's health when the cause + * is the connection, and says what to do next (WALM-396). + */ +import assert from "node:assert/strict"; +import http from "node:http"; +import type { AddressInfo } from "node:net"; +import test from "node:test"; +import type { MemWalSession } from "../auth.js"; +import { + classifyToolError, + describeFailure, + probeRelayerHealth, +} from "../tools/failure.js"; +import { wrapTool } from "../tools/util.js"; + +/** A relayer whose `/health` answers however `respond` says. */ +async function fakeRelayer( + respond: (res: http.ServerResponse) => void, +): Promise<{ url: string; close: () => Promise }> { + const server = http.createServer((req, res) => { + if (req.url === "/health") respond(res); + else res.writeHead(404).end(); + }); + await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)); + const { port } = server.address() as AddressInfo; + return { + url: `http://127.0.0.1:${port}`, + close: () => + new Promise((resolve) => { + server.closeAllConnections(); + server.close(() => resolve()); + }), + }; +} + +/** A port nothing listens on, so a connect is refused at once. */ +async function closedPortUrl(): Promise { + const probe = await fakeRelayer(() => {}); + await probe.close(); + return probe.url; +} + +const healthy = (res: http.ServerResponse) => + res + .writeHead(200, { "content-type": "application/json" }) + .end(JSON.stringify({ status: "ok", version: "1.4.2" })); + +function sessionAt(relayerUrl: string): MemWalSession { + return { + accountId: `0x${"b".repeat(64)}`, + relayerUrl, + agentClient: "claude-code", + } as unknown as MemWalSession; +} + +function recallTimeout(stage: string): Error { + const body = { + error: `Recall timed out after 14001ms during ${stage}`, + message: `Recall timed out after 14001ms during ${stage}`, + code: "RECALL_TIMEOUT", + stage, + elapsed_ms: 14001, + }; + const err = new Error(`Walrus Memory server error (504): ${body.message}`) as Error & { + status: number; + serverCode: string; + }; + err.status = 504; + err.serverCode = "RECALL_TIMEOUT"; + err.cause = JSON.stringify(body); + return err; +} + +function named(name: string, message = "boom"): Error { + const err = new Error(message); + err.name = name; + return err; +} + +async function textOf( + session: MemWalSession, + tool: string, + err: unknown, +): Promise<{ text: string; isError?: boolean }> { + const originalError = console.error; + console.error = () => {}; + const originalWrite = process.stderr.write.bind(process.stderr); + (process.stderr as unknown as { write: unknown }).write = () => true; + try { + const result = await wrapTool(session, tool, async () => { + throw err; + })({}); + return { text: result.content[0].text, isError: result.isError }; + } finally { + console.error = originalError; + (process.stderr as unknown as { write: unknown }).write = originalWrite; + } +} + +// ── classifyToolError ──────────────────────────────────────────────────── + +test("a relayer RECALL_TIMEOUT is read with its stage and time", () => { + assert.deepEqual(classifyToolError(recallTimeout("walrus_download")), { + kind: "recall_timeout", + stage: "walrus_download", + elapsedMs: 14001, + }); +}); + +test("every shape the SDK gives up in reads as a timeout", () => { + // 0.1.7 aborts recall with a bare AbortError; 0.1.8 names its own. + for (const name of ["AbortError", "TimeoutError", "MemWalRequestTimeout"]) { + assert.deepEqual(classifyToolError(named(name)), { kind: "timeout" }, name); + } +}); + +test("a failed connect reads as unreachable, with the socket's code", () => { + const err = new TypeError("fetch failed", { + cause: Object.assign(new Error("connect ECONNREFUSED"), { code: "ECONNREFUSED" }), + }); + assert.deepEqual(classifyToolError(err), { + kind: "unreachable", + code: "ECONNREFUSED", + host: null, + }); +}); + +test("a failed DNS lookup keeps the host it was for", () => { + const err = new TypeError("fetch failed", { + cause: Object.assign(new Error("getaddrinfo ENOTFOUND fullnode.example"), { + code: "ENOTFOUND", + hostname: "fullnode.example", + }), + }); + assert.deepEqual(classifyToolError(err), { + kind: "unreachable", + code: "ENOTFOUND", + host: "fullnode.example", + }); +}); + +test("anything else keeps today's handling", () => { + assert.deepEqual(classifyToolError(new Error("bad input")), { kind: "other" }); + // The remember tools name their own timeouts; those keep their prefixes. + assert.deepEqual(classifyToolError(named("MemWalRememberJobTimeout")), { kind: "other" }); + assert.deepEqual(classifyToolError("not even an error"), { kind: "other" }); +}); + +// ── probeRelayerHealth ─────────────────────────────────────────────────── + +test("a healthy relayer is reported with its version", async (t) => { + const relayer = await fakeRelayer(healthy); + t.after(relayer.close); + const probe = await probeRelayerHealth(relayer.url, 2000); + assert.equal(probe.kind, "ok"); + assert.equal(probe.kind === "ok" && probe.version, "1.4.2"); +}); + +test("an unhealthy relayer is reported by status", async (t) => { + const relayer = await fakeRelayer((res) => res.writeHead(503).end()); + t.after(relayer.close); + const probe = await probeRelayerHealth(relayer.url, 2000); + assert.equal(probe.kind, "http"); + assert.equal(probe.kind === "http" && probe.status, 503); +}); + +test("a relayer that never answers is reported as a timeout, on time", async (t) => { + const relayer = await fakeRelayer(() => {}); + t.after(relayer.close); + const started = Date.now(); + const probe = await probeRelayerHealth(relayer.url, 200); + assert.equal(probe.kind, "timeout"); + assert.ok(Date.now() - started < 2000, "the probe must not outlive its budget"); +}); + +test("paused writes show in the health line, since /health still answers 200", async (t) => { + const relayer = await fakeRelayer((res) => + res + .writeHead(200, { "content-type": "application/json" }) + .end(JSON.stringify({ status: "ok", version: "1.4.2", writes: "paused" })), + ); + t.after(relayer.close); + const probe = await probeRelayerHealth(relayer.url, 2000); + assert.equal(probe.kind === "ok" && probe.writesUnavailable, true); + assert.match( + describeFailure("memwal_remember", { kind: "timeout" }, probe), + /Relayer health: ok \(\d+ms, v1\.4\.2, writes unavailable\)/, + ); +}); + +test("a probe given a budget AbortSignal cannot take still resolves", async () => { + const probe = await probeRelayerHealth("http://127.0.0.1:1", 2500.5); + assert.ok(probe.kind === "unreachable" || probe.kind === "timeout", probe.kind); +}); + +test("a relayer that refuses the connection is reported as unreachable", async () => { + const probe = await probeRelayerHealth(await closedPortUrl(), 2000); + assert.equal(probe.kind, "unreachable"); + assert.equal(probe.kind === "unreachable" && probe.code, "ECONNREFUSED"); +}); + +// ── describeFailure ────────────────────────────────────────────────────── + +test("every message carries the three labelled lines", () => { + const text = describeFailure("memwal_recall", { kind: "timeout" }, { kind: "ok", ms: 38 }); + assert.match(text, /^Cause: /m); + assert.match(text, /^Relayer health: /m); + assert.match(text, /^Next step: /m); +}); + +test("each stage gets its own advice", () => { + const expectations: Record = { + embed: /embedding/i, + vector_search: /database|index/i, + walrus_download: /lower `limit`/, + seal_decrypt: /SEAL/, + auth: /delegate key/, + }; + for (const [stage, advice] of Object.entries(expectations)) { + const text = describeFailure( + "memwal_recall", + { kind: "recall_timeout", stage, elapsedMs: 14001 }, + null, + ); + assert.match(text, advice, stage); + assert.match(text, /14\.0s/, stage); + } +}); + +// ── wrapTool end to end ────────────────────────────────────────────────── + +test("a recall the relayer cut short names the step and what to change", async () => { + const { text, isError } = await textOf( + sessionAt("http://127.0.0.1:1"), + "memwal_recall", + recallTimeout("walrus_download"), + ); + assert.equal(isError, true); + assert.match(text, /downloading memories from Walrus/); + assert.match(text, /lower `limit`/); + // The relayer answered this call, so there is nothing to probe. + assert.match(text, /Relayer health: up/); +}); + +test("a recall the SDK gave up on says the relayer is up, and that a retry is safe", async (t) => { + const relayer = await fakeRelayer(healthy); + t.after(relayer.close); + const { text, isError } = await textOf( + sessionAt(relayer.url), + "memwal_recall", + named("AbortError", "This operation was aborted"), + ); + assert.equal(isError, true); + assert.match(text, /Relayer health: ok/); + assert.match(text, /safe to retry/i); + assert.doesNotMatch(text, /This operation was aborted$/); +}); + +test("a recall that cannot reach a dead relayer says so", async () => { + const err = new TypeError("fetch failed", { + cause: Object.assign(new Error("connect ECONNREFUSED"), { code: "ECONNREFUSED" }), + }); + const { text } = await textOf(sessionAt(await closedPortUrl()), "memwal_recall", err); + assert.match(text, /ECONNREFUSED/); + assert.match(text, /Relayer health: unreachable/); +}); + +test("a failed request is not blamed on a relayer that answers its health check", async (t) => { + // SDK 0.1.7 builds the SEAL session on the Sui fullnode before the + // recall request goes out, and a failure there is a bare "fetch failed". + const relayer = await fakeRelayer(healthy); + t.after(relayer.close); + const err = new TypeError("fetch failed", { + cause: Object.assign(new Error("getaddrinfo ENOTFOUND fullnode.example"), { + code: "ENOTFOUND", + hostname: "fullnode.example", + }), + }); + const { text } = await textOf(sessionAt(relayer.url), "memwal_recall", err); + assert.match(text, /could not complete a network request/); + assert.match(text, /fullnode\.example/); + assert.match(text, /answered its health check/); + assert.doesNotMatch(text, /could not reach the relayer|stalled inside/); +}); + +test("a timeout on a healthy relayer does not claim the relayer is where it stalled", async (t) => { + const relayer = await fakeRelayer(healthy); + t.after(relayer.close); + const { text } = await textOf(sessionAt(relayer.url), "memwal_recall", named("AbortError")); + assert.match(text, /inside the relayer, or on a service it waits on first/); + assert.doesNotMatch(text, /stalled inside it\./); +}); + +test("a write that timed out is never offered a blind retry", async (t) => { + const relayer = await fakeRelayer(healthy); + t.after(relayer.close); + for (const tool of ["memwal_remember", "memwal_remember_bulk", "memwal_analyze"]) { + const { text } = await textOf(sessionAt(relayer.url), tool, named("AbortError")); + assert.match(text, /may already be stored/, tool); + assert.match(text, /memwal_recall/, tool); + assert.doesNotMatch(text, /safe to retry|please retry/i, tool); + } +}); + +test("an error that is not about time or reachability keeps its old wording", async () => { + const { text } = await textOf( + sessionAt("http://127.0.0.1:1"), + "memwal_recall", + new Error("namespace is invalid"), + ); + assert.equal(text, "Tool error: namespace is invalid"); +}); diff --git a/services/server/scripts/mcp/__tests__/tool-scope.test.ts b/services/server/scripts/mcp/__tests__/tool-scope.test.ts index e902001ac..1cdc272eb 100644 --- a/services/server/scripts/mcp/__tests__/tool-scope.test.ts +++ b/services/server/scripts/mcp/__tests__/tool-scope.test.ts @@ -8,6 +8,7 @@ import { createMcpServer } from "../server.js"; const WRITE_TOOLS = [ "memwal_remember", "memwal_remember_bulk", + "memwal_remember_status", "memwal_analyze", "memwal_restore", ]; diff --git a/services/server/scripts/mcp/__tests__/write-path-redaction.test.ts b/services/server/scripts/mcp/__tests__/write-path-redaction.test.ts new file mode 100644 index 000000000..eb909fbcc --- /dev/null +++ b/services/server/scripts/mcp/__tests__/write-path-redaction.test.ts @@ -0,0 +1,496 @@ +/** + * No credential reaches the SDK from any write tool (WALM-642). + * + * The ticket was reproduced exactly here: a preference plus a fake password URL + * sent through the real MCP server, with the SDK mocked, and `remember`, + * `remember_bulk` and `analyze` all passed the full text through unchanged. + * These tests are that repro, inverted — the mock now records everything the + * handler forwards, and every assertion is about what the SDK was *given*, not + * about what the tool said afterwards. A handler that redacted its reply but + * still forwarded the secret would pass a message-only test and fail these. + * + * Walrus storage is append-only and immutable, which is why the check has to be + * in front of the write: there is no delete to fall back on. + * + * The wait budget is zeroed so each tool returns at accept. That is the shortest + * path through each handler and it exercises the same pre-forward screen; the + * bounded-wait branch is covered by the `*-fast-return` files. + */ +process.env.MEMWAL_MCP_REMEMBER_WAIT_MS = "0"; + +import assert from "node:assert/strict"; +import test, { type TestContext } from "node:test"; +import { Client } from "@modelcontextprotocol/sdk/client/index.js"; +import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js"; +import type { MemWalSession } from "../auth.js"; + +const { createMcpServer } = await import("../server.js"); + +/** Everything the handlers handed to the SDK, in call order. */ +interface Forwarded { + remember: string[]; + bulk: string[][]; + analyze: string[]; +} + +function sessionWith(forwarded: Forwarded): MemWalSession { + return { + oauthScope: "memwal:read memwal:write", + namespace: "default", + memwal: { + async rememberAsync(text: string) { + forwarded.remember.push(text); + return { job_id: "job-1", status: "running" }; + }, + async rememberBulkAsync(items: Array<{ text: string }>) { + forwarded.bulk.push(items.map((i) => i.text)); + return { + job_ids: items.map((_, i) => `bulk-job-${i + 1}`), + total: items.length, + status: "accepted", + }; + }, + async analyze(text: string) { + forwarded.analyze.push(text); + // Echo the passage back as one extracted fact, so a leak that + // slipped through would also show up in the reply. + return { + job_ids: ["analyze-job-1"], + facts: [{ text }], + fact_count: 1, + status: "accepted", + owner: "0xowner", + }; + }, + }, + } as unknown as MemWalSession; +} + +async function clientFor(session: MemWalSession, t: TestContext): Promise { + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair(); + const server = createMcpServer(session); + const client = new Client({ name: "write-path-redaction-test", version: "1.0.0" }); + t.after(async () => { + await client.close(); + await server.close(); + }); + await server.connect(serverTransport); + await client.connect(clientTransport); + return client; +} + +function textOf(result: unknown): string { + return (result as { content: Array<{ text: string }> }).content + .map((c) => c.text) + .join("\n"); +} + +/** Everything the session was ever handed, flattened, plus the tool's reply. */ +function allForwarded(forwarded: Forwarded): string { + return [ + ...forwarded.remember, + ...forwarded.bulk.flat(), + ...forwarded.analyze, + ].join("\n"); +} + +const PASSWORD = "hunter2"; +const MIXED = + "I prefer dark mode in every editor, and the staging db is " + + `postgres://admin:${PASSWORD}@db.internal:5432/app`; +const PREFERENCE = "I prefer dark mode in every editor"; + +// ── the reproduction, on all three write paths ────────────────────────────── + +test("memwal_remember keeps the preference and never forwards the password", async (t) => { + const forwarded: Forwarded = { remember: [], bulk: [], analyze: [] }; + const client = await clientFor(sessionWith(forwarded), t); + const result = await client.callTool({ + name: "memwal_remember", + arguments: { text: MIXED }, + }); + + assert.equal(forwarded.remember.length, 1, "the write must still happen"); + const sent = forwarded.remember[0]; + assert.ok(!sent.includes(PASSWORD), "the password was forwarded to the SDK"); + assert.ok(!sent.includes("admin:"), "the userinfo was forwarded to the SDK"); + // The point of redacting rather than dropping: the fact survives. + assert.ok(sent.includes(PREFERENCE), "the preference was lost with the credential"); + assert.ok(sent.includes("db.internal:5432/app"), "the host was lost too"); + + const text = textOf(result); + assert.ok(!text.includes(PASSWORD), "the reply echoed the password back"); + assert.match(text, /url-credentials/); + assert.match(text, /credential span\(s\) were removed/); +}); + +test("memwal_remember_bulk screens every entry, and one bad entry does not sink the batch", async (t) => { + const forwarded: Forwarded = { remember: [], bulk: [], analyze: [] }; + const client = await clientFor(sessionWith(forwarded), t); + const result = await client.callTool({ + name: "memwal_remember_bulk", + arguments: { + facts: [ + "I always use pnpm", + MIXED, + "sk-abcdefghijklmnopqrstuvwxyz0123", + "Deploy on Thursdays", + ], + }, + }); + + assert.equal(forwarded.bulk.length, 1); + const sent = forwarded.bulk[0]; + const joined = sent.join("\n"); + assert.ok(!joined.includes(PASSWORD), "a password reached the SDK"); + assert.ok(!joined.includes("sk-abcdefghijklmnopqrstuvwxyz0123"), "a key reached the SDK"); + + // Three survive: the two clean facts, plus the redacted mixed one. The + // bare key had no fact around it, so it is dropped rather than stored as + // an empty placeholder. + assert.equal(sent.length, 3); + assert.ok(sent.includes("I always use pnpm"), "a clean fact was altered or dropped"); + assert.ok(sent.includes("Deploy on Thursdays"), "a clean fact was altered or dropped"); + assert.ok(sent.some((s) => s.includes(PREFERENCE))); + + const text = textOf(result); + assert.ok(!text.includes(PASSWORD)); + assert.match(text, /NOT SAVED \(1\)/); + assert.match(text, /#3/, "the dropped entry must be identified by position"); +}); + +test("memwal_analyze strips the passage before the extractor ever sees it", async (t) => { + const forwarded: Forwarded = { remember: [], bulk: [], analyze: [] }; + const client = await clientFor(sessionWith(forwarded), t); + const result = await client.callTool({ + name: "memwal_analyze", + arguments: { + text: + `${MIXED}\nAlso, my GitHub token is ghp_abcdefghijklmnopqrstuvwxyz0123456789 ` + + "and I review PRs on Fridays.", + }, + }); + + assert.equal(forwarded.analyze.length, 1); + const sent = forwarded.analyze[0]; + assert.ok(!sent.includes(PASSWORD), "the password reached the extractor LLM"); + assert.ok( + !sent.includes("ghp_abcdefghijklmnopqrstuvwxyz0123456789"), + "the GitHub token reached the extractor LLM", + ); + assert.ok(sent.includes(PREFERENCE)); + assert.ok(sent.includes("review PRs on Fridays")); + + const text = textOf(result); + assert.ok(!text.includes(PASSWORD)); + assert.match(text, /credential span\(s\) were removed/); +}); + +test("the delegate private key never reaches the SDK from any write path", async (t) => { + // The one secret that matters most for this product: the Ed25519 seed in + // ~/.memwal/credentials.json, which grants read AND write to the user's + // memories until the delegate is revoked. It is pure lowercase hex, so it + // is invisible to the entropy rule by design — the label beside it is what + // catches it. A user pasting their credentials file into chat is the + // realistic way this arrives. + const SEED = "4f3c2b1a9e8d7c6b5a4f3e2d1c0b9a8f7e6d5c4b3a2f1e0d9c8b7a6f5e4d3c2b"; + const pasted = + `I set up a second laptop today. From credentials.json: ` + + `"delegatePrivateKey": "${SEED}", and I use the work namespace there.`; + + const forwarded: Forwarded = { remember: [], bulk: [], analyze: [] }; + const client = await clientFor(sessionWith(forwarded), t); + + for (const [name, args] of [ + ["memwal_remember", { text: pasted }], + ["memwal_remember_bulk", { facts: [pasted] }], + ["memwal_analyze", { text: pasted }], + ] as Array<[string, Record]>) { + const result = await client.callTool({ name, arguments: args }); + assert.ok( + !textOf(result).includes(SEED), + `${name} echoed the delegate key back in its reply`, + ); + } + + const sent = allForwarded(forwarded); + assert.ok(sent.length > 0, "the writes must still happen"); + assert.ok(!sent.includes(SEED), "the delegate private key reached the SDK"); + // Redacted, not dropped: the fact around it survives on every path. + assert.equal(forwarded.remember.length, 1); + assert.equal(forwarded.bulk.length, 1); + assert.equal(forwarded.analyze.length, 1); + for (const text of [forwarded.remember[0], forwarded.bulk[0][0], forwarded.analyze[0]]) { + assert.ok(text.includes("second laptop"), "the fact was lost with the key"); + assert.ok(text.includes("work namespace"), "the fact was lost with the key"); + } +}); + +test("an unlabelled hex identifier still reaches the SDK unchanged", async (t) => { + // The other half of the label gate, asserted at the handler boundary: a + // 64-hex string nobody called a key is a digest, an object id or a blob id + // — the facts this product exists to remember. + const forwarded: Forwarded = { remember: [], bulk: [], analyze: [] }; + const client = await clientFor(sessionWith(forwarded), t); + const fact = + "My Sui package id is 0xe80f2feec1c139616a86c9f71210152e2a7ca552b20841f2e192f99f75864437 " + + "and the release digest is 4f3c2b1a9e8d7c6b5a4f3e2d1c0b9a8f7e6d5c4b3a2f1e0d9c8b7a6f5e4d3c2b"; + + const result = await client.callTool({ + name: "memwal_remember", + arguments: { text: fact }, + }); + assert.deepEqual(forwarded.remember, [fact]); + assert.doesNotMatch(textOf(result), /redacted|NOT SAVED/i); +}); + +// ── the rules that are not about credentials ──────────────────────────────── + +test("an explicit do-not-save is honoured on every write path", async (t) => { + const forwarded: Forwarded = { remember: [], bulk: [], analyze: [] }; + const client = await clientFor(sessionWith(forwarded), t); + const text = "My bank PIN is 4821 — don't save this."; + + for (const [name, args] of [ + ["memwal_remember", { text }], + ["memwal_remember_bulk", { facts: [text] }], + ["memwal_analyze", { text }], + ] as Array<[string, Record]>) { + const result = await client.callTool({ name, arguments: args }); + assert.match( + textOf(result), + /NOT SAVED|Nothing was saved/, + `${name} did not say it withheld the text`, + ); + } + + assert.deepEqual(forwarded.remember, []); + assert.deepEqual(forwarded.bulk, []); + assert.deepEqual(forwarded.analyze, []); + assert.ok(!allForwarded(forwarded).includes("4821")); +}); + +test("pasted third-party content is not saved as a user fact", async (t) => { + const forwarded: Forwarded = { remember: [], bulk: [], analyze: [] }; + const client = await clientFor(sessionWith(forwarded), t); + const pasted = "```\nERROR 500 from the vendor API\n at handler (index.js:42)\n```"; + + const remembered = await client.callTool({ + name: "memwal_remember", + arguments: { text: pasted }, + }); + assert.match(textOf(remembered), /pasted third-party content/); + + const analyzed = await client.callTool({ + name: "memwal_analyze", + arguments: { text: pasted }, + }); + assert.match(textOf(analyzed), /pasted third-party content/); + + assert.deepEqual(forwarded.remember, []); + assert.deepEqual(forwarded.analyze, []); +}); + +// ── the ordinary case, which must be untouched ────────────────────────────── + +test("a plain preference is forwarded byte-for-byte, with no note attached", async (t) => { + const forwarded: Forwarded = { remember: [], bulk: [], analyze: [] }; + const client = await clientFor(sessionWith(forwarded), t); + const clean = "I always use pnpm, TypeScript strict mode, and deploy on Thursdays."; + + const remembered = await client.callTool({ + name: "memwal_remember", + arguments: { text: clean }, + }); + assert.deepEqual(forwarded.remember, [clean]); + assert.doesNotMatch(textOf(remembered), /redacted|NOT SAVED/i); + + await client.callTool({ + name: "memwal_remember_bulk", + arguments: { facts: [clean, "My coffee order is a matcha oat latte"] }, + }); + assert.deepEqual(forwarded.bulk, [ + [clean, "My coffee order is a matcha oat latte"], + ]); + + await client.callTool({ name: "memwal_analyze", arguments: { text: clean } }); + assert.deepEqual(forwarded.analyze, [clean]); +}); + +test("the idempotency key is derived from what is actually written", async (t) => { + // Keyed on the original, a retry of a redacted fact would derive a key for + // text that was never sent — and the accept-timeout message promises a + // retry is safe. + const seen: string[] = []; + const session = { + oauthScope: "memwal:read memwal:write", + namespace: "default", + memwal: { + async rememberAsync(text: string, _ns: unknown, opts: { idempotencyKey: string }) { + seen.push(`${text}::${opts.idempotencyKey}`); + return { job_id: "job-1", status: "running" }; + }, + }, + } as unknown as MemWalSession; + + const client = await clientFor(session, t); + await client.callTool({ name: "memwal_remember", arguments: { text: MIXED } }); + await client.callTool({ name: "memwal_remember", arguments: { text: MIXED } }); + + assert.equal(seen.length, 2); + assert.equal(seen[0], seen[1], "the same fact must derive the same key twice"); + assert.ok(!seen[0].includes(PASSWORD)); +}); + +/* ─────────────────────────────────────────────────────────────────────────── + * Review findings on the WALM-642 branch, asserted where it counts: at what + * the SDK was handed. + * ------------------------------------------------------------------------ */ + +test("splitting a secret across bulk entries does not reach the SDK", async (t) => { + // Every label-gated rule searched a window inside ONE string, so the label + // in one entry and its value in the next walked past all of them — with + // MemWal's own delegate private key, the worst thing this product can leak. + // An agent that paraphrases a user across two entries is enough; nothing + // here needs malice. + const SEED = "4f3c2b1a9e8d7c6b5a4f3e2d1c0b9a8f7e6d5c4b3a2f1e0d9c8b7a6f5e4d3c2b"; + const forwarded: Forwarded = { remember: [], bulk: [], analyze: [] }; + const client = await clientFor(sessionWith(forwarded), t); + + const result = await client.callTool({ + name: "memwal_remember_bulk", + arguments: { + facts: [ + "I set up a second laptop today", + "my delegate private key for the mainnet account", + SEED, + "and I use the work namespace there", + ], + }, + }); + + assert.equal(forwarded.bulk.length, 1, "the other facts must still be written"); + assert.ok( + !forwarded.bulk[0].join("\n").includes(SEED), + "the delegate private key reached the SDK from a split batch", + ); + assert.ok(!textOf(result).includes(SEED), "the reply echoed the delegate key back"); + // The facts either side are still saved: this is a scalpel, not a batch + // refusal. + assert.ok(forwarded.bulk[0].some((f) => f.includes("second laptop"))); + assert.ok(forwarded.bulk[0].some((f) => f.includes("work namespace"))); +}); + +test("a bulk batch of plain identifiers is still forwarded byte-for-byte", async (t) => { + // The cross-entry screen only fires when a credential label is somewhere in + // the batch. Without one, nothing about the batch path may differ from the + // single-fact path. + const forwarded: Forwarded = { remember: [], bulk: [], analyze: [] }; + const client = await clientFor(sessionWith(forwarded), t); + const facts = [ + "My Sui package id is 0xe80f2feec1c139616a86c9f71210152e2a7ca552b20841f2e192f99f75864437", + "The release digest is 4f3c2b1a9e8d7c6b5a4f3e2d1c0b9a8f7e6d5c4b3a2f1e0d9c8b7a6f5e4d3c2b", + "I always use pnpm", + ]; + const result = await client.callTool({ + name: "memwal_remember_bulk", + arguments: { facts }, + }); + assert.deepEqual(forwarded.bulk, [facts]); + assert.doesNotMatch(textOf(result), /redacted|NOT SAVED/i); +}); + +test("one do-not-save line does not discard a whole transcript", async (t) => { + // `NO_SAVE_DIRECTIVE` is a whole-string predicate applied to a whole + // passage: one "don't save this part" line refused all forty turns, with + // `text: ""` and an instruction not to retry — and no `isError`, so the + // client read it as a successful call that had saved nothing. + const forwarded: Forwarded = { remember: [], bulk: [], analyze: [] }; + const client = await clientFor(sessionWith(forwarded), t); + const transcript = [ + "user: I always use pnpm for every project", + "assistant: noted", + "user: my bank PIN is 4821 - don't save this part", + "assistant: understood", + "user: I deploy on Thursdays and never on Fridays", + "assistant: makes sense", + ].join("\n"); + + const result = await client.callTool({ + name: "memwal_analyze", + arguments: { text: transcript }, + }); + + assert.equal(forwarded.analyze.length, 1, "the whole transcript was discarded again"); + const sent = forwarded.analyze[0]; + assert.ok(!sent.includes("4821"), "the withheld line reached the extractor"); + assert.ok(sent.includes("I always use pnpm for every project")); + assert.ok(sent.includes("I deploy on Thursdays")); + + const text = textOf(result); + assert.ok(!text.includes("4821")); + assert.match(text, /span\(s\) were dropped/, "the drop was not reported to the agent"); +}); + +test("a fenced transcript is extracted from, and a refusal is flagged as one", async (t) => { + const forwarded: Forwarded = { remember: [], bulk: [], analyze: [] }; + const client = await clientFor(sessionWith(forwarded), t); + const fenced = + "```\n" + + [ + "user: I always use pnpm for every project", + "assistant: noted", + "user: I deploy on Thursdays and never on Fridays", + "assistant: makes sense", + ].join("\n") + + "\n```"; + + await client.callTool({ name: "memwal_analyze", arguments: { text: fenced } }); + assert.equal(forwarded.analyze.length, 1, "a fenced transcript was refused as a paste"); + assert.ok(forwarded.analyze[0].includes("I always use pnpm for every project")); + + // A passage with nothing usable left is still refused — and now says so as + // an error, rather than looking like a call that succeeded and saved zero. + const refused = await client.callTool({ + name: "memwal_analyze", + arguments: { text: "My bank PIN is 4821 — don't save this." }, + }); + assert.equal(forwarded.analyze.length, 1, "a refused passage was forwarded anyway"); + assert.equal( + (refused as { isError?: boolean }).isError, + true, + "a call that saved nothing reported success", + ); + assert.match(textOf(refused), /NOT SAVED/); +}); + +test("memwal_analyze will not take an unbounded passage", async (t) => { + // The schema was `z.string().min(1)` with no maximum while the tool is + // documented as accepting a whole transcript, so the work the sidecar's + // single thread did in front of every other caller was the caller's choice. + const forwarded: Forwarded = { remember: [], bulk: [], analyze: [] }; + const client = await clientFor(sessionWith(forwarded), t); + const result = await client.callTool({ + name: "memwal_analyze", + arguments: { text: "a".repeat(200_001) }, + }); + assert.equal((result as { isError?: boolean }).isError, true); + assert.deepEqual(forwarded.analyze, [], "an over-long passage was forwarded"); +}); + +test("an ordinary URL is not mangled on the way to the SDK", async (t) => { + // The userinfo groups excluded `/` but not `?` or `=`, so this came out as + // `https://[redacted:url-credentials]@corp.com` — host and port destroyed, + // the mangled fact written to append-only storage, and the agent told a + // credential had been removed when there was none. + const forwarded: Forwarded = { remember: [], bulk: [], analyze: [] }; + const client = await clientFor(sessionWith(forwarded), t); + const fact = + "Our dashboard is at https://app.example.com:8443?owner=alice@corp.com and we deploy Fridays"; + const result = await client.callTool({ + name: "memwal_remember", + arguments: { text: fact }, + }); + assert.deepEqual(forwarded.remember, [fact]); + assert.doesNotMatch(textOf(result), /redacted|credential span/i); +}); diff --git a/services/server/scripts/mcp/server.ts b/services/server/scripts/mcp/server.ts index ad1c1c7e3..49708d64b 100644 --- a/services/server/scripts/mcp/server.ts +++ b/services/server/scripts/mcp/server.ts @@ -3,6 +3,10 @@ import { createRequire } from "node:module"; import { applyAgentClientFromServer } from "./agent-client.js"; import type { MemWalSession } from "./auth.js"; import { registerTools } from "./tools/index.js"; +import { + SECRET_EXCLUSION_RULES, + AUTO_SAVE_OPT_IN_RULE, +} from "./tools/memory-policy.js"; const requirePkg = createRequire(import.meta.url); @@ -25,6 +29,10 @@ const PACKAGE_VERSION: string = * `instructions` travels with `initialize`, before any `tools/list`, so lazy * loading cannot strip it. * + * The secret-exclusion and opt-in paragraphs are not written out here: they are + * pulled from tools/memory-policy.ts, the copy this package owns of the block + * shared with the MCP client and the plugin hooks (WALM-642). + * * Keep roughly in sync with the plugin's SessionStart hook, which delivers * equivalent text on the plugin install path: * packages/mcp/plugin/scripts/on_session_start.mjs @@ -50,6 +58,20 @@ const INSTRUCTIONS = [ "summary. Skip one-off tasks, the current file or bug, and small talk. Use", "memwal_remember_bulk when several distinct facts arrived at once.", "", + // WALM-642. Verbatim from tools/memory-policy.ts, which the tool + // descriptions state too, so instructions and descriptions cannot drift + // into telling the model two different things about secrets. + AUTO_SAVE_OPT_IN_RULE, + "", + SECRET_EXCLUSION_RULES, + "", + "By default memwal_remember and memwal_remember_bulk return in ~1s once the relayer has", + "accepted the job (job_id / job_ids). The Walrus write continues in the background (~30-60s)", + "and the fact is NOT stored yet. That is the normal result. Do not claim it is saved.", + "Do NOT re-send the same text — that queues duplicates. Resolve with memwal_remember_status", + "(job_id, or job_ids for a whole batch). Only a blob_id in the tool reply means the fact is", + "already stored (that happens when an optional wait budget was set and the write finished).", + "", "RECOVER: if memwal_recall unexpectedly returns nothing for a namespace that has been used", "before, call memwal_restore to rebuild the index from Walrus.", "", diff --git a/services/server/scripts/mcp/tools/analyze.ts b/services/server/scripts/mcp/tools/analyze.ts index 2ea48f8d3..7bf5f5be6 100644 --- a/services/server/scripts/mcp/tools/analyze.ts +++ b/services/server/scripts/mcp/tools/analyze.ts @@ -3,13 +3,45 @@ import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; import type { MemWalSession } from "../auth.js"; import { TOOL_METADATA } from "./annotations.js"; import { wrapTool, explorerFooter } from "./util.js"; +import { SECRET_EXCLUSION_RULES, AUTO_SAVE_OPT_IN_RULE } from "./memory-policy.js"; +import { + sanitizePassage, + redactionNotice, + refusalNotice, + droppedSpanNotice, +} from "./redaction.js"; +import { + ANALYZE_EXTRACTION_DEADLINE_MS, + REMEMBER_POLL_INTERVAL_MS, + REMEMBER_WAIT_MS, + pendingBulkMessage, + withAcceptDeadline, + withRelayerRetry, + withWaitDeadline, + isStillRunning, +} from "./remember-wait.js"; + +/** + * Ceiling on one passage, in characters. + * + * This schema had a `.min(1)` and no maximum while the tool is documented as + * taking a whole transcript, which made the input length entirely the caller's + * choice — and the redactor runs over every character of it on the sidecar's + * single thread, in front of every other in-flight tool call. The regexes are + * linear now (see `URL_USERINFO`), so this is a backstop rather than the fix: + * 200k characters screens in tens of milliseconds, is far more than any real + * transcript, and is well under what the extractor LLM behind `/api/analyze` + * would accept anyway. + */ +const MAX_ANALYZE_CHARS = 200_000; const ANALYZE_INPUT = { text: z .string() .min(1) + .max(MAX_ANALYZE_CHARS) .describe( - "Conversation transcript, note, or arbitrary text from which to extract memorable facts." + `Conversation transcript, note, or arbitrary text from which to extract memorable facts (max ${MAX_ANALYZE_CHARS} characters). Credential shapes (passwords, API keys, tokens, private keys, seed phrases, auth headers, URLs with an embedded user:password) are stripped from this text before it is sent for extraction, so no secret reaches the extractor or storage. A span the user asked not to save, or that is pasted third-party material, is dropped on its own — the rest of the passage is still extracted from.` ), namespace: z .string() @@ -20,9 +52,22 @@ const ANALYZE_INPUT = { } as const; /** - * memwal_analyze — let Walrus Memory's LLM extract distinct facts from a piece of - * text and persist each as its own memory. Resolves only after all extracted - * facts have been written end-to-end (or the call times out). + * memwal_analyze — let Walrus Memory's LLM extract distinct facts from a piece + * of text and persist each as its own memory. + * + * Returns once the relayer has extracted the facts and durably queued a write + * for each, the same contract as `memwal_remember_bulk`. + * + * Blocking to terminal was left in place while the two remember tools moved to + * a bounded wait, which made this the slowest tool in the set by a wide + * margin: measured at 37.0s against dev after `memwal_remember` had dropped to + * 0.2s there. The shape of the wait is the same as bulk's — N Walrus writes, + * one upload per wallet — so there was no reason for the answer to be shaped + * differently. + * + * Extraction itself is worth waiting for, and this still does: `analyze()` + * resolves after the LLM has run, so the facts it found are in the reply. + * Only the upload of those facts is handed back as job_ids. */ export function registerAnalyzeTool( server: McpServer, @@ -33,29 +78,188 @@ export function registerAnalyzeTool( { ...TOOL_METADATA.memwal_analyze, description: - "Extract memorable facts from a longer passage of text (preferences, habits, biographical info, constraints) and save each as a separate Walrus Memory memory. Use this when you want MemWal's LLM to split the facts out of a transcript or notes for you; if you already know the exact facts, use memwal_remember or memwal_remember_bulk instead.", + "Extract memorable facts from a longer passage of text (preferences, habits, biographical info, constraints) and save each as a separate Walrus Memory memory. Use this when you want MemWal's LLM to split the facts out of a transcript or notes for you; if you already know the exact facts, use memwal_remember or memwal_remember_bulk instead. The extracted facts come back immediately; if the result says the writes are still in flight it carries job_ids — confirm them with memwal_remember_status rather than telling the user they are saved. " + + AUTO_SAVE_OPT_IN_RULE + + " " + + SECRET_EXCLUSION_RULES + + " This tool forwards a whole passage, so it is the easiest way to leak a credential that happened to sit next to a fact: the passage is stripped of credential shapes before it is sent for extraction. A span the user asked not to save, or that is pasted third-party material, is dropped on its own and named in the reply; only a passage with nothing usable left is refused outright.", inputSchema: ANALYZE_INPUT, }, wrapTool<{ text: string; namespace?: string }>(session, "memwal_analyze", async ({ text, namespace }) => { - const result = await session.memwal.analyzeAndWait(text, namespace, { - timeoutMs: 180_000, + // Runs BEFORE the passage reaches the SDK, and therefore before it + // reaches the extractor LLM. Everything this tool stores is derived + // from this text, so a credential left in it can be copied into any + // number of extracted facts — on append-only storage (WALM-642). + // + // `sanitizePassage`, not `sanitizeFact`: the refusal predicates are + // whole-string, and applied to a transcript one "don't save this + // part" line threw away every other turn with it. They are scoped + // per span here, so the offending span is dropped and named and the + // rest is still extracted from. + const safe = sanitizePassage(text); + if (safe.refusal) { + return { + // Flagged as an error, because it is not a successful call: + // nothing was extracted and nothing was saved, and a bare + // text result reads to a client exactly like one that did. + isError: true, + content: [ + { type: "text" as const, text: refusalNotice(safe.refusal) }, + ], + }; + } + const notice = [ + redactionNotice(safe.kinds, safe.count), + droppedSpanNotice(safe.dropped, safe.segments), + ] + .filter(Boolean) + .join("\n\n"); + const safeText = safe.text; + + // `analyze` (not `analyzeAndWait`) returns once extraction is done + // and every fact has a queued job, which is the point this tool can + // usefully answer at. + const accepted = await withAcceptDeadline( + // Same reasoning as bulk: the retry only fires on rejections + // that never reached the handler, so no job row can exist yet + // to duplicate. + withRelayerRetry( + () => session.memwal.analyze(safeText, namespace), + "analyze this text", + ), + "memwal_analyze extraction", + // Not an accept. `/api/analyze` runs the extractor LLM inline + // before it answers — which is why the SDK allows this call 60s + // where it allows a remember 30s. The 15s accept ceiling would + // have cut off healthy extraction on any transcript long enough + // to be worth extracting from. + // + // Budgeted just under the MCP client's own ceiling rather than + // level with it: at 60_000 this leg raced the host's abort and + // the caller lost the error message, on the one endpoint whose + // retry duplicates every extracted fact. + { idempotent: false, deadlineMs: ANALYZE_EXTRACTION_DEADLINE_MS }, + ); + + const withNotice = (body: string) => + notice ? `${body}\n\n${notice}` : body; + + const facts = accepted.facts ?? []; + // Nothing to wait on, and nothing to confirm later. Say so plainly + // rather than handing back an empty job list. + if (accepted.job_ids.length === 0) { + return { + content: [ + { + type: "text" as const, + text: withNotice( + `Extracted 0 facts from that text — nothing was saved.`, + ), + }, + ], + }; + } + + const entries = accepted.job_ids.map((jobId, i) => ({ + jobId, + text: facts[i]?.text ?? "", + })); + + // The extraction result is the part of this call an agent can act + // on immediately, so it leads — the write status follows it. + const extracted = `Extracted ${facts.length} fact(s):\n${entries + .map((e, i) => `${i + 1}. ${e.text || "(unknown fact)"}`) + .join("\n")}`; + + const pending = (waitedMs: number) => ({ + content: [ + { + type: "text" as const, + text: withNotice( + `${extracted}\n\n${pendingBulkMessage(entries, waitedMs)}`, + ), + }, + ], }); - const lines = result.results.map( - (r, i) => - `${i + 1}. [${r.status}]${r.blob_id ? ` blob_id=${r.blob_id}` : ""} ${ - result.facts[i]?.text ?? "(unknown fact)" - }` + + if (REMEMBER_WAIT_MS === 0) return pending(0); + + const startedAt = Date.now(); + const namespaces = entries.map( + () => namespace ?? session.namespace ?? "default" ); - const summary = `Extracted ${result.facts.length} fact(s) — succeeded=${result.succeeded} failed=${result.failed}`; + // The wait is a courtesy; the accept above is the part that had to + // succeed. If the relayer goes quiet mid-poll `withWaitDeadline` + // raises MemWalRelayerUnresponsive, and letting that propagate + // would discard both the job_ids and the extracted facts — the + // caller would have no way to settle writes that are still running + // and no way to get the extraction back without paying for it + // again. `memwal_remember` already degrades this way; so does this. + let result; + try { + result = await withWaitDeadline( + session.memwal.waitForRememberJobs(accepted.job_ids, namespaces, { + timeoutMs: REMEMBER_WAIT_MS, + pollIntervalMs: REMEMBER_POLL_INTERVAL_MS, + }), + REMEMBER_WAIT_MS, + ); + } catch (err) { + if (!isStillRunning(err)) throw err; + return pending(Date.now() - startedAt); + } + const waitedMs = Date.now() - startedAt; + + const unfinished = result.results.filter((r) => r.status === "timeout"); + if (unfinished.length === result.results.length) return pending(waitedMs); + + const lines = result.results.map((r, i) => { + // `timeout` is not a failure — the write is still running and + // its job_id is how the caller settles it later. Rendered the + // same way memwal_remember_bulk renders it. + const state = + r.status === "timeout" ? `still uploading, job_id=${r.id}` : r.status; + // Only a settled row may show a blob_id — see remember-status.ts + // for why an unfinished write can already carry a real one. + const blob = r.status === "done" && r.blob_id ? ` blob_id=${r.blob_id}` : ""; + return `${i + 1}. [${state}]${blob} ${ + entries[i]?.text || "(unknown fact)" + }`; + }); + // `result.failed` is total-minus-succeeded, so it counts a + // still-uploading write as failed — while the stragglers block + // below tells the agent those same jobs are on their way and must + // not be re-sent. An agent reading `failed=` re-sends an in-flight + // write, which is a duplicate paid Walrus blob queued behind the + // original. Count only what actually reached a terminal failure. + // Same correction as remember-bulk.ts; this file was written from + // the same template one commit earlier and missed it. + const reallyFailed = result.results.filter( + (r) => r.status !== "done" && r.status !== "timeout", + ).length; + const summary = + `Extracted ${facts.length} fact(s) — succeeded=${result.succeeded}` + + (reallyFailed ? ` failed=${reallyFailed}` : "") + + (unfinished.length ? ` (${unfinished.length} still uploading)` : ""); + const stragglers = + unfinished.length > 0 + ? `\n\n${pendingBulkMessage( + result.results.flatMap((r, i) => + r.status === "timeout" + ? [{ jobId: r.id, text: entries[i]?.text ?? "" }] + : [], + ), + waitedMs, + )}` + : ""; const footer = result.succeeded > 0 ? `\n\n${explorerFooter()}` : ""; return { content: [ { - type: "text", - text: - lines.length > 0 - ? `${summary}\n\n${lines.join("\n")}${footer}` - : `${summary}${footer}`, + type: "text" as const, + text: withNotice( + `${summary}\n\n${lines.join("\n")}${stragglers}${footer}`, + ), }, ], }; diff --git a/services/server/scripts/mcp/tools/annotations.ts b/services/server/scripts/mcp/tools/annotations.ts index 12e85fbde..71465d66c 100644 --- a/services/server/scripts/mcp/tools/annotations.ts +++ b/services/server/scripts/mcp/tools/annotations.ts @@ -15,6 +15,11 @@ export const TOOL_METADATA = { title: "Remember Multiple Facts", annotations: { readOnlyHint: false, destructiveHint: false }, }, + memwal_remember_status: { + title: "Check a Remember Job", + // Reads the state of a write already in flight; starts no new work. + annotations: { readOnlyHint: true, destructiveHint: false }, + }, memwal_analyze: { title: "Analyze and Remember", // Context recall may remove stale vector rows for blobs confirmed absent. diff --git a/services/server/scripts/mcp/tools/failure.ts b/services/server/scripts/mcp/tools/failure.ts new file mode 100644 index 000000000..41c2b513d --- /dev/null +++ b/services/server/scripts/mcp/tools/failure.ts @@ -0,0 +1,229 @@ +/** + * Why a tool call failed, in words an agent can act on. Every message carries + * the same three lines — `Cause`, `Relayer health`, `Next step` — because the + * agent only reads text. + */ + +export type HealthProbe = + | { kind: "ok"; ms: number; version?: string; writesUnavailable?: boolean } + | { kind: "http"; ms: number; status: number } + | { kind: "unreachable"; ms: number; code: string } + | { kind: "timeout"; ms: number }; + +export type ToolFailure = + | { kind: "recall_timeout"; stage: string | null; elapsedMs: number | null } + | { kind: "timeout" } + | { kind: "unreachable"; code: string; host: string | null } + | { kind: "other" }; + +/** Tools whose call may already have stored something by the time it fails. + * Mirrors the bridge's `MUTATING_TOOLS`: a retry of one of these can mint a + * second paid blob, so none is ever told a retry is safe. */ +export const WRITE_TOOLS: ReadonlySet = new Set([ + "memwal_remember", + "memwal_remember_bulk", + "memwal_analyze", +]); + +/** How the SDK says it gave up waiting: 0.1.7 aborts recall with a bare + * `AbortError`, 0.1.8 throws its own `MemWalRequestTimeout`. */ +const TIMEOUT_NAMES = new Set(["AbortError", "TimeoutError", "MemWalRequestTimeout"]); + +export function classifyToolError(err: unknown): ToolFailure { + if (err == null || typeof err !== "object") return { kind: "other" }; + const e = err as { name?: unknown; message?: unknown; serverCode?: unknown; cause?: unknown }; + if (e.serverCode === "RECALL_TIMEOUT") { + // The SDK keeps the raw body on `cause`; the stage is only there. + const body = parseObject(e.cause); + return { + kind: "recall_timeout", + stage: typeof body?.stage === "string" ? body.stage : null, + elapsedMs: typeof body?.elapsed_ms === "number" ? body.elapsed_ms : null, + }; + } + if (typeof e.name === "string" && TIMEOUT_NAMES.has(e.name)) return { kind: "timeout" }; + if (e.message === "fetch failed") { + return { kind: "unreachable", code: codeOf(e.cause), host: hostOf(e.cause) }; + } + return { kind: "other" }; +} + +/** `GET {baseUrl}/health`, bounded by `timeoutMs`. Never rejects: a probe + * that fails is itself the answer. */ +export async function probeRelayerHealth( + baseUrl: string, + timeoutMs: number, +): Promise { + const started = Date.now(); + let signal: AbortSignal | undefined; + try { + // Inside the `try`: `AbortSignal.timeout` throws on a value it + // cannot take, and this function must not reject. + signal = AbortSignal.timeout(timeoutMs); + const res = await fetch(`${baseUrl.replace(/\/+$/, "")}/health`, { signal }); + const ms = Date.now() - started; + if (!res.ok) { + await res.body?.cancel(); + return { kind: "http", ms, status: res.status }; + } + const body = (await res.json().catch(() => null)) as { + version?: unknown; + write_ready?: unknown; + writes?: unknown; + } | null; + const version = typeof body?.version === "string" ? body.version : undefined; + // `/health` answers 200 while writes are paused or Postgres is + // full; "ok" alone would hide the cause of a failed write. + const writesUnavailable = body?.write_ready === false || body?.writes === "paused"; + return { kind: "ok", ms, version, ...(writesUnavailable ? { writesUnavailable } : {}) }; + } catch (err) { + const ms = Date.now() - started; + if (signal?.aborted) return { kind: "timeout", ms }; + return { kind: "unreachable", ms, code: codeOf((err as { cause?: unknown })?.cause) }; + } +} + +const STAGES: Record = { + embed: { + doing: "turning the query into an embedding", + next: "The embedding provider is slow. Wait a minute, then retry.", + }, + vector_search: { + doing: "searching the memory index", + next: "The memory database is slow. Wait a minute, then retry.", + }, + walrus_download: { + doing: "downloading memories from Walrus", + next: + "Walrus is slow to return memories. Retry once; if it happens again, " + + "lower `limit` so fewer memories are downloaded.", + }, + seal_decrypt: { + doing: "decrypting memories with SEAL", + next: "The SEAL key servers are slow. Wait a minute, then retry.", + }, + auth: { + doing: "checking credentials, before the recall could start", + next: + "Checking the delegate key used up the whole deadline. Retry once; if it " + + "happens again, report it with the time of the call.", + }, +}; + +export function describeFailure( + tool: string, + failure: Exclude, + probe: HealthProbe | null, +): string { + if (failure.kind === "recall_timeout") { + const known = failure.stage === null ? undefined : STAGES[failure.stage]; + const doing = known?.doing ?? (failure.stage ? `at step "${failure.stage}"` : ""); + const after = + failure.elapsedMs === null ? "" : ` after ${(failure.elapsedMs / 1000).toFixed(1)}s`; + return lines( + "❌ Walrus Memory recall timed out.", + `the relayer stopped${after}${doing ? ` while ${doing}` : ""}.`, + "up (it answered this call).", + known?.next ?? "Retry once.", + ); + } + + // Before a recall reaches the relayer, the SDK builds a SEAL session on + // the Sui fullnode, and on 0.1.7 its 15s clock is already running. So a + // timeout or a failed request is only the relayer's when the relayer's + // own `/health` also fails; otherwise say what is actually known. + const healthy = probe?.kind === "ok"; + let headline: string; + let cause: string; + if (failure.kind === "timeout") { + headline = `❌ Walrus Memory ${tool} timed out.`; + cause = healthy + ? "no answer within the SDK's time limit. The relayer answered its health check, " + + "so the time went on this call: inside the relayer, or on a service it waits " + + "on first (such as the Sui fullnode)." + : "no answer within the SDK's time limit, and the relayer is not healthy."; + } else if (healthy) { + const where = failure.host ? ` to ${failure.host}` : ""; + headline = `❌ Walrus Memory ${tool} could not complete a network request.`; + cause = + `a request this call depends on failed${where} (${failure.code}). The relayer ` + + "answered its health check, so the failure is on the way to another service " + + "this call needs (such as the Sui fullnode), or it was brief."; + } else { + headline = `❌ Walrus Memory ${tool} could not reach the relayer.`; + cause = `the connection to the relayer failed (${failure.code}).`; + } + + // Not "this call only reads": `memwal_restore` re-indexes. What makes a + // retry safe is that none of these can store a duplicate. + let next: string; + if (WRITE_TOOLS.has(tool)) { + next = + "this call writes, so it may already be stored. Check with `memwal_recall` " + + "before saving it again." + + (healthy ? "" : " The relayer is not healthy right now, so wait a minute first."); + } else if (healthy) { + next = + "repeating this call cannot store a duplicate, so it is safe to retry once. If " + + "it keeps happening, report it with the time of the call."; + } else { + next = + "repeating this call cannot store a duplicate, so it is safe to retry, but " + + "wait a minute first: the relayer is not healthy right now."; + } + return lines(headline, cause, healthLine(probe), capitalize(next)); +} + +function healthLine(probe: HealthProbe | null): string { + if (probe === null) return "not checked."; + switch (probe.kind) { + case "ok": + return ( + `ok (${probe.ms}ms${probe.version ? `, v${probe.version}` : ""}` + + `${probe.writesUnavailable ? ", writes unavailable" : ""}).` + ); + case "http": + return `HTTP ${probe.status} (${probe.ms}ms) — up, but not healthy.`; + case "unreachable": + return `unreachable (${probe.code}).`; + case "timeout": + return `no answer within ${(probe.ms / 1000).toFixed(1)}s.`; + } +} + +function lines(headline: string, cause: string, health: string, next: string): string { + return [ + headline, + `Cause: ${capitalize(cause)}`, + `Relayer health: ${health}`, + `Next step: ${next}`, + ].join("\n"); +} + +function capitalize(text: string): string { + return text.charAt(0).toUpperCase() + text.slice(1); +} + +function parseObject(raw: unknown): Record | null { + if (typeof raw !== "string") return null; + try { + const parsed = JSON.parse(raw); + return parsed && typeof parsed === "object" ? parsed : null; + } catch { + return null; + } +} + +/** The host a failed connect was for, when Node says (`getaddrinfo` names + * the hostname, a refused connect the address). */ +function hostOf(cause: unknown): string | null { + const c = cause as { hostname?: unknown; address?: unknown } | null; + if (typeof c?.hostname === "string") return c.hostname; + if (typeof c?.address === "string") return c.address; + return null; +} + +function codeOf(cause: unknown): string { + const code = (cause as { code?: unknown } | null)?.code; + return typeof code === "string" ? code : "unknown"; +} diff --git a/services/server/scripts/mcp/tools/health.ts b/services/server/scripts/mcp/tools/health.ts index b1ca2a8d0..dd403170f 100644 --- a/services/server/scripts/mcp/tools/health.ts +++ b/services/server/scripts/mcp/tools/health.ts @@ -40,8 +40,18 @@ export function registerHealthTool( const relayerNote = session.publicRelayerUrl ? ` relayer=${session.publicRelayerUrl}` : ""; - const pausedNote = extra.writes === "paused" ? " writes=paused" : ""; - const writeNote = `${readyNote}${pausedNote}`; + // "degraded" is the case write_ready cannot express: the relayer + // still accepts and durably queues a write, so write_ready stays + // true, but recent durable writes are failing and none are + // landing. Say so, or an agent keeps queueing writes that fail + // minutes later. + const writesStateNote = + extra.writes === "paused" + ? " writes=paused" + : extra.writes === "degraded" + ? " writes=degraded (accepted, but recent writes are failing downstream)" + : ""; + const writeNote = `${readyNote}${writesStateNote}`; return { content: [ { diff --git a/services/server/scripts/mcp/tools/index.ts b/services/server/scripts/mcp/tools/index.ts index e230ee39c..763ed0a73 100644 --- a/services/server/scripts/mcp/tools/index.ts +++ b/services/server/scripts/mcp/tools/index.ts @@ -2,6 +2,7 @@ import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; import type { MemWalSession } from "../auth.js"; import { registerRememberTool } from "./remember.js"; +import { registerRememberStatusTool } from "./remember-status.js"; import { registerRememberBulkTool } from "./remember-bulk.js"; import { registerRecallTool } from "./recall.js"; import { registerAnalyzeTool } from "./analyze.js"; @@ -26,6 +27,7 @@ export function registerTools(server: McpServer, session: MemWalSession): void { if (canWrite) { registerRememberTool(server, session); + registerRememberStatusTool(server, session); registerRememberBulkTool(server, session); registerAnalyzeTool(server, session); registerRestoreTool(server, session); @@ -38,6 +40,7 @@ export function registerTools(server: McpServer, session: MemWalSession): void { export { registerRememberTool, + registerRememberStatusTool, registerRememberBulkTool, registerRecallTool, registerAnalyzeTool, diff --git a/services/server/scripts/mcp/tools/memory-policy.ts b/services/server/scripts/mcp/tools/memory-policy.ts new file mode 100644 index 000000000..924b5e5e2 --- /dev/null +++ b/services/server/scripts/mcp/tools/memory-policy.ts @@ -0,0 +1,82 @@ +/** + * Shared automatic-memory policy — the relayer sidecar's copy. + * + * Consumed by the live tool descriptions (`remember`, `remember-bulk`, + * `analyze`) and by the sidecar's `initialize` instructions in `server.ts`. + * Descriptions are the layer every MCP client sees once tools are loaded, so + * the rules have to be stated here too and not only in the client package. + * + * WALM-642. + */ + +// ─── memwal:policy-block:start ─────────────────────────────────────────────── +// WALM-642. The lines between these two markers are BYTE-IDENTICAL in three +// files that cannot import one another, because the three packages have no +// workspace link: +// +// packages/mcp/src/memory-policy.ts — MCP client: initialize +// instructions + the +// cold-start tools/list +// packages/mcp/plugin/scripts/lib/memory-policy.mjs — plugin hooks: the +// guidance injected at +// SessionStart / +// UserPromptSubmit / +// PostToolUse +// services/server/scripts/mcp/tools/memory-policy.ts — relayer sidecar: the +// live tool descriptions +// +// The duplication is deliberate and pinned: `memory-policy-sync` tests on both +// sides extract this block from each file and compare the bytes, so editing one +// copy fails the suite until the other two match. Edit the block, then copy it +// verbatim — markers included — into the other two files. + +/** + * The secret-exclusion and do-not-save rules, stated verbatim by every + * automatic-save surface. + * + * These are model-facing rules, not enforcement. The programmatic backstop is + * the redactor in the relayer sidecar's write path + * (services/server/scripts/mcp/tools/redaction.ts), which runs before any text + * reaches the SDK. + */ +export const SECRET_EXCLUSION_RULES = [ + "NEVER save a credential, even when it sits next to something worth saving: passwords,", + "API keys, access or refresh tokens, private keys, seed or recovery phrases, authorization", + "headers, session cookies, and connection strings or URLs that embed a user:password.", + "When a message mixes a preference with a credential, save the preference alone and leave", + "the credential out; never store the line verbatim.", + "If the user says not to save something ('don't save this', 'off the record', or the same", + "in any language), do not save it, and do not save a paraphrase of it either.", + "Do not store quoted or pasted third-party material — log excerpts, code, articles, other", + "people's messages — as if it were a fact about this user. Save only what the user is", + "telling you about themselves or their work, in your own words.", +].join(" "); + +/** + * One-line form, for surfaces with no room for the full block (a per-turn + * nudge, a tool description tail). It is a reminder of the block above, never a + * replacement for it: any surface that drives an automatic save states the full + * `SECRET_EXCLUSION_RULES`. + */ +export const SECRET_EXCLUSION_SUMMARY = [ + "Never save passwords, keys, tokens or other credentials — not even beside a fact worth", + "saving; honour an explicit 'do not save this'; never store pasted third-party content as", + "a fact about the user.", +].join(" "); + +/** + * Whether to save unprompted is the user's standing choice, and this is the + * sentence that says so. A direct request ("remember that ...") is never gated + * by it — the gate is only on saving something the user did not ask you to save. + */ +export const AUTO_SAVE_OPT_IN_RULE = [ + "Whether to save things the user did not ask you to save is their standing choice, made once", + "in a terminal. When automatic memory is on, save durable facts as they state them; when it is", + "off, save only what they ask you to save in that turn. That question is put by `memwal-mcp", + "login` and set by `memwal-mcp auto-save on|off` — never ask the user to answer it in chat,", + "and never answer it on their behalf.", +].join(" "); + +/** Bumped whenever the text above changes, so a stale copy is identifiable. */ +export const MEMORY_POLICY_VERSION = "2026-09-17.2"; +// ─── memwal:policy-block:end ───────────────────────────────────────────────── diff --git a/services/server/scripts/mcp/tools/redaction.ts b/services/server/scripts/mcp/tools/redaction.ts new file mode 100644 index 000000000..1a118c230 --- /dev/null +++ b/services/server/scripts/mcp/tools/redaction.ts @@ -0,0 +1,906 @@ +/** + * Credential redaction for the memory write path (WALM-642). + * + * Every tool that forwards free text to the SDK — `memwal_remember`, + * `memwal_remember_bulk`, `memwal_analyze` — runs its input through + * `sanitizeFact` FIRST. The model-facing rules in `memory-policy.ts` are + * guidance; this module is the backstop that does not depend on a model having + * read them. Walrus storage is append-only and immutable: a secret that reaches + * it cannot be deleted, so the check has to sit in front of the write rather + * than behind it. + * + * Two shapes of answer: + * + * - REDACT (the common case, and what the ticket asks for): a message that + * mixes a durable preference with a credential keeps the preference and + * loses only the credential span, replaced by `[redacted:]`. Dropping + * the whole fact would lose the thing the user actually wanted stored. + * - REFUSE: nothing safe is left, the user said not to save it, or the text + * is plainly pasted third-party material. The caller forwards nothing. + * + * ── False positives, deliberately ─────────────────────────────────────────── + * The patterns below are SHAPE-based, not entropy-based. That is a choice, and + * it costs recall: + * + * - There is no free-standing "long random-looking string" rule. MemWal's own + * durable facts are exactly that shape — Walrus blob ids (43-char + * base64url), Sui object and account ids (`0x` + 64 hex), git SHAs (40 hex), + * content digests. A generic high-entropy rule would redact the product's + * primary nouns. The one entropy rule that survived + * (`HIGH_ENTROPY_SECRET`) demands ≥64 characters AND mixed case AND a + * digit, which excludes every one of those (all-lowercase hex, or shorter), + * while still catching a raw base64 key blob. + * - `password:`-style assignments redact whatever follows the separator, so + * "password: ask Marta" loses "ask Marta". Over-redacting a sentence about + * a credential is cheap; under-redacting the credential is permanent. + * - Hex key material is caught by the LABEL next to it, never by its shape — + * see `HEX_RUN`. That is what lets MemWal's own 64-hex delegate private key + * be removed while a bare 40-hex commit SHA or a `0x`-prefixed Sui object id + * is left alone. The cost is that a hex secret pasted with no label at all + * still passes; the model-facing rules are what cover that. + * - Quoted-content detection only fires on unambiguous pastes (a fenced + * block, a multi-line `>` quotation, or a ≥200-char fully quoted passage). + * A short quoted sentence is left to the model-facing rules, because an + * agent legitimately quotes the user's own words back when saving a fact. + * + * Nothing here logs, echoes, or returns the matched secret. Callers get the + * redacted text and the KINDS that were removed — never the values. + */ + +export type RedactionKind = + | "url-credentials" + | "vendor-api-key" + | "private-key-block" + | "jwt" + | "credential-assignment" + | "auth-header" + | "seed-phrase" + | "labelled-key-material" + | "high-entropy-secret"; + +/** Why a text was refused outright instead of redacted. */ +export type RefusalReason = + /** The user said not to save it. */ + | "no-save-directive" + /** After redaction there was no fact left — the text was only a secret. */ + | "credential-only" + /** Pasted third-party material, not a fact about this user. */ + | "pasted-content"; + +export interface SanitizedText { + /** Text safe to forward. Empty when `refusal` is set. */ + text: string; + /** True when the text was changed or refused. */ + changed: boolean; + /** Kinds removed, first-seen order. Never contains a secret value. */ + kinds: RedactionKind[]; + /** Number of spans replaced. */ + count: number; + /** Set when the caller must forward nothing at all. */ + refusal?: RefusalReason; +} + +function placeholder(kind: RedactionKind): string { + return `[redacted:${kind}]`; +} + +/** + * An explicit instruction not to save, from the user. + * + * Deliberately demands a demonstrative object ("this", "that", "it"): without + * it, "remember that I don't save screenshots to the Desktop" — a perfectly + * good durable preference — would be refused as a do-not-save directive. + */ +const NO_SAVE_DIRECTIVE = + /\b(?:do\s+not|don'?t|dont|never|please\s+do\s+not|please\s+don'?t)\s+(?:save|store|remember|record|keep|persist|log)\s+(?:this|that|it|these|those|any\s+of\s+(?:this|that|it))\b/i; + +/** The idiom, which carries the same instruction without naming saving. */ +const OFF_THE_RECORD = /\boff[-\s]the[-\s]record\b/i; + +/** + * PEM private key blocks. The second pattern is not redundant: a paste that was + * cut off mid-key has a BEGIN line and no END line, and without the open-ended + * form the key body would survive untouched. + */ +const PEM_CLOSED = + /-----BEGIN [A-Z0-9 ]*PRIVATE KEY-----[\s\S]*?-----END [A-Z0-9 ]*PRIVATE KEY-----/g; +const PEM_OPEN = /-----BEGIN [A-Z0-9 ]*PRIVATE KEY-----[\s\S]*/g; + +/** + * `scheme://user:password@host` — the exact shape in the WALM-642 repro, and + * the shape of every connection string that carries its own credentials + * (`postgres://`, `mongodb+srv://`, `amqp://`, `redis://`, ...). + * + * Only the userinfo is replaced: scheme, host, port and path are the part of a + * connection string worth remembering. + * + * Two bounds on this pattern, both load-bearing: + * + * - The scheme repeat is CAPPED. Unbounded (`[A-Za-z0-9+.-]*`) it had to be + * tried and abandoned at every start offset, which is quadratic in the + * input: 30 KB took 317 ms, 60 KB 1254 ms and 120 KB 4814 ms on one core. + * `memwal_analyze` forwards a whole transcript and the sidecar is + * single-threaded, so a long paste was a stall for every other caller, not + * just for itself. 30 characters is longer than any registered scheme. + * - The userinfo groups exclude `?`, `=` and `&`, not just `/` and + * whitespace. Without that, + * `https://app.example.com:8443?owner=alice@corp.com` matched with + * `app.example.com` as the user and `8443?owner=alice` as the password: + * the host and port were destroyed, `corp.com` was promoted to hostname, + * the mangled fact was written to append-only storage, and the result told + * the agent a credential had been removed when there was none. A query + * string cannot appear before the userinfo in a real URL, so excluding + * them costs nothing but a password that literally contains one. + */ +const URL_USERINFO = + /([A-Za-z][A-Za-z0-9+.-]{0,30}:\/\/)([^\s/@:?=&]+):([^\s/@?=&]+)@/g; + +/** + * Authorization / cookie headers, value dropped, header name kept. + * + * The cookie value stops at whitespace rather than at the end of the line, and + * then continues across `;`-separated pairs. A real header (`Cookie: a=1; b=2`) + * is matched whole; a header quoted mid-sentence loses the cookie and not the + * rest of the sentence. Cookie values are token-shaped by spec, so the only + * thing this gives up is a value that contains a raw space — which is not + * legal in one anyway. + */ +const AUTH_HEADER = + /\b((?:proxy-)?authorization)(\s*[:=]\s*)(?:bearer|basic|token|digest)?\s*\S+/gi; +const COOKIE_HEADER = + /\b((?:set-)?cookie)(\s*[:=]\s*)[^\s;]+(?:\s*;\s*[^\s;]+)*/gi; + +/** + * `key=value` / `key: value` where the key names a credential. + * + * The separator must still follow the keyword (allowing one closing quote), + * which is what keeps ordinary prose out: "my password manager is 1Password" + * has no separator after "password" and does not match. + * + * ── What the gate has to let in ──────────────────────────────────────────── + * The keyword may start the identifier, sit mid-identifier, or follow a + * separator character, so the left gate is `\b` OR a lookbehind covering + * letters, digits, `_` and `-`: + * + * - `(?<=[a-z])` is what makes `delegatePrivateKey` match — MemWal's own + * worst secret, where `PrivateKey` sits mid-identifier and a plain `\b` + * never fires. + * - `_` is a WORD character, so `\b` does not fire between `_` and `P` + * either, and `_` is not `[a-z]`. That left the entire SCREAMING_SNAKE + * namespace open: `POSTGRES_PASSWORD=`, `AWS_SECRET_ACCESS_KEY=`, + * `DB_PASSWORD=`, `X_AUTH_TOKEN:`, `SESSION_SECRET=` and `my_api_key=` + * all passed through verbatim — the exact spelling a credential arrives in + * when someone pastes an env file or a shell export. + * + * And a closing quote may sit between the keyword and the separator, because + * that is what a JSON object looks like: `{"username": "alice", "password": + * "hunter2-prod-9xQ"}` matched nothing at all before. The quote is captured + * with the separator and written back, so the shape of the line survives; only + * the value is replaced. + * + * `access[_-]?keys?` is in the list for `AWS_SECRET_ACCESS_KEY`: `secret` is + * there, but the separator does not follow it, and no alternative covered the + * `ACCESS_KEY` that does precede the `=`. + */ +const CREDENTIAL_ASSIGNMENT = + /(?:\b|(?<=[a-z0-9_-]))(passwords?|passwd|pwd|passphrases?|api[_-]?keys?|apikeys?|secret[_-]?keys?|access[_-]?keys?|client[_-]?secrets?|secrets?|access[_-]?tokens?|refresh[_-]?tokens?|auth[_-]?tokens?|bearer[_-]?tokens?|tokens?|private[_-]?keys?|credentials?)(["'`]?\s*[:=]\s*)("[^"\n]*"|'[^'\n]*'|`[^`\n]*`|[^\s,;]+)/gi; + +/** + * Vendor-prefixed keys. Each prefix is issued by exactly one service and never + * appears at the head of ordinary text, so these are the highest-confidence + * patterns in the file — no context needed. + */ +const VENDOR_KEYS: RegExp[] = [ + /\bsk-ant-[A-Za-z0-9_-]{16,}/g, // Anthropic + /\bsk-proj-[A-Za-z0-9_-]{16,}/g, // OpenAI project + /\bsk-[A-Za-z0-9]{20,}/g, // OpenAI classic + /\b(?:ghp|gho|ghu|ghs|ghr)_[A-Za-z0-9]{20,}/g, // GitHub + /\bgithub_pat_[A-Za-z0-9_]{20,}/g, // GitHub fine-grained + /\b(?:AKIA|ASIA)[0-9A-Z]{16}\b/g, // AWS access key id + /\bxox[baprs]-[A-Za-z0-9-]{10,}/g, // Slack + /\bAIza[0-9A-Za-z_-]{35}\b/g, // Google API + /\bglpat-[A-Za-z0-9_-]{16,}/g, // GitLab + /\bnpm_[A-Za-z0-9]{36}\b/g, // npm + /\bSG\.[A-Za-z0-9_-]{16,}\.[A-Za-z0-9_-]{16,}/g, // SendGrid + /\b[sprk]k_(?:live|test)_[A-Za-z0-9]{16,}/g, // Stripe + /\bshp(?:at|ss|ca|pa)_[a-fA-F0-9]{32}\b/g, // Shopify + /\bdop_v1_[a-f0-9]{64}\b/g, // DigitalOcean + /\bhf_[A-Za-z0-9]{30,}\b/g, // Hugging Face +]; + +/** Three base64url segments — a signed JWT, whatever it encodes. */ +const JWT = /\beyJ[A-Za-z0-9_-]{8,}\.[A-Za-z0-9_-]{8,}\.[A-Za-z0-9_-]{8,}/g; + +/** + * A BIP-39 mnemonic, gated on the user naming it. + * + * Twelve consecutive lowercase words is also what an ordinary sentence looks + * like, so the words alone prove nothing; "seed phrase" / "mnemonic" / + * "recovery phrase" in front of them is what makes the match safe. A mnemonic + * pasted with no label is left to `SECRET_EXCLUSION_RULES` — carrying the + * 2048-word list here to catch it would be a lot of weight for a case the model + * rules already cover. + */ +const SEED_PHRASE = + /(?:\b|(?<=[a-z]))((?:seed|recovery|secret|mnemonic)[\s_-]*(?:phrase|words)|mnemonic)([^A-Za-z0-9]{0,4}(?:is|are)?[^A-Za-z0-9]{0,4})((?:[a-z]{3,8}[ \t]+){11,23}[a-z]{3,8})\b/gi; + +/** + * Key material in hex, identified by the LABEL beside it rather than by how + * random it looks. + * + * This exists because of one specific secret: `delegatePrivateKey` in + * `~/.memwal/credentials.json` is a 64-hex Ed25519 seed, the thing auth.ts + * marks "NEVER log this", and whoever holds it can read and write the user's + * memories until the delegate is revoked. It is the worst thing this product + * can leak — and it is pure lowercase hex, so it is deliberately excluded by + * `looksLikeSecretBlob` below and was sailing straight through. + * + * The exclusion is still right: a bare hex run is a git SHA, a Walrus blob id, + * a Sui object or account id, a content digest — the identifiers users most + * want remembered. So the discriminator is not entropy, it is the label. A hex + * run is removed only when a credential word sits next to it, which catches + * every shape the secret actually arrives in: + * + * delegatePrivateKey 4f3c... (prose / a pasted line) + * "delegatePrivateKey": "4f3c..." (the credentials.json file itself) + * my delegate private key is 4f3c... + * 4f3c... is my private key (label after the value) + * + * while `Pin the build to commit 4f2b8c1e...` and `my account id is 0x7f3a...` + * keep passing through untouched. That asymmetry is the whole design, and it is + * pinned from both sides in secret-redaction.test.ts. + */ +const HEX_RUN = /\b(?:0x)?[0-9a-fA-F]{32,}\b/g; + +/** How far either side of a hex run a label may sit — a few tokens. */ +const HEX_LABEL_WINDOW = 48; + +/** + * Words that make an adjacent hex run key material. + * + * Not anchored on a word boundary, so it matches inside a camelCase identifier + * (`delegatePrivateKey`). `credential(s)` is deliberately ABSENT: MemWal's own + * prose says "credentials.json" constantly, and a sentence naming that file + * next to a commit SHA would lose the SHA. + */ +const HEX_CREDENTIAL_LABEL = + /private[\s_-]*key|secret[\s_-]*key|delegate[\s_-]*key|delegate[\s_-]*private|signing[\s_-]*key|priv[\s_-]*key|api[\s_-]*key|access[\s_-]*key|auth[\s_-]*key|secret|seed|mnemonic|passphrase/i; + +/** + * The same, for a label that FOLLOWS the value ("4f3c... is my private key"). + * Tighter than the backward form — it has to march through the small joining + * phrase rather than search a window — because a trailing window would sweep + * in whatever sentence happens to come next. + */ +const HEX_LABEL_AFTER = + /^[^A-Za-z0-9]{0,4}(?:is|was)?[^A-Za-z0-9]{0,4}(?:my|the|our|his|her|their)?[^A-Za-z0-9]{0,4}(?:delegate[\s_-]*)?(?:private[\s_-]*key|secret[\s_-]*key|seed|mnemonic|passphrase|api[\s_-]*key)/i; + +/** Placeholders this module writes, for stripping out of a context window. */ +const PLACEHOLDER_RUN = /\[redacted:[a-z-]+\]/g; + +/** + * True when the span at [start, end) is itself part of a placeholder an earlier + * rule already wrote. + * + * `labelled-key-material` and `credential-assignment` are both 21 characters of + * exactly the alphabet {@link OPAQUE_RUN} scans for, so without this a run of + * redactions would start redacting its own output. + */ +function isInsidePlaceholder(text: string, start: number): boolean { + const before = text.slice(Math.max(0, start - 10), start); + return /\[redacted:$/.test(before); +} + +/** + * True when a credential word sits within `window` characters of [start, end). + * + * `window` is a parameter because the batch screen (see + * {@link sanitizeFactBatch}) looks across entry boundaries, where the label and + * the value are further apart than "a few tokens" by construction. + */ +function hasAdjacentCredentialLabel( + text: string, + start: number, + end: number, + window: number = HEX_LABEL_WINDOW, +): boolean { + // Placeholders left by earlier rules carry the words "secret" and "key", + // so a run of redactions would otherwise start labelling its own + // neighbours — and the neighbour after `private_key=[redacted:...]` is + // exactly the kind of bare SHA this rule must not touch. + const before = text + .slice(Math.max(0, start - window), start) + .replace(PLACEHOLDER_RUN, " "); + if (HEX_CREDENTIAL_LABEL.test(before)) return true; + return HEX_LABEL_AFTER.test( + text.slice(end, end + HEX_LABEL_WINDOW).replace(PLACEHOLDER_RUN, " "), + ); +} + +/** + * The same label gate, for key material that is not hex. + * + * `HEX_RUN` only ever looked at hex, and `HIGH_ENTROPY_CANDIDATE` demands 64+ + * characters, so everything in between passed with a label in front of it: + * `My AWS secret access key is wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY for + * prod` came back completely unchanged, three label words and all. In an AWS + * key pair that meant the non-secret `AKIA...` id was redacted by the vendor + * rule while the 40-character secret half survived — the wrong half of the + * pair, every time. + * + * Same discriminator as the hex rule, for the same reason: the LABEL, never the + * shape. A bare 43-character blob id, a commit SHA or a Sui object id with no + * credential word near it still passes, which is the property the rest of this + * file is built around. + */ +const OPAQUE_RUN = /[A-Za-z0-9+/_=-]{20,}/g; + +/** + * A shape guard on top of the label, so an ordinary long word next to the word + * "secret" is not mistaken for key material. + * + * A generated credential is mixed case, or carries digits, or carries base64 + * padding and separators; `my_api_key_rotation_policy` is none of those. This + * is not an entropy test and is not trying to be one — the label is still what + * decides. It only keeps prose out. + */ +function looksLikeOpaqueToken(token: string): boolean { + if (token.length < 20) return false; + const hasLower = /[a-z]/.test(token); + const hasUpper = /[A-Z]/.test(token); + const hasDigit = /[0-9]/.test(token); + const hasSymbol = /[+/=]/.test(token); + return (hasLower && hasUpper) || hasDigit || hasSymbol; +} + +/** + * A ≥64-character base64/base64url run carrying lower case, upper case AND a + * digit. See the false-positive note at the top: the three conditions together + * are what exclude blob ids, Sui addresses, git SHAs and hex digests, all of + * which are shorter, single-case, or both. + */ +const HIGH_ENTROPY_CANDIDATE = /\b[A-Za-z0-9+/_-]{64,}={0,2}/g; + +function looksLikeSecretBlob(token: string): boolean { + const body = token.replace(/=+$/, ""); + if (body.length < 64) return false; + if (!/[a-z]/.test(body)) return false; + if (!/[A-Z]/.test(body)) return false; + if (!/[0-9]/.test(body)) return false; + // Hex is single-case by convention and covers digests, SHAs and Sui ids; + // a mixed-case hex string is still far more likely a digest than a key. + if (/^[0-9a-fA-F]+$/.test(body)) return false; + return true; +} + +/** + * Unambiguously pasted third-party material. + * + * Narrow on purpose — see the false-positive note. An agent saving a fact + * routinely quotes the user's own sentence, so a short quoted string is NOT + * treated as a paste. + */ +function isPastedContent(text: string): boolean { + const s = text.trim(); + if (/^```[\s\S]*```$/.test(s)) return true; + const lines = s.split(/\r?\n/).filter((l) => l.trim() !== ""); + if (lines.length >= 2 && lines.every((l) => /^\s*>/.test(l))) return true; + if (s.length >= 200 && /^["“][\s\S]*["”]$/.test(s)) return true; + return false; +} + +/** + * Is there still a fact here once the placeholders are taken out? + * + * Three words of two or more letters. "sk-live-..." on its own redacts to a + * bare placeholder and has none; "I keep the staging key in 1Password, + * api_key=..." keeps its sentence and has plenty. + */ +function hasSalvageableContent(text: string): boolean { + const withoutPlaceholders = text.replace(/\[redacted:[a-z-]+\]/g, " "); + const words = withoutPlaceholders.match(/[\p{L}\p{N}]{2,}/gu) ?? []; + return words.length >= 3; +} + +/** + * Strip credentials from one piece of text before it is forwarded to the SDK. + * + * Pure, synchronous and side-effect free: no logging, no I/O. The secret exists + * only in the caller's argument and never leaves this function. + */ +export function sanitizeFact(input: string): SanitizedText { + const original = input ?? ""; + + if (NO_SAVE_DIRECTIVE.test(original) || OFF_THE_RECORD.test(original)) { + return { + text: "", + changed: true, + kinds: [], + count: 0, + refusal: "no-save-directive", + }; + } + + if (isPastedContent(original)) { + return { + text: "", + changed: true, + kinds: [], + count: 0, + refusal: "pasted-content", + }; + } + + const kinds: RedactionKind[] = []; + let count = 0; + const hit = (kind: RedactionKind) => { + if (!kinds.includes(kind)) kinds.push(kind); + count += 1; + return placeholder(kind); + }; + + let text = original; + + // PEM first: it spans lines, and running the line-oriented patterns over a + // key body would shred it into several partial matches instead of one. + text = text.replace(PEM_CLOSED, () => hit("private-key-block")); + text = text.replace(PEM_OPEN, () => hit("private-key-block")); + + text = text.replace(URL_USERINFO, (_m, scheme: string) => { + hit("url-credentials"); + return `${scheme}${placeholder("url-credentials")}@`; + }); + + text = text.replace(AUTH_HEADER, (_m, name: string, sep: string) => { + hit("auth-header"); + return `${name}${sep}${placeholder("auth-header")}`; + }); + text = text.replace(COOKIE_HEADER, (_m, name: string, sep: string) => { + hit("auth-header"); + return `${name}${sep}${placeholder("auth-header")}`; + }); + + // Before the vendor patterns, so `api_key=sk-...` is reported once as an + // assignment rather than twice. + text = text.replace(CREDENTIAL_ASSIGNMENT, (_m, key: string, sep: string) => { + hit("credential-assignment"); + return `${key}${sep}${placeholder("credential-assignment")}`; + }); + + for (const pattern of VENDOR_KEYS) { + text = text.replace(pattern, () => hit("vendor-api-key")); + } + + text = text.replace(JWT, () => hit("jwt")); + + text = text.replace(SEED_PHRASE, (_m, label: string, sep: string) => { + hit("seed-phrase"); + return `${label}${sep}${placeholder("seed-phrase")}`; + }); + + // Label-gated, and therefore run over the text as it stands now: the + // window check reads the characters on either side, so it has to see the + // real neighbours rather than a half-rewritten string. + text = text.replace(HEX_RUN, (match: string, offset: number, whole: string) => + hasAdjacentCredentialLabel(whole, offset, offset + match.length) + ? hit("labelled-key-material") + : match, + ); + + // Everything else a label makes into key material: the 20+ character + // base64/base64url runs that are too short for the entropy rule and not hex + // enough for HEX_RUN. Same window, same gate, same asymmetry — a bare run + // with nothing calling it a key still passes. + text = text.replace(OPAQUE_RUN, (match: string, offset: number, whole: string) => { + if (isInsidePlaceholder(whole, offset)) return match; + if (!looksLikeOpaqueToken(match)) return match; + return hasAdjacentCredentialLabel(whole, offset, offset + match.length) + ? hit("labelled-key-material") + : match; + }); + + text = text.replace(HIGH_ENTROPY_CANDIDATE, (token: string) => + looksLikeSecretBlob(token) ? hit("high-entropy-secret") : token, + ); + + if (count === 0) { + return { text: original, changed: false, kinds: [], count: 0 }; + } + + // Collapse the whitespace a removed block leaves behind, so the stored fact + // does not carry the shape of what was taken out. + text = text.replace(/[ \t]{2,}/g, " ").replace(/\n{3,}/g, "\n\n").trim(); + + if (!hasSalvageableContent(text)) { + return { + text: "", + changed: true, + kinds, + count, + refusal: "credential-only", + }; + } + + return { text, changed: true, kinds, count }; +} + +/** Human-readable reason, for the note handed back to the agent. */ +export function refusalMessage(reason: RefusalReason): string { + switch (reason) { + case "no-save-directive": + return "the text says not to save it"; + case "credential-only": + return "the text was a credential with no fact around it"; + case "pasted-content": + return "the text is pasted third-party content, not a fact about the user"; + } +} + +/** + * The line appended to a tool result when something was removed. + * + * Names the kinds and nothing else — an agent needs to know a redaction + * happened so it does not tell the user the whole line was stored, and it never + * needs the value back. + */ +export function redactionNotice(kinds: RedactionKind[], count: number): string { + if (count === 0) return ""; + return ( + `Note: ${count} credential span(s) were removed before saving ` + + `(${kinds.join(", ")}). What was stored is the redacted text — the secret ` + + `was never sent to Walrus Memory and is not logged. Tell the user the ` + + `credential was left out; do not re-send it.` + ); +} + +/** + * The line for a text that was not saved at all. + */ +export function refusalNotice(reason: RefusalReason): string { + return ( + `NOT SAVED: ${refusalMessage(reason)}. Nothing was written to Walrus ` + + `Memory. Do not retry this text — if there is a durable fact in it, ` + + `restate the fact without the sensitive part and save that instead.` + ); +} + +/* ─────────────────────────────────────────────────────────────────────────── + * Screening a BATCH, not one string at a time (WALM-642). + * + * `sanitizeFact` sees one entry, and every label-gated rule in this file + * searches a window inside that one string. `memwal_remember_bulk` takes up to + * twenty of them, which is a way around all of it: put the label in one entry + * and the value in the next and each is individually unremarkable. + * + * ["my delegate private key for the mainnet account", + * "4f3c...789"] + * + * Both entries passed untouched, while the same words as ONE string were + * correctly redacted. That value is MemWal's own delegate private key — the + * thing this module calls the worst secret the product can leak — and an agent + * that paraphrases a user across two entries is all it takes. Nothing about it + * requires malice. + * + * So a batch is screened as a batch. The per-entry pass runs first and is + * unchanged, which keeps the refusal granularity that already works (one bad + * entry is dropped, the rest of the batch still lands). Then a second pass asks + * a question the first one cannot: is there a credential label in the text + * ADJACENT to this value, where adjacent means the entries either side of it as + * well as its own? + * + * A value that is an entry all by itself is screened against the whole batch + * rather than its neighbours. A bare token with no words around it has no + * meaning of its own to lose, and "which entry did the label end up in" is not + * something the user controls. + * ------------------------------------------------------------------------ */ + +/** What separates two entries when the batch is viewed as one passage. */ +const BATCH_SEPARATOR = "\n"; + +/** + * A label anywhere in `context`, on either side of [start, end). + * + * The per-entry gate measures a character window because it is looking inside + * one sentence. Across entries the unit is the entry: the caller passes exactly + * the text that counts as adjacent and this searches all of it, so there is no + * second magic number to keep in step with the first. + */ +function hasCredentialLabelInContext( + context: string, + start: number, + end: number, +): boolean { + const before = context.slice(0, start).replace(PLACEHOLDER_RUN, " "); + if (HEX_CREDENTIAL_LABEL.test(before)) return true; + const after = context.slice(end).replace(PLACEHOLDER_RUN, " "); + return HEX_CREDENTIAL_LABEL.test(after); +} + +/** Words left once the placeholders and one candidate token are removed. */ +function wordsAround(text: string, token: string): number { + const rest = text + .replace(PLACEHOLDER_RUN, " ") + .replace(token, " ") + .match(/[\p{L}\p{N}]{2,}/gu); + return rest?.length ?? 0; +} + +/** + * Screen a whole batch. One result per input, in input order. + * + * Drop-in for `inputs.map(sanitizeFact)` — every entry comes back with the same + * shape and the same per-entry refusals — plus the cross-entry pass above. + */ +export function sanitizeFactBatch(inputs: string[]): SanitizedText[] { + const perEntry = inputs.map((text) => sanitizeFact(text ?? "")); + // Nothing survived the per-entry pass, so there is nothing left to screen. + if (perEntry.every((r) => r.refusal)) return perEntry; + + // What each entry contributes as CONTEXT. A refused entry is never + // forwarded, but its words still say what the batch is about, so it keeps + // supplying label context from its original text. + const views = perEntry.map((r, i) => (r.refusal ? (inputs[i] ?? "") : r.text)); + const wholeBatch = views.join(BATCH_SEPARATOR); + // Placeholders carry the words "secret" and "key" (`high-entropy-secret` + // most obviously), so a batch where one entry was already redacted would + // otherwise label itself. + if (!HEX_CREDENTIAL_LABEL.test(wholeBatch.replace(PLACEHOLDER_RUN, " "))) { + return perEntry; + } + + return perEntry.map((result, i) => { + if (result.refusal) return result; + const own = views[i]; + const prev = i > 0 ? views[i - 1] : ""; + const next = i + 1 < views.length ? views[i + 1] : ""; + // The neighbours, with this entry in the middle, and this entry's + // offset inside it. + const neighbourhood = [prev, own, next].join(BATCH_SEPARATOR); + const ownStart = prev.length + BATCH_SEPARATOR.length; + + const kinds = [...result.kinds]; + let count = result.count; + let hitAny = false; + const text = own.replace( + OPAQUE_RUN, + (match: string, offset: number, whole: string) => { + if (isInsidePlaceholder(whole, offset)) return match; + if (!looksLikeOpaqueToken(match)) return match; + // An entry that is essentially just this token is screened + // against every entry; one with a sentence around it, against + // the entries either side. + const bare = wordsAround(own, match) < 3; + const context = bare ? wholeBatch : neighbourhood; + const start = bare + ? views.slice(0, i).reduce( + (n, v) => n + v.length + BATCH_SEPARATOR.length, + 0, + ) + offset + : ownStart + offset; + if (!hasCredentialLabelInContext(context, start, start + match.length)) { + return match; + } + if (!kinds.includes("labelled-key-material")) { + kinds.push("labelled-key-material"); + } + count += 1; + hitAny = true; + return placeholder("labelled-key-material"); + }, + ); + + if (!hitAny) return result; + const collapsed = text.replace(/[ \t]{2,}/g, " ").trim(); + if (!hasSalvageableContent(collapsed)) { + return { text: "", changed: true, kinds, count, refusal: "credential-only" }; + } + return { text: collapsed, changed: true, kinds, count }; + }); +} + +/* ─────────────────────────────────────────────────────────────────────────── + * Screening a PASSAGE (WALM-642). + * + * `NO_SAVE_DIRECTIVE`, `OFF_THE_RECORD` and `isPastedContent` are whole-string + * predicates, written for one fact and correct there. `memwal_analyze` applies + * them to a whole transcript, where "whole string" means something else + * entirely: one "don't save this part" line in a forty-turn conversation + * refused the entire passage — `refusal`, `text: ""`, all forty turns + * discarded, and a note telling the agent not to retry it. The call carried no + * `isError`, so the client saw a successful call that had saved nothing. + * + * The fix is the one `memwal_remember_bulk` already uses on entries: scope the + * refusal to the span that earned it. A passage is split into segments, each is + * screened on its own, the offending ones are dropped and named, and everything + * else is extracted from. + * + * Two deliberate wrinkles: + * + * - A no-save directive drops its NEIGHBOURS too, within its own paragraph. + * "My bank PIN is 4821" on one line and "don't save this" on the next is + * the ordinary way people write it, and dropping only the second line would + * save the first — a far worse outcome than losing a turn either side. + * - A fenced block WRAPPING THE WHOLE PASSAGE is unwrapped rather than + * refused, when it has no language tag and holds a transcript's worth of + * lines. That shape is what this tool documents as its canonical input, and + * refusing it discarded exactly the passages people most wanted analysed. A + * tagged fence (```json, ```sh) is code, a short one is a snippet, and a + * fence INSIDE the passage is still dropped as a paste — none of those + * change. + * ------------------------------------------------------------------------ */ + +/** + * Non-blank lines an untagged outer fence must hold before it is read as a + * transcript rather than a pasted snippet. + */ +const TRANSCRIPT_FENCE_MIN_LINES = 4; + +/** One dropped span, named by position and reason. Never carries its text. */ +export interface DroppedSpan { + /** 1-based line number in the passage as it was received. */ + line: number; + reason: RefusalReason; +} + +export interface SanitizedPassage extends SanitizedText { + /** Spans removed before extraction. Empty when the passage came through whole. */ + dropped: DroppedSpan[]; + /** Non-blank segments the passage was split into. */ + segments: number; +} + +/** Strip one outer fence when it is a transcript rather than a code paste. */ +function unwrapTranscriptFence(text: string): string { + const s = text.trim(); + const m = /^```([^\n]*)\n([\s\S]*?)\n?```$/.exec(s); + if (!m) return text; + // A language tag is an author saying "this is code", so take them at their + // word and leave it to the per-segment paste rule. + if (m[1].trim() !== "") return text; + const body = m[2]; + const lines = body.split(/\r?\n/).filter((l) => l.trim() !== ""); + if (lines.length < TRANSCRIPT_FENCE_MIN_LINES) return text; + return body; +} + +interface PassageSegment { + text: string; + /** 1-based line the segment starts on. */ + line: number; + blank: boolean; +} + +/** + * Split a passage into the units a refusal may apply to. + * + * One line per segment, except that a fenced block and a run of `>` quotation + * stay whole — both are multi-line by nature, and `isPastedContent` can only + * recognise them as one piece. + */ +function splitPassage(text: string): PassageSegment[] { + const lines = text.split(/\r?\n/); + const segments: PassageSegment[] = []; + for (let i = 0; i < lines.length; i++) { + const line = lines[i]; + if (/^\s*```/.test(line)) { + const start = i; + const block = [line]; + i++; + while (i < lines.length) { + block.push(lines[i]); + if (/^\s*```/.test(lines[i])) break; + i++; + } + segments.push({ text: block.join("\n"), line: start + 1, blank: false }); + continue; + } + if (/^\s*>/.test(line)) { + const start = i; + const block = [line]; + while (i + 1 < lines.length && /^\s*>/.test(lines[i + 1])) { + block.push(lines[++i]); + } + segments.push({ text: block.join("\n"), line: start + 1, blank: false }); + continue; + } + segments.push({ text: line, line: i + 1, blank: line.trim() === "" }); + } + return segments; +} + +/** + * Strip credentials from a passage, dropping only the spans that must go. + * + * `text` is what is safe to forward, `dropped` says what was removed and why, + * and `refusal` is set only when nothing survived at all. + */ +export function sanitizePassage(input: string): SanitizedPassage { + const segments = splitPassage(unwrapTranscriptFence(input ?? "")); + const results = segments.map((s) => + s.blank ? null : sanitizeFact(s.text), + ); + const dropped: (RefusalReason | null)[] = results.map((r) => r?.refusal ?? null); + + // A directive takes its immediate neighbours with it, unless a blank line + // stands between them — see the note above. Computed against the ORIGINAL + // drop list so one directive cannot cascade down a whole passage. + const directive = dropped.map((d) => d === "no-save-directive"); + for (let i = 0; i < segments.length; i++) { + if (!directive[i]) continue; + for (const j of [i - 1, i + 1]) { + if (j < 0 || j >= segments.length) continue; + if (segments[j].blank) continue; + if (dropped[j] === null) dropped[j] = "no-save-directive"; + } + } + + const kinds: RedactionKind[] = []; + let count = 0; + const kept: string[] = []; + const droppedSpans: DroppedSpan[] = []; + let nonBlank = 0; + for (let i = 0; i < segments.length; i++) { + const segment = segments[i]; + if (segment.blank) { + kept.push(""); + continue; + } + nonBlank += 1; + const reason = dropped[i]; + if (reason) { + droppedSpans.push({ line: segment.line, reason }); + continue; + } + const result = results[i]!; + count += result.count; + for (const kind of result.kinds) { + if (!kinds.includes(kind)) kinds.push(kind); + } + kept.push(result.text); + } + + const text = kept + .join("\n") + .replace(/[ \t]{2,}/g, " ") + .replace(/\n{3,}/g, "\n\n") + .trim(); + const changed = count > 0 || droppedSpans.length > 0; + + if (!hasSalvageableContent(text)) { + // Nothing usable left. Report the reason that took the most of it, so + // the agent is told the truth about why rather than a generic refusal. + const tally = new Map(); + for (const span of droppedSpans) { + tally.set(span.reason, (tally.get(span.reason) ?? 0) + 1); + } + const [top] = [...tally.entries()].sort((a, b) => b[1] - a[1]); + return { + text: "", + changed: true, + kinds, + count, + refusal: top?.[0] ?? "credential-only", + dropped: droppedSpans, + segments: nonBlank, + }; + } + + return { text, changed, kinds, count, dropped: droppedSpans, segments: nonBlank }; +} + +/** + * The line naming spans that were dropped from a passage but not the whole of + * it. Positions and reasons only — never the text that was removed. + */ +export function droppedSpanNotice(dropped: DroppedSpan[], segments: number): string { + if (dropped.length === 0) return ""; + const detail = dropped + .map((d) => `line ${d.line} (${refusalMessage(d.reason)})`) + .join("; ") + return ( + `Note: ${dropped.length} of ${segments} span(s) were dropped before ` + + `extraction and nothing from them was saved — ${detail}. The rest of the ` + + `passage was extracted from normally. Tell the user which part was left ` + + `out; do not re-send it.` + ); +} diff --git a/services/server/scripts/mcp/tools/remember-bulk.ts b/services/server/scripts/mcp/tools/remember-bulk.ts index 69ed965f2..133c126b3 100644 --- a/services/server/scripts/mcp/tools/remember-bulk.ts +++ b/services/server/scripts/mcp/tools/remember-bulk.ts @@ -3,6 +3,22 @@ import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; import type { MemWalSession } from "../auth.js"; import { TOOL_METADATA } from "./annotations.js"; import { wrapTool, explorerFooter } from "./util.js"; +import { SECRET_EXCLUSION_RULES, AUTO_SAVE_OPT_IN_RULE } from "./memory-policy.js"; +import { + sanitizeFactBatch, + redactionNotice, + refusalMessage, + type RedactionKind, +} from "./redaction.js"; +import { + REMEMBER_POLL_INTERVAL_MS, + REMEMBER_WAIT_MS, + pendingBulkMessage, + withAcceptDeadline, + withWaitDeadline, + withRelayerRetry, + isStillRunning, +} from "./remember-wait.js"; const REMEMBER_BULK_INPUT = { facts: z @@ -10,7 +26,7 @@ const REMEMBER_BULK_INPUT = { .min(1) .max(20) .describe( - "Array of complete, detailed fact statements to save (1-20). Each entry is one full fact — do not summarize or merge them." + "Array of complete, detailed fact statements to save (1-20). Each entry is one full fact — do not summarize or merge them. Leave credentials out: passwords, API keys, tokens, private keys, seed phrases, auth headers and URLs with an embedded user:password are stripped from each entry before the write and never stored." ), namespace: z .string() @@ -21,11 +37,20 @@ const REMEMBER_BULK_INPUT = { } as const; /** - * memwal_remember_bulk — persist several durable facts in one batched request - * and return only once every job reaches a terminal state. Wraps the SDK's - * `rememberBulkAndWait` (embed + SEAL-encrypt all items concurrently, upload - * N blobs in parallel). Prefer this over N separate `memwal_remember` calls - * when you learned multiple distinct facts at once. + * memwal_remember_bulk — persist several durable facts in one batched request. + * + * Mirrors `memwal_remember`: returns once every job reaches a terminal state + * if that happens inside `REMEMBER_WAIT_MS`, and otherwise hands back the + * job_ids saying plainly that the facts are not saved yet. + * + * Blocking here was worse than blocking on a single fact, not better. The + * server instructions steer an agent to this tool whenever it learned more + * than one thing, so it is the common path — and a batch is N separate Walrus + * writes contending for the same upload slots + * (`WALRUS_UPLOAD_PER_WALLET_CONCURRENCY` defaults to 1), so they land one + * after another rather than together. Against a 30–75s single-write spread a + * five-fact batch could exhaust the old fixed 120s budget outright and return + * nothing but timeouts, having blocked the agent for two minutes first. */ export function registerRememberBulkTool( server: McpServer, @@ -36,33 +61,184 @@ export function registerRememberBulkTool( { ...TOOL_METADATA.memwal_remember_bulk, description: - "Save multiple durable facts in one call. Use when you learned several distinct facts at once (onboarding details, a list of preferences, decisions from a discussion). Pass an array of complete fact statements (max 20) — do not summarize. Prefer this over repeated memwal_remember calls.", + "Save multiple durable facts in one call. Use when you learned several distinct facts at once (onboarding details, a list of preferences, decisions from a discussion). Pass an array of complete fact statements (max 20) — do not summarize. Prefer this over repeated memwal_remember calls. By default this returns in ~1s once the relayer has accepted the batch (job_ids) — the Walrus writes are still in flight and the facts are NOT stored yet. Do not claim they are saved. Resolve with memwal_remember_status(job_ids). blob_ids in the same reply mean they landed inside an optional wait budget (MEMWAL_MCP_REMEMBER_WAIT_MS). " + + AUTO_SAVE_OPT_IN_RULE + + " " + + SECRET_EXCLUSION_RULES + + " Walrus storage is append-only: a stored secret cannot be deleted, so each entry is stripped of credential shapes before writing and entries that are nothing but a secret are dropped, with a note saying which. The batch is screened as a whole, so splitting a credential's label into one entry and its value into another does not get it past the filter.", inputSchema: REMEMBER_BULK_INPUT, }, wrapTool<{ facts: string[]; namespace?: string }>(session, "memwal_remember_bulk", async ({ facts, namespace }) => { - const items = facts.map((text) => ({ text, namespace })); - const result = await session.memwal.rememberBulkAndWait(items, { - timeoutMs: 120_000, + // Every entry is sanitized BEFORE the batch is handed to the SDK. + // Walrus is append-only, so a credential that lands cannot be + // taken back (WALM-642). An entry that survives keeps its safe + // part; an entry that is only a secret — or that the user asked + // not to save — is dropped from the batch rather than the whole + // call failing, so the other facts still land. + // + // Screened as a BATCH, not entry by entry: every label-gated rule + // searches a window inside one string, so a label in one entry and + // its value in the next defeated all of them — including the one + // that exists for MemWal's own delegate private key. See + // `sanitizeFactBatch`. + const screened = sanitizeFactBatch(facts).map((result, index) => ({ + index, + result, + })); + const kept = screened.filter((s) => !s.result.refusal); + const dropped = screened.filter((s) => s.result.refusal); + const droppedNote = dropped.length + ? `\n\nNOT SAVED (${dropped.length}): ` + + dropped + .map((d) => `#${d.index + 1} — ${refusalMessage(d.result.refusal!)}`) + .join("; ") + + ". Do not re-send those; restate any durable fact without the sensitive part instead." + : ""; + + if (kept.length === 0) { + return { + content: [ + { + type: "text" as const, + text: + `Nothing was saved to Walrus Memory: every fact in this batch was ` + + `withheld.${droppedNote}`, + }, + ], + }; + } + + const redactedKinds: RedactionKind[] = []; + let redactedCount = 0; + for (const s of kept) { + redactedCount += s.result.count; + for (const kind of s.result.kinds) { + if (!redactedKinds.includes(kind)) redactedKinds.push(kind); + } + } + const policyNote = + [redactionNotice(redactedKinds, redactedCount), droppedNote.trim()] + .filter(Boolean) + .join("\n\n"); + + // The only texts anything below may echo or forward. The originals + // still hold the secret and must not reach a result line. + const safeFacts = kept.map((s) => s.result.text); + const items = safeFacts.map((text) => ({ text, namespace })); + // Two steps rather than `rememberBulkAndWait`, for the same reason + // `memwal_remember` splits them: acceptance is the part that must + // succeed, the wait is a courtesy we cut short. + const accepted = await withAcceptDeadline( + // Safe to wrap despite bulk having no idempotency key: the + // retry only fires on rejections that never reached the + // handler, so no job row can exist to duplicate. + withRelayerRetry( + () => session.memwal.rememberBulkAsync(items), + "save these facts", + ), + "memwal_remember_bulk batch", + { idempotent: false }, + ); + + // Pair each job with its fact up front. Every later branch needs + // it, and the relayer returns job_ids in input order. + const entries = accepted.job_ids.map((jobId, i) => ({ + jobId, + text: safeFacts[i] ?? "", + })); + + const withNotice = (body: string) => + policyNote ? `${body}\n\n${policyNote}` : body; + + const pending = (waitedMs: number) => ({ + content: [ + { + type: "text" as const, + text: withNotice(pendingBulkMessage(entries, waitedMs)), + }, + ], }); + + if (REMEMBER_WAIT_MS === 0) return pending(0); + + const startedAt = Date.now(); + const namespaces = items.map( + (item) => item.namespace ?? session.namespace ?? "default" + ); + // `waitForRememberJobs` never throws on expiry — it reports the + // stragglers as `timeout` per item, so a batch can come back part + // landed and part still in flight. `withWaitDeadline` around it + // does throw, though: a relayer that goes quiet mid-poll raises + // MemWalRelayerUnresponsive, and letting that propagate discards + // every job_id in the batch, leaving the caller nothing to settle + // accepted writes with. `memwal_remember` already degrades to its + // pending branch here; so does this. + let result; + try { + result = await withWaitDeadline( + session.memwal.waitForRememberJobs(accepted.job_ids, namespaces, { + timeoutMs: REMEMBER_WAIT_MS, + pollIntervalMs: REMEMBER_POLL_INTERVAL_MS, + }), + REMEMBER_WAIT_MS, + ); + } catch (err) { + if (!isStillRunning(err)) throw err; + return pending(Date.now() - startedAt); + } + const waitedMs = Date.now() - startedAt; + + const unfinished = result.results.flatMap((r, i) => + r.status === "timeout" ? [{ jobId: r.id, text: safeFacts[i] ?? "" }] : [] + ); + // Nothing landed inside the budget — the ordinary outcome when the + // queue is busy. Say so once rather than printing N timeout rows. + if (unfinished.length === result.results.length) return pending(waitedMs); + const lines = result.results.map((r, i) => { // Label each result with its source fact by index. The SDK // returns results in input order, but guard against a length / // ordering mismatch so we never print "— undefined". - const text = facts[i] ?? ""; - const blob = r.blob_id ? ` blob_id=${r.blob_id}` : ""; + // Echo sanitized text only (WALM-642) — never the original facts[]. + const text = safeFacts[i] ?? ""; + // Only a settled row may show a blob_id — see remember-status.ts + // for why an unfinished write can already carry a real one. + const blob = r.status === "done" && r.blob_id ? ` blob_id=${r.blob_id}` : ""; const err = r.error ? ` error=${r.error}` : ""; - return `${i + 1}. [${r.status}]${blob}${err}${text ? ` — ${text}` : ""}`; + // `timeout` is not a failure — the write is still running and + // its job_id is how the caller settles it later. + const state = r.status === "timeout" ? `still uploading, job_id=${r.id}` : r.status; + return `${i + 1}. [${state}]${blob}${err}${text ? ` — ${text}` : ""}`; }); - const summary = `Saved ${result.succeeded}/${result.total} fact(s) to Walrus Memory (failed=${result.failed}).`; + // `result.failed` is total-minus-succeeded, so it counts a + // still-uploading write as failed — while the tail below says that + // same job is on its way. An agent reading `failed=` re-sends an + // in-flight write, which is the duplicate this branch exists to + // avoid. Count only what actually reached a terminal failure. + const reallyFailed = result.results.filter( + (r) => r.status !== "done" && r.status !== "timeout", + ).length; + const summary = + `Saved ${result.succeeded}/${result.total} fact(s) to Walrus Memory` + + (reallyFailed ? ` (failed=${reallyFailed})` : "") + + (unfinished.length ? ` (${unfinished.length} still uploading)` : "") + + "."; const footer = result.succeeded > 0 ? `\n\n${explorerFooter()}` : ""; + const tail = unfinished.length + ? `\n\n${unfinished.length} write(s) are STILL UPLOADING and are NOT saved yet. ` + + `Do not claim those facts are stored; resolve them with memwal_remember_status ` + + `using job_ids=[${unfinished.map((u) => u.jobId).join(", ")}]. Do not re-send ` + + `them — that queues duplicates behind the originals.` + : ""; return { content: [ { type: "text", - text: - lines.length > 0 + text: withNotice( + (lines.length > 0 ? `${summary}\n\n${lines.join("\n")}${footer}` - : `${summary}${footer}`, + : `${summary}${footer}`) + tail, + ), }, ], }; diff --git a/services/server/scripts/mcp/tools/remember-status.ts b/services/server/scripts/mcp/tools/remember-status.ts new file mode 100644 index 000000000..56ec5bfd2 --- /dev/null +++ b/services/server/scripts/mcp/tools/remember-status.ts @@ -0,0 +1,381 @@ +import { z } from "zod"; +import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; +import type { MemWalSession } from "../auth.js"; +import { TOOL_METADATA } from "./annotations.js"; +import { wrapTool, walruscanBlobUrl, explorerFooter } from "./util.js"; +import { + REMEMBER_POLL_INTERVAL_MS, + isStillRunning, + isRateLimited, + nameJobError, + rateLimitedError, + withAcceptDeadline, + withWaitDeadline, +} from "./remember-wait.js"; + +/** + * Ceiling on a single status wait, held under the MCP client's own deadline. + * + * `@modelcontextprotocol/sdk` times a request out after + * `DEFAULT_REQUEST_TIMEOUT_MSEC` = 60s unless the caller overrides it. A tool + * that waits the full 60s therefore loses every race it enters: the client + * gives up first and the agent sees `MCP error -32001: Request timed out` + * instead of the answer the tool was about to return. Confirmed live against + * the production relayer — `waitMs: 60000` on a three-job batch returned + * exactly that, with no way for the caller to tell a slow write from a broken + * tool. + * + * 45s leaves room for the round trip and still covers the median write. A job + * that outlives it is not lost: the job_id stays valid and the caller asks + * again, which is the whole point of this tool being separate from the write. + */ +const MAX_STATUS_WAIT_MS = 45_000; + +/** Wait applied when the caller does not choose one. Declared above + * `STATUS_INPUT` because the tool description interpolates it. */ +const DEFAULT_STATUS_WAIT_MS = 10_000; + +/** Matches the bulk write cap, so a whole batch settles in one call. */ +const MAX_STATUS_JOB_IDS = 20; + +const STATUS_INPUT = { + job_id: z + .string() + .min(1) + .optional() + .describe("The job_id returned by memwal_remember when the write had not landed yet."), + job_ids: z + .array(z.string().min(1)) + .min(1) + .max(MAX_STATUS_JOB_IDS) + .optional() + .describe( + "The job_ids returned by memwal_remember_bulk when the writes had not landed yet. Pass the whole batch in one call rather than polling each id separately. Supply either this or job_id." + ), + waitMs: z + .number() + .int() + .min(0) + .max(MAX_STATUS_WAIT_MS) + .optional() + .describe( + `How long to wait for the job to finish, in milliseconds (0-${MAX_STATUS_WAIT_MS}, default ${DEFAULT_STATUS_WAIT_MS}). Pass 0 to read the current state without waiting.` + ), +} as const; + +/** + * memwal_remember_status — resolve a remember job that `memwal_remember` + * handed back as still in flight. + * + * This is the other half of the bounded wait: `memwal_remember` refuses to + * claim a fact is saved when it isn't, so something has to be able to say + * whether it landed. Three outcomes, kept distinct because an agent acts + * differently on each — saved (blob_id), still running (ask again), failed + * (the fact is NOT stored and must be re-sent). + */ +export function registerRememberStatusTool( + server: McpServer, + session: MemWalSession +): void { + server.registerTool( + "memwal_remember_status", + { + ...TOOL_METADATA.memwal_remember_status, + description: + "Check whether in-flight Walrus Memory writes have landed. Call this with the job_id memwal_remember returned, or job_ids from memwal_remember_bulk, when the write was reported NOT saved yet. Returns the blob_id once stored, reports that it is still uploading (call again with the ids still listed), or reports that it failed — in which case the fact was never stored and you should send it again. A batch can come back mixed, so read every line before telling the user anything is saved.", + inputSchema: STATUS_INPUT, + }, + wrapTool<{ job_id?: string; job_ids?: string[]; waitMs?: number }>( + session, + "memwal_remember_status", + async ({ job_id, job_ids, waitMs }) => { + const budget = waitMs ?? DEFAULT_STATUS_WAIT_MS; + + // Exactly one of the two. Enforced here rather than in the + // schema because `registerTool` takes a raw shape, which has + // nowhere to hang a cross-field refinement. + const batch = job_ids ?? []; + if (batch.length > 0 && job_id) { + throw new Error( + "Pass job_id or job_ids, not both — they would describe different writes." + ); + } + if (batch.length > 0) return await settleBatch(session, batch, budget); + if (!job_id) { + throw new Error( + "Pass job_id (from memwal_remember) or job_ids (from memwal_remember_bulk)." + ); + } + + // A zero budget means "read the current state", which is a + // single GET. waitForRememberJob cannot express that: it + // sleeps before its first poll, so a 0ms deadline would + // return "still running" without ever asking the relayer. + if (budget === 0) { + return settleFromStatus( + job_id, + await withAcceptDeadline( + session.memwal.getRememberStatus(job_id), + "status read", + { idempotent: true }, + ), + ); + } + + try { + const result = await withWaitDeadline( + session.memwal.waitForRememberJob(job_id, { + timeoutMs: budget, + pollIntervalMs: REMEMBER_POLL_INTERVAL_MS, + }), + budget, + ); + return saved(result.blob_id, result.namespace); + } catch (err) { + if (isStillRunning(err)) { + // Same masking as the batch path: the SDK's poll loop + // hides a refused poll behind its own timeout, so + // confirm with one direct read before calling it + // progress. Interpret the probe OUTSIDE the catch so + // a terminal job (`failed` / `not_found`) is not + // swallowed back into "still uploading". + let status; + try { + status = await withAcceptDeadline( + session.memwal.getRememberStatus(job_id), + "status read", + { idempotent: true }, + ); + } catch (probe) { + if (isRateLimited(probe)) { + throw rateLimitedError( + probe, + "check whether the write landed" + ); + } + // A failed read is not evidence about the job. + return stillRunning(job_id); + } + return settleFromStatus(job_id, status); + } + throw nameJobError(err); + } + } + ) + ); +} + +/** + * Settle a batch of job_ids in one report. + * + * Unlike the single-job path this never throws on a failed job: a batch + * routinely comes back mixed, and throwing on the first failure would hide the + * blob_ids of the writes that did land — the agent would have no way to tell + * which facts still need re-sending. + */ +async function settleBatch( + session: MemWalSession, + jobIds: string[], + budgetMs: number +) { + // A zero budget is a single batched read, the same shortcut the one-job + // path takes: `waitForRememberJobs` sleeps before its first poll, so a 0ms + // deadline would report everything as still running without ever asking. + let rows = + budgetMs === 0 + ? ( + await withAcceptDeadline( + session.memwal.getRememberBulkStatus(jobIds), + "batch status read", + { idempotent: true }, + ) + ).results.map((r) => ({ + id: r.job_id, + status: r.status, + blob_id: r.blob_id ?? "", + error: r.error, + })) + : ( + await withWaitDeadline( + session.memwal.waitForRememberJobs(jobIds, [], { + timeoutMs: budgetMs, + pollIntervalMs: REMEMBER_POLL_INTERVAL_MS, + }), + budgetMs, + ) + ).results.map((r) => ({ + id: r.id, + status: r.status, + blob_id: r.blob_id, + error: r.error, + })); + + // `waitForRememberJobs` polls internally and swallows each poll's error, + // stamping every row "polling timed out" when the budget runs out. A batch + // whose polls were all REFUSED — a 429 on the status endpoint — is + // therefore indistinguishable from one that is genuinely still uploading, + // and reporting the refusal as progress is what sends an agent back to + // poll again on a budget it has already spent. + // + // Nothing moving at all is the shape that refusal takes, so confirm it + // with one direct read. Only then: if any row settled, the polls were + // clearly getting through and no probe is warranted. + const nothingMoved = + budgetMs > 0 && rows.length > 0 && rows.every((r) => r.status === "timeout"); + if (nothingMoved) { + try { + const probed = await withAcceptDeadline( + session.memwal.getRememberBulkStatus(jobIds), + "batch status read", + { idempotent: true }, + ); + // Merge by job id rather than replacing the list. The relayer is + // not obliged to echo one row per requested id in the order asked: + // omit the ids it cannot find and a wholesale replace would drop + // them from the report entirely, so a job that vanished between + // the wait and the probe is never mentioned and the caller never + // learns it has to re-send that fact. Keep every id, and take the + // probe's answer only where it gave one. + const byId = new Map(probed.results.map((r) => [r.job_id, r])); + rows = rows.map((row) => { + const fresh = byId.get(row.id); + if (!fresh) return row; + return { + id: row.id, + status: fresh.status, + blob_id: fresh.blob_id ?? "", + error: fresh.error, + }; + }); + } catch (err) { + if (isRateLimited(err)) { + throw rateLimitedError(err, "check whether the writes landed"); + } + // Any other probe failure is not evidence about the jobs; fall + // back to what the wait already reported rather than inventing an + // outcome from a failed read. + } + } + + // `timeout` (the waited path) and pending/running/uploaded (the immediate + // read) are the same thing to a caller: still in flight, ask again. + const inFlight = rows.filter( + (r) => r.status !== "done" && r.status !== "failed" && r.status !== "not_found" + ); + const failed = rows.filter((r) => r.status === "failed" || r.status === "not_found"); + const done = rows.filter((r) => r.status === "done"); + + const lines = rows.map((r, i) => { + // Only a settled row may show a blob_id. An `uploaded` row already has + // a real one — the blob IS on Walrus — but the job is not finished: + // `persist_uploaded_state` writes blob_id at that point, and the + // SetMetadataAndTransfer + insert_vector legs that make the fact + // recallable, and hand the blob object to the user, still have to run + // before `done`. Printing it next to "still uploading" contradicts this + // server's own instruction that "only a blob_id in the tool reply means + // the fact is already stored", and an agent that believes the blob_id + // tells the user a fact is saved that a later failure can still lose. + // The single-job path never shows one on a running job (`stillRunning`); + // neither does this. + const blob = r.status === "done" && r.blob_id ? ` blob_id=${r.blob_id}` : ""; + // `waitForRememberJobs` stamps "polling timed out after Nms" on rows + // that simply had not landed when the budget ran out. That is our + // clock expiring, not the job failing, so showing it as `error=` next + // to "still uploading" reads like the write broke when it is still on + // its way. Only a terminal row gets to explain itself. + const terminal = r.status === "failed" || r.status === "not_found"; + const err = terminal && r.error ? ` error=${r.error}` : ""; + const state = + r.status === "done" + ? "saved" + : r.status === "failed" || r.status === "not_found" + ? `NOT STORED (${r.status})` + : "still uploading"; + return `${i + 1}. [${state}] job_id=${r.id}${blob}${err}`; + }); + + const parts = [ + `${done.length}/${rows.length} saved` + + (inFlight.length ? `, ${inFlight.length} still uploading` : "") + + (failed.length ? `, ${failed.length} NOT stored` : "") + + ".", + lines.join("\n"), + ]; + if (inFlight.length) { + parts.push( + `Still uploading — call memwal_remember_status again with job_ids=[${inFlight + .map((r) => r.id) + .join(", ")}]. Do not re-send those facts; that queues duplicates.` + ); + } + if (failed.length) { + parts.push( + `NOT stored — these facts were never saved and must be sent again with ` + + `memwal_remember or memwal_remember_bulk.` + ); + } + if (done.length) parts.push(explorerFooter()); + + return { content: [{ type: "text" as const, text: parts.join("\n\n") }] }; +} + +/** Map a job-status read onto the three outcomes this tool keeps distinct. + * + * Shared by the zero-budget path and the confirming probe after a wait + * timeout, so a terminal job cannot be reported as still uploading on one + * path and as an error on the other. */ +function settleFromStatus( + jobId: string, + status: { status: string; blob_id?: string; namespace?: string; error?: string }, +) { + if (status.status === "done") { + return saved(status.blob_id ?? "", status.namespace); + } + if (status.status === "failed") { + throw nameJobError( + Object.assign( + new Error(`remember job failed: ${status.error ?? "unknown error"}`), + { status: 500, jobId }, + ), + ); + } + if (status.status === "not_found") { + throw nameJobError( + Object.assign(new Error(`remember job not found: ${jobId}`), { + status: 404, + jobId, + }), + ); + } + return stillRunning(jobId, status.status); +} + +function saved(blobId: string, namespace?: string) { + return { + content: [ + { + type: "text" as const, + text: + `Saved to Walrus Memory. blob_id=${blobId}` + + (namespace ? ` namespace=${namespace}` : "") + + `\nExplorer: ${walruscanBlobUrl(blobId)}`, + }, + ], + }; +} + +function stillRunning(jobId: string, state?: string) { + return { + content: [ + { + type: "text" as const, + text: + `STILL UPLOADING — not saved yet${state ? ` (state: ${state})` : ""}.\n` + + `job_id=${jobId}\n` + + `The job is still queued or uploading. Call memwal_remember_status again ` + + `with this job_id. Do not re-send the fact with memwal_remember — that ` + + `queues a duplicate behind this one.`, + }, + ], + }; +} diff --git a/services/server/scripts/mcp/tools/remember-wait.ts b/services/server/scripts/mcp/tools/remember-wait.ts new file mode 100644 index 000000000..8483a2b6b --- /dev/null +++ b/services/server/scripts/mcp/tools/remember-wait.ts @@ -0,0 +1,563 @@ +/** + * Shared wait-budget and job-error handling for the two remember tools. + * + * `memwal_remember` used to block on `rememberAndWait` until the write reached + * `done`. Measured against the production relayer that is 30–75s for a single + * short fact, and the agent can do nothing with the wait: the job is durably + * accepted ~1s in, and everything after that is upload queue time + * (`WALRUS_UPLOAD_PER_WALLET_CONCURRENCY` defaults to 1, so a second write + * waits for the first). + * + * So the tool now waits a bounded budget and then hands the caller a job_id. + * It does NOT claim the fact is saved when it isn't — a job can still fail + * after acceptance (one observed failure: "Memory encryption backend is + * unavailable" 31.6s in, from the SEAL sidecar being unreachable). + */ +import { createHash } from "node:crypto"; + +import { createLogger } from "../logger.js"; + +const log = createLogger("mcp"); + +/** Window over which the same (namespace, text) resolves to one key. */ +const IDEMPOTENCY_BUCKET_MS = 30 * 60 * 1000; + +/** + * Content-derived idempotency key for a single `remember`. + * + * Computed HERE rather than relied on from the SDK. The sidecar installs the + * published `@mysten-incubation/memwal` (0.1.7), whose `rememberAsync` mints a + * `crypto.randomUUID()` and caches it per client instance — and the sidecar + * builds a fresh client per transport session, so a retry after a reconnect + * gets a brand-new key and the relayer stores the fact a second time at full + * cost. The accept-timeout message promises the caller a retry is safe, so the + * key that makes it safe has to exist in the version actually running, not in + * an unreleased source tree. + * + * The bucket bounds the collapse: `remember_jobs` rows are never pruned, so an + * unbucketed key would dedupe against a job from any point in history and a + * deliberate re-save of a since-deleted fact would hand back the old blob id. + * Retries happen seconds after the original, so 30 minutes covers them. + * + * `/api/remember/bulk` takes no key at all, which is why the bulk tool stays + * `idempotent: false` instead of calling this. + */ +export function derivedIdempotencyKey(namespace: string | undefined, text: string): string { + const bucket = Math.floor(Date.now() / IDEMPOTENCY_BUCKET_MS); + const digest = createHash("sha256") + .update(`${bucket}\0${namespace ?? ""}\0${text}`) + .digest("hex"); + return `r1-${digest}`; +} + +/** + * Hard ceiling on the wait budget. 90s matches the timeout the tool used + * while it still blocked to terminal, so an operator can restore the old + * always-block behaviour but cannot push the call past what MCP clients + * are willing to wait for. + */ +export const MAX_REMEMBER_WAIT_MS = 90_000; + +/** + * Default wait before `memwal_remember` hands back a job_id. + * + * Zero — return at accept. The MCP SDK's default `tools/call` timeout is + * 60s; a budget that actually waits for Walrus (30–75s) overruns that, and + * any in-between budget that fits still returns pending on nearly every + * call. Restore block-until-done with `MEMWAL_MCP_REMEMBER_WAIT_MS=90000`. + */ +const DEFAULT_REMEMBER_WAIT_MS = 0; + + + +/** + * How long `memwal_remember` waits for the write to land before returning a + * job_id instead. `0` returns at accept (~1s). + * + * Read once — it cannot change mid-process — but validated, because a typo'd + * value must not silently pick a wait nobody asked for. `Number.parseInt` + * alone accepts "10s" as 10 (a 10ms wait, effectively fire-and-forget) and + * yields NaN for "" or "abc", and every NaN comparison is false. + */ +export function parseWaitBudget(raw: string | undefined): number { + if (raw === undefined || raw.trim() === "") return DEFAULT_REMEMBER_WAIT_MS; + const parsed = Number(raw); + if (!Number.isFinite(parsed) || parsed < 0) { + log.warn("remember.wait_budget_invalid", { + value: raw, + usingMs: DEFAULT_REMEMBER_WAIT_MS, + }); + return DEFAULT_REMEMBER_WAIT_MS; + } + if (parsed > MAX_REMEMBER_WAIT_MS) { + log.warn("remember.wait_budget_clamped", { + value: raw, + usingMs: MAX_REMEMBER_WAIT_MS, + }); + return MAX_REMEMBER_WAIT_MS; + } + return Math.floor(parsed); +} + +export const REMEMBER_WAIT_MS = parseWaitBudget( + process.env.MEMWAL_MCP_REMEMBER_WAIT_MS +); + +/** + * Poll interval for a bounded wait. + * + * `waitForRememberJob` sleeps BEFORE its first poll and defaults to 1500ms + * with 1.5^attempt backoff, which spends a 10s budget on ~5 polls and can + * miss a write that landed at 1.2s. 400ms catches the fast case while + * staying far below the relayer's 60 weighted-requests/min delegate-key + * limit — the backoff reaches ~6 polls in 10s, not 25. + */ +export const REMEMBER_POLL_INTERVAL_MS = 400; + +/** + * `waitForRememberJob` signals outcome through `status` on a plain Error: + * 504 = still running at the deadline, 500 = the job failed, 404 = no such + * job. Only 504 is a non-error for us — the write is still in flight and the + * caller gets a job_id to resolve later. + */ +export function isStillRunning(err: unknown): boolean { + // 504: `waitForRememberJob` hit its own budget — the job is still going. + if ((err as { status?: number } | null)?.status === 504) return true; + // Our deadline firing means the RELAYER stopped answering us, not that the + // job stopped. It was accepted, it is a row in `remember_jobs`, and the + // caller needs the job_id to settle it. Treating this as a plain error + // threw that id away — the one thing the pending result exists to return. + return (err as { name?: string } | null)?.name === "MemWalRelayerUnresponsive"; +} + +/** + * Give a job error the name `wrapTool` routes on, so the agent can tell a + * failed write from a missing one without parsing the message. The SDK throws + * an unnamed Error with a status code; `wrapTool` cannot classify that. + */ +export function nameJobError(err: unknown): unknown { + if (!(err instanceof Error)) return err; + const status = (err as { status?: number }).status; + if (status === 500) err.name = "MemWalRememberJobFailed"; + else if (status === 404) err.name = "MemWalRememberJobNotFound"; + else if (status === 504) err.name = "MemWalRememberJobTimeout"; + return err; +} + +/** + * The line shown when a write is accepted but has not landed inside the wait + * budget. Worded so an agent cannot read it as success: the fact is NOT saved + * yet, and there is exactly one way to find out whether it lands. + */ +export function pendingMessage(jobId: string, waitedMs: number): string { + // The zero-budget path never waited, so saying "has not finished after + // 0.0s" would misdescribe it — and that is the default path, the one an + // agent reads on nearly every call. + const opening = + waitedMs === 0 + ? "ACCEPTED, NOT YET SAVED — the relayer has durably queued this write." + : `NOT SAVED YET — the write was accepted but has not finished after ${(waitedMs / 1000).toFixed(1)}s.`; + + return ( + `${opening}\n` + + `job_id=${jobId}\n` + + `Walrus uploads queue, so storing typically takes another 30-60s. Do NOT tell the ` + + `user the fact is stored — say it is being saved. Call memwal_remember_status with ` + + `this job_id to get the blob_id once it lands, or to learn that it failed; a job ` + + `CAN fail after acceptance, and this is the only way to find out. Do not re-send ` + + `the same fact with memwal_remember — that queues a second copy behind this one.` + ); +} + +/** Longest fact echoed back in a pending listing. The line exists so the + * agent can tell which job_id belongs to which fact, not to reproduce the + * fact — and 20 of them at full length would crowd out the instructions + * underneath. */ +const PENDING_FACT_PREVIEW_CHARS = 80; + +function previewFact(text: string): string { + const flat = text.replace(/\s+/g, " ").trim(); + return flat.length > PENDING_FACT_PREVIEW_CHARS + ? `${flat.slice(0, PENDING_FACT_PREVIEW_CHARS - 1)}…` + : flat; +} + +/** + * The bulk counterpart of `pendingMessage`. Same contract — an agent must not + * read it as success — with the one addition bulk needs: each job_id is paired + * with the fact it carries, because "one of these five failed" is only + * actionable if the agent can tell which. + */ +export function pendingBulkMessage( + entries: Array<{ jobId: string; text: string }>, + waitedMs: number, +): string { + const n = entries.length; + const opening = + waitedMs === 0 + ? `ACCEPTED, NOT YET SAVED — the relayer has durably queued ${n} write(s).` + : `NOT SAVED YET — ${n} write(s) were accepted but had not finished after ${(waitedMs / 1000).toFixed(1)}s.`; + + const lines = entries + .map((e, i) => `${i + 1}. job_id=${e.jobId} — ${previewFact(e.text)}`) + .join("\n"); + + return ( + `${opening}\n${lines}\n` + + `Walrus uploads queue and are written one at a time per wallet, so a batch takes ` + + `longer than a single fact. Do NOT tell the user these facts are stored — say they ` + + `are being saved. Call memwal_remember_status with job_ids=[...] to get the blob_ids ` + + `once they land, or to learn that one failed; a job CAN fail after acceptance, and ` + + `this is the only way to find out. Do not re-send these facts with ` + + `memwal_remember_bulk — that queues a second copy behind them.` + ); +} + +/** + * Longest we let the relayer take to ACCEPT a write before giving up on it. + * + * The SDK's `signedRequest` only aborts a request when the caller hands it a + * signal, and of the memory methods only `recall()` does (15s). `rememberAsync` + * and every job-status poll call it with no signal at all, so the underlying + * `fetch` has no deadline of its own. That makes `timeoutMs` a loop-entry + * check rather than a bound: `waitForRememberJob` tests `Date.now() < deadline` + * at the top of each iteration, so one stalled HTTP request runs as long as the + * socket stays open and a tool documented as capping at 90s is observed past + * 120s. Returning at accept does not fix that on its own — the accept POST is + * exactly one of the unbounded calls. + * + * 15s matches the only deadline the SDK sets for itself. A healthy accept is + * ~1.1s, so this fires only when something is genuinely wrong. + * + * The SDK grows its own 30s per-request backstop in the release after the + * pinned 0.1.7, which does not make this redundant: that one is a floor for + * every consumer, this is the tighter bound an interactive agent needs, and + * whichever is smaller fires first. + */ +export const DEFAULT_ACCEPT_DEADLINE_MS = 15_000; + +/** + * Read once, validated the same way as the wait budget: a typo must not + * silently pick a deadline nobody asked for. Exposed as an env knob because an + * operator on a slow link is the one person who can tell a hung relayer from a + * merely distant one. + */ +export const ACCEPT_DEADLINE_MS = (() => { + const raw = process.env.MEMWAL_MCP_ACCEPT_DEADLINE_MS; + if (raw === undefined || raw.trim() === "") return DEFAULT_ACCEPT_DEADLINE_MS; + const parsed = Number(raw); + if (!Number.isFinite(parsed) || parsed <= 0) { + log.warn("remember.accept_deadline_invalid", { + value: raw, + usingMs: DEFAULT_ACCEPT_DEADLINE_MS, + }); + return DEFAULT_ACCEPT_DEADLINE_MS; + } + return Math.floor(parsed); +})(); + +/** + * Grace added to a wait budget before we stop believing the SDK will return. + * + * The budget bounds when the SDK starts its last poll, not when that poll + * finishes, so a stalled request can overshoot by an unbounded amount. This + * caps the overshoot instead. + */ +const WAIT_OVERSHOOT_GRACE_MS = 10_000; + +class DeadlineExceededError extends Error { + constructor(message: string) { + super(message); + this.name = "MemWalRelayerUnresponsive"; + } +} + +/** + * Bound an SDK call that has no deadline of its own. + * + * The underlying request is NOT cancelled — the SDK gives us no way to pass a + * signal, so `fetch` keeps running until it settles or the socket dies. What + * this bounds is how long the agent waits on it, which is the part the user + * experiences as a hang. The orphaned request costs one socket and resolves + * into a promise nobody reads. + */ +export async function withDeadline( + work: Promise, + ms: number, + message: string, +): Promise { + let timer: NodeJS.Timeout | undefined; + try { + return await Promise.race([ + work, + new Promise((_, reject) => { + timer = setTimeout(() => reject(new DeadlineExceededError(message)), ms); + // Deliberately NOT unref'd — the same correction a4a94e41 made + // in the SDK, which this had copied. A deadline is the one + // timer somebody IS waiting on: unref'd it stops firing the + // moment nothing else holds the loop open, so the stalled + // request it exists to bound hangs forever instead. The + // `finally` below clears it, so it cannot outlive its work. + }), + ]); + } finally { + if (timer) clearTimeout(timer); + } +} + +/** Bound the accept leg of a write. + * + * `idempotent` is not cosmetic. `POST /api/remember` carries a content-derived + * idempotency_key, so a retry collapses onto the job already in flight. + * `POST /api/remember/bulk` carries none at all — the handler mints a fresh + * uuid per item and inserts with no conflict clause — so a retry there is N + * more paid Walrus blobs for the same N facts. + * + * That distinction decides what we may tell the agent, and the deadline makes + * it urgent rather than theoretical: `withDeadline` does not cancel the + * underlying request, so when it fires the relayer has usually accepted + * already. Inviting a blind retry on the bulk path is close to guaranteeing + * the duplicate. + */ +export function withAcceptDeadline( + work: Promise, + what: string, + opts: { idempotent: boolean; deadlineMs?: number }, +): Promise { + const ms = opts.deadlineMs ?? ACCEPT_DEADLINE_MS; + const shared = + `Walrus Memory did not accept the ${what} within ${ms / 1000}s — the ` + + `relayer is unreachable or not responding. The write may or may not have been queued, ` + + `so do NOT tell the user it was saved.`; + + return withDeadline( + work, + ms, + opts.idempotent + ? `${shared} Retrying is safe: this write carries a content-derived idempotency ` + + `key, so a retry attaches to the job already in flight instead of queueing a ` + + `second paid copy.` + : `${shared} Do NOT retry blindly — this endpoint carries no idempotency key, so a ` + + `re-send stores every fact a SECOND time at full cost. Check with memwal_recall ` + + `first, and only re-send what is genuinely missing.`, + ); +} + +/** + * The MCP TypeScript SDK's default `tools/call` timeout — + * `DEFAULT_REQUEST_TIMEOUT_MSEC` in @modelcontextprotocol/sdk + * (shared/protocol.js:8, applied at :712 as + * `options?.timeout ?? DEFAULT_REQUEST_TIMEOUT_MSEC`). + * + * A host that does not raise its own ceiling aborts the call at this mark. An + * abort is strictly worse than a timeout we raise ourselves: the caller gets no + * message it can act on, and on a non-idempotent endpoint its retry writes + * everything twice. So every leg this file bounds must finish inside it. + */ +export const MCP_CLIENT_DEFAULT_TIMEOUT_MS = 60_000; + +/** + * Headroom left under the client ceiling for JSON-RPC framing, transport and + * the relayer's own response write. A leg budgeted at exactly the ceiling loses + * the race it was meant to win. + */ +const CLIENT_TIMEOUT_HEADROOM_MS = 15_000; + +/** + * Deadline for `memwal_analyze`'s extraction leg. + * + * `/api/analyze` runs the extractor LLM inline before it answers, so this leg + * is real work, not an accept — 15s would cut off healthy extraction on any + * transcript worth extracting from. But it was budgeted at the full 60_000, + * exactly the client ceiling, so a slow extraction raced the host's abort and + * usually lost: the agent saw a dead call instead of a message, and `analyze` + * is the one endpoint with no idempotency key, so retrying it stores every + * extracted fact a second time. + * + * Derived from the ceiling rather than written as a literal so the two cannot + * drift apart. + */ +export const ANALYZE_EXTRACTION_DEADLINE_MS = + MCP_CLIENT_DEFAULT_TIMEOUT_MS - CLIENT_TIMEOUT_HEADROOM_MS; + +/** Bound a status wait at its own budget plus the overshoot grace. */ +export function withWaitDeadline(work: Promise, budgetMs: number): Promise { + return withDeadline( + work, + budgetMs + WAIT_OVERSHOOT_GRACE_MS, + `Walrus Memory stopped responding while waiting for the write to land. The job is ` + + `still queued relayer-side — do NOT tell the user it was saved, and do NOT re-send ` + + `the fact. Call memwal_remember_status with the job_id to settle it.`, + ); +} + +/** + * Longest we will sit on a relayer-advised cooldown before handing the problem + * back to the agent. + * + * The relayer answers a spent rate-limit budget with `retry_after_seconds: 60`. + * Sleeping that out inside a tool call is not a fix — it is the 60s hang this + * whole change set exists to remove, and the MCP client would time out first. + * So a short cooldown is absorbed and a long one is reported, with the wait + * named so the agent can come back rather than guess. + */ +export const MAX_ABSORBED_COOLDOWN_MS = 8_000; + +/** Attempts, including the first. Two retries is enough for a transient blip; + * more just delays an answer the agent could act on. */ +const RELAYER_RETRY_ATTEMPTS = 3; + +/** + * Errors where the request provably did NOT reach the handler, so re-sending + * cannot duplicate work. + * + * This matters most for `/api/remember/bulk`, which carries no idempotency key + * — a blind retry there would store every fact twice. Both cases below are + * rejections BEFORE any job row exists: 429 comes from the rate limiter, and + * AUTH_UPSTREAM_UNAVAILABLE from the delegate-key lookup failing open. Any + * other 5xx could have been thrown after a write started, so it is not retried. + */ +function isSafelyRetryable(err: unknown): boolean { + const e = err as { status?: number; serverCode?: string } | null; + if (e?.status === 429) return true; + return e?.status === 503 && e?.serverCode === "AUTH_UPSTREAM_UNAVAILABLE"; +} + +function advisedCooldownMs(err: unknown): number { + const secs = (err as { retryAfterSeconds?: number } | null)?.retryAfterSeconds; + return typeof secs === "number" && secs > 0 ? secs * 1000 : 1_000; +} + +/** + * What the relayer actually said it denied on. + * + * `rate_limit_response` (services/server/src/rate_limit.rs) answers a 429 with + * `{error, layer, limit, retry_after_seconds}`, and the SDK embeds that body + * verbatim in the error message. Read it rather than asserting a scope: there + * are three layers with different windows — `delegate_key` and `account_burst` + * per minute, `account_sustained` per HOUR — and naming the wrong one sends + * the caller to wait out a window that was never the one it hit. + */ +function rateLimitFacts(err: unknown): { layer?: string; limit?: string } { + const text = String((err as { message?: string } | null)?.message ?? ""); + const open = text.indexOf("{"); + const close = text.lastIndexOf("}"); + if (open < 0 || close <= open) return {}; + try { + const body = JSON.parse(text.slice(open, close + 1)) as Record; + return { + layer: typeof body.layer === "string" ? body.layer : undefined, + limit: typeof body.limit === "string" ? body.limit : undefined, + }; + } catch { + // A body we cannot parse is not a reason to lose the 429 itself. + return {}; + } +} + +/** Whose budget the named layer belongs to. */ +function scopeOfLayer(layer: string | undefined): string { + if (layer === "delegate_key") return "per delegate key"; + if (layer === "account_burst" || layer === "account_sustained") return "per account"; + return "on this account"; +} + +/** + * The retry advice. + * + * `retry_after_seconds` is a refill hint from a token bucket, NOT a reset. On + * the hourly `account_sustained` layer the relayer still advises ~300s, which + * buys back only a slice of a 1000/hour budget — waiting it out and retrying + * re-trips the limit whenever the budget is genuinely spent. Saying "resets in + * Ns; retry after that" turned that into a retry loop that never converges, so + * the message now says what the number is worth. + */ +function retryAdvice(secs: number, layer: string | undefined): string { + return layer === "account_sustained" + ? `The relayer advises retrying in ~${secs}s, but that is a partial refill of an ` + + `hourly budget, not a reset — if the budget is spent, a retry then fails again. ` + + `Report the failure rather than waiting and retrying in a loop.` + : `Retry after ~${secs}s.`; +} + +/** True if the relayer refused this call with a rate limit. */ +export function isRateLimited(err: unknown): boolean { + return (err as { status?: number } | null)?.status === 429; +} + +/** + * The rate-limit error an agent can act on, built from what the relayer said. + * + * Shared with `memwal_remember_status` so a throttled poll cannot be reported + * as "still uploading": a denied poll tells us nothing about the job, and + * rendering it as progress is what turns a rate limit into an agent that keeps + * polling and spends more of the budget it has already exhausted. + */ +export function rateLimitedError(err: unknown, what: string): Error { + const secs = Math.ceil(advisedCooldownMs(err) / 1000); + const { layer, limit } = rateLimitFacts(err); + const hit = layer + ? `Limit hit: ${layer} ${scopeOfLayer(layer)}` + (limit ? ` (${limit})` : "") + ". " + : ""; + const e = new Error( + `Walrus Memory rate limit reached while trying to ${what}. ${hit}` + + `${retryAdvice(secs, layer)} To spend less of the budget, save several facts ` + + `with one memwal_remember_bulk call instead of repeated memwal_remember calls, ` + + `and settle a batch with a single memwal_remember_status(job_ids=[...]).`, + ); + e.name = "MemWalRelayerUnavailable"; + (e as Error & { status?: number }).status = 429; + return e; +} + +/** Honour the relayer's `retry_after` instead of surfacing a raw 429. + * + * Once the per-delegate-key budget is spent the write is simply never made, + * and nothing retried — the quietest failure here. A short cooldown is + * absorbed; a long one is reported, because sleeping it out inside a tool call + * is the hang this file exists to remove. */ +export async function withRelayerRetry(work: () => Promise, what: string): Promise { + let last: unknown; + for (let attempt = 1; attempt <= RELAYER_RETRY_ATTEMPTS; attempt++) { + try { + return await work(); + } catch (err) { + last = err; + if (!isSafelyRetryable(err)) throw err; + + const cooldown = advisedCooldownMs(err); + if (attempt === RELAYER_RETRY_ATTEMPTS || cooldown > MAX_ABSORBED_COOLDOWN_MS) { + const secs = Math.ceil(cooldown / 1000); + const limited = (err as { status?: number }).status === 429; + const { layer, limit } = limited ? rateLimitFacts(err) : {}; + const hit = layer + ? `Limit hit: ${layer} ${scopeOfLayer(layer)}` + + (limit ? ` (${limit})` : "") + ". " + : ""; + const e = new Error( + limited + ? `Walrus Memory rate limit reached while trying to ${what}. THE FACT WAS NOT ` + + `SAVED — tell the user it could not be stored rather than that it is being ` + + `saved. ${hit}${retryAdvice(secs, layer)} To spend less of the budget, save ` + + `several facts with one memwal_remember_bulk call instead of repeated ` + + `memwal_remember calls, and settle a batch with a single ` + + `memwal_remember_status(job_ids=[...]).` + : `Walrus Memory could not ${what}: the relayer's credential check is ` + + `temporarily unavailable. THE FACT WAS NOT SAVED. Retry in about ${secs}s.`, + ); + e.name = "MemWalRelayerUnavailable"; + (e as Error & { status?: number }).status = (err as { status?: number }).status; + throw e; + } + + log.warn("remember.relayer_retry", { + what, + attempt, + status: (err as { status?: number }).status, + cooldownMs: cooldown, + }); + await new Promise((r) => setTimeout(r, cooldown)); + } + } + throw last; +} diff --git a/services/server/scripts/mcp/tools/remember.ts b/services/server/scripts/mcp/tools/remember.ts index 2e9b523e9..91979ba68 100644 --- a/services/server/scripts/mcp/tools/remember.ts +++ b/services/server/scripts/mcp/tools/remember.ts @@ -3,13 +3,26 @@ import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; import type { MemWalSession } from "../auth.js"; import { TOOL_METADATA } from "./annotations.js"; import { wrapTool, walruscanBlobUrl } from "./util.js"; +import { SECRET_EXCLUSION_RULES, AUTO_SAVE_OPT_IN_RULE } from "./memory-policy.js"; +import { sanitizeFact, redactionNotice, refusalNotice } from "./redaction.js"; +import { + REMEMBER_WAIT_MS, + REMEMBER_POLL_INTERVAL_MS, + isStillRunning, + nameJobError, + pendingMessage, + withAcceptDeadline, + withWaitDeadline, + withRelayerRetry, + derivedIdempotencyKey, +} from "./remember-wait.js"; const REMEMBER_INPUT = { text: z .string() .min(1) .describe( - "The full, detailed fact to save. Pass the COMPLETE statement — do not summarize." + "The full, detailed fact to save. Pass the COMPLETE statement — do not summarize. Leave credentials out: passwords, API keys, tokens, private keys, seed phrases, auth headers and URLs with an embedded user:password are stripped before the write and never stored." ), namespace: z .string() @@ -20,8 +33,13 @@ const REMEMBER_INPUT = { } as const; /** - * memwal_remember — persist a durable fact to MemWal and return only when the - * blob is written end-to-end (embed → SEAL encrypt → Walrus upload → on-chain). + * memwal_remember — persist a durable fact to MemWal. + * + * Returns as soon as the blob is written end-to-end (embed → SEAL encrypt → + * Walrus upload → on-chain) when that happens inside `REMEMBER_WAIT_MS`. + * Otherwise it returns the job_id and says plainly that the fact is not saved + * yet — see `remember-wait.ts` for why the old always-block behaviour cost + * 30–75s per call. * * Call this PROACTIVELY whenever the user reveals a durable fact about * themselves or the project (preference, decision, constraint, correction, @@ -38,23 +56,88 @@ export function registerRememberTool( { ...TOOL_METADATA.memwal_remember, description: - "Save a durable fact about the user or project to their Walrus Memory. Call this PROACTIVELY whenever the user states a preference, decision, constraint, correction, identity detail, or recurring workflow — even if they did not say 'remember this'. Skip one-off tasks, the current file or bug, and small talk. Pass the full statement; do not summarize. To save several facts at once, use memwal_remember_bulk instead.", + "Save a durable fact about the user or project to their Walrus Memory. Call this whenever the user states a preference, decision, constraint, correction, identity detail, or recurring workflow — PROACTIVELY, without being asked, when they have turned automatic memory on. Skip one-off tasks, the current file or bug, and small talk. Pass the full statement; do not summarize. To save several facts at once, use memwal_remember_bulk instead. By default this returns in ~1s once the relayer has accepted the job (job_id) — the Walrus write is still in flight and the fact is NOT stored yet. Do not claim it is saved. Resolve with memwal_remember_status. A blob_id in the same reply means it did land inside an optional wait budget (MEMWAL_MCP_REMEMBER_WAIT_MS). " + + AUTO_SAVE_OPT_IN_RULE + + " " + + SECRET_EXCLUSION_RULES + + " Walrus storage is append-only: a stored secret cannot be deleted, so this tool strips credential shapes from the text before writing and tells you what it removed.", inputSchema: REMEMBER_INPUT, }, wrapTool<{ text: string; namespace?: string }>(session, "memwal_remember", async ({ text, namespace }) => { - const result = await session.memwal.rememberAndWait( - text, - namespace, - { timeoutMs: 90_000 } + // Runs BEFORE anything reaches the SDK. Walrus is append-only, so a + // credential that gets written cannot be taken back (WALM-642). + const safe = sanitizeFact(text); + if (safe.refusal) { + return { + content: [ + { type: "text" as const, text: refusalNotice(safe.refusal) }, + ], + }; + } + const notice = redactionNotice(safe.kinds, safe.count); + const safeText = safe.text; + + // Two steps rather than `rememberAndWait`, because the accept and + // the wait need separate budgets: acceptance is the part that + // must succeed, the wait is a courtesy we cut short. + const accepted = await withAcceptDeadline( + withRelayerRetry( + () => + session.memwal.rememberAsync(safeText, namespace, { + // Ours, not the SDK's random one — see + // derivedIdempotencyKey. This is what makes the + // accept-timeout message's retry promise true. + // Keyed on the REDACTED text, so a retry of the + // same fact derives the same key. + idempotencyKey: derivedIdempotencyKey(namespace, safeText), + }), + "save this fact", + ), + "memwal_remember write", + { idempotent: true }, ); - return { + + const withNotice = (body: string) => + notice ? `${body}\n\n${notice}` : body; + + const pending = (waitedMs: number) => ({ content: [ { - type: "text", - text: `Saved to Walrus Memory. blob_id=${result.blob_id} namespace=${result.namespace}\nExplorer: ${walruscanBlobUrl(result.blob_id)}`, + type: "text" as const, + text: withNotice(pendingMessage(accepted.job_id, waitedMs)), }, ], - }; + }); + + // A zero budget is the documented fire-and-accept mode. Skip the + // wait entirely instead of entering a loop that cannot poll. + if (REMEMBER_WAIT_MS === 0) return pending(0); + + const startedAt = Date.now(); + try { + const result = await withWaitDeadline( + session.memwal.waitForRememberJob(accepted.job_id, { + timeoutMs: REMEMBER_WAIT_MS, + pollIntervalMs: REMEMBER_POLL_INTERVAL_MS, + }), + REMEMBER_WAIT_MS, + ); + return { + content: [ + { + type: "text" as const, + text: withNotice( + `Saved to Walrus Memory. blob_id=${result.blob_id} namespace=${result.namespace}\nExplorer: ${walruscanBlobUrl(result.blob_id)}`, + ), + }, + ], + }; + } catch (err) { + // Still running at the deadline is the expected path, not a + // failure — the job is durably accepted and keeps going. + if (isStillRunning(err)) return pending(Date.now() - startedAt); + throw nameJobError(err); + } }) ); } diff --git a/services/server/scripts/mcp/tools/util.ts b/services/server/scripts/mcp/tools/util.ts index 7068c7fa2..b82589f11 100644 --- a/services/server/scripts/mcp/tools/util.ts +++ b/services/server/scripts/mcp/tools/util.ts @@ -3,6 +3,7 @@ */ import type { MemWalSession } from "../auth.js"; import { createLogger } from "../logger.js"; +import { classifyToolError, describeFailure, probeRelayerHealth } from "./failure.js"; const log = createLogger("mcp"); @@ -48,6 +49,10 @@ export function explorerFooter(): string { const DEFAULT_SLOW_TOOL_WARN_MS = 5000; +/** How long a failed call waits on the relayer's `/health` before saying it + * got no answer. The call has already failed; this only buys its diagnosis. */ +const HEALTH_PROBE_TIMEOUT_MS = 2000; + /** * Above this, a tool call is reported at `warn` rather than `info`. * Tuned to sit above a healthy `memwal_health` (single unsigned GET to the @@ -165,13 +170,32 @@ export function wrapTool( // Name the failure in the structured line too. Without this the log // says a call failed and the operator still has to go find the // separate console.error below to learn how. + // Prefer an explicitly set `name` over the constructor's. The SDK + // signals a job outcome with a status code on a plain Error, whose + // constructor is always `Error` — routing on that alone left the + // switch below unreachable for exactly the cases it names. + const name = err?.name && err.name !== "Error" + ? err.name + : err?.constructor?.name ?? "Error"; + // Timed before the probe below, which is diagnosis, not the call. + const fields = outcomeFields(); + // A timeout or a failed connect says nothing about the relayer + // on its own; ask its `/health` so the agent learns whether it is + // down, unhealthy, or up with this one call stuck. + const failure = classifyToolError(err); + const probe = + failure.kind === "timeout" || failure.kind === "unreachable" + ? await probeRelayerHealth(session.relayerUrl, HEALTH_PROBE_TIMEOUT_MS) + : null; log.warn("tool.failed", { - ...outcomeFields(), - errName: err?.constructor?.name ?? "Error", + ...fields, + errName: name, errMessage: err?.message ?? String(err), causeCode: err?.cause?.code ?? null, + failureKind: failure.kind, + stage: failure.kind === "recall_timeout" ? failure.stage : null, + health: probe?.kind ?? null, }); - const name = err?.constructor?.name ?? "Error"; const msg = err?.message ?? String(err); const cause = err?.cause; const causeStr = cause @@ -186,6 +210,13 @@ export function wrapTool( : "") ); + if (failure.kind !== "other") { + return { + content: [{ type: "text", text: describeFailure(tool, failure, probe) }], + isError: true, + }; + } + let prefix = "Tool error"; switch (name) { case "MemWalRememberJobFailed": diff --git a/services/server/scripts/package-lock.json b/services/server/scripts/package-lock.json index f68a66a28..306bde7f4 100644 --- a/services/server/scripts/package-lock.json +++ b/services/server/scripts/package-lock.json @@ -9,7 +9,7 @@ "version": "0.1.0", "dependencies": { "@modelcontextprotocol/sdk": "1.29.0", - "@mysten-incubation/memwal": "0.0.3", + "@mysten-incubation/memwal": "0.1.7", "@mysten/seal": "1.1.0", "@mysten/sui": "2.17.0", "@mysten/walrus": "1.1.7", @@ -621,20 +621,23 @@ } }, "node_modules/@mysten-incubation/memwal": { - "version": "0.0.3", - "resolved": "https://registry.npmjs.org/@mysten-incubation/memwal/-/memwal-0.0.3.tgz", - "integrity": "sha512-TafWL5MPEOPXmQstap5YP/5j6tm6x+ZJXTnEG6KKoIfyE7krKZ8lIZCsa6EcS4hrzY1rDtr+El5RxjmKlGhCig==", + "version": "0.1.7", + "resolved": "https://registry.npmjs.org/@mysten-incubation/memwal/-/memwal-0.1.7.tgz", + "integrity": "sha512-bVkk0+Zsjl6A4g6j7gvVhyFjmtHwhXLimgm8CoQGhRoM1Pz1XrR8/R5Uio+vah1JB3s6ICTtt1ycCcOAamyTfg==", "license": "Apache-2.0", "dependencies": { "@noble/ed25519": "^2.3.0", "@noble/hashes": "^2.0.0" }, + "engines": { + "node": ">=20.0.0" + }, "peerDependencies": { "@mysten/seal": ">=1.1.0", "@mysten/sui": ">=2.5.0", "@mysten/walrus": ">=1.0.3", "ai": ">=4.0.0", - "zod": "^3.23.0" + "zod": "^3.23.0 || ^4.0.0" }, "peerDependenciesMeta": { "@mysten/walrus": { diff --git a/services/server/scripts/package.json b/services/server/scripts/package.json index 8cb3f4980..78384fafe 100644 --- a/services/server/scripts/package.json +++ b/services/server/scripts/package.json @@ -5,12 +5,12 @@ "type": "module", "scripts": { "sidecar": "tsx sidecar-server.ts", - "typecheck": "tsc --noEmit", + "typecheck": "tsc --noEmit", "test": "node --test --import tsx './mcp/__tests__/*.test.ts' './__tests__/*.test.ts'" }, "dependencies": { "@modelcontextprotocol/sdk": "1.29.0", - "@mysten-incubation/memwal": "0.0.3", + "@mysten-incubation/memwal": "0.1.7", "@mysten/seal": "1.1.0", "@mysten/sui": "2.17.0", "@mysten/walrus": "1.1.7", diff --git a/services/server/scripts/sidecar-server.ts b/services/server/scripts/sidecar-server.ts index a76986b6e..8ec5acef8 100644 --- a/services/server/scripts/sidecar-server.ts +++ b/services/server/scripts/sidecar-server.ts @@ -21,6 +21,7 @@ * GET /health → local liveness (no auth) * GET /ready → Sui/Walrus execution identity + limits (no auth) * GET /metrics/wallet → aggregate wallet-execution metrics (no auth) + * GET /metrics/uploads → upload-queue counters + limits, no I/O (no auth) * GET /internal/wallet-balances → per-wallet balances (sidecar auth) * /mcp/* → MCP session routes (own auth; see mcp/) * POST /seal/encrypt → { data, owner, packageId, accountId } → { encryptedData } diff --git a/services/server/scripts/sidecar/app.ts b/services/server/scripts/sidecar/app.ts index b9b9eadb2..dc35a245e 100644 --- a/services/server/scripts/sidecar/app.ts +++ b/services/server/scripts/sidecar/app.ts @@ -3,7 +3,7 @@ * * Registration order is load-bearing: * 1. request-id + CORS-strip middleware run for every request. - * 2. /health, /ready, /metrics/wallet, and full-mode MCP routes are mounted BEFORE + * 2. /health, /ready, /metrics/*, and full-mode MCP routes are mounted BEFORE * the shared-secret middleware — they must stay reachable without the * sidecar token (probes, scrapers, and MCP traffic that carries the * end-user's own Bearer token instead). @@ -22,6 +22,7 @@ import { import { registerHealthRoute, registerInternalWalletBalancesRoute, + registerUploadMetricsRoute, registerWalletMetricsRoute, } from "./routes/health.js"; import { registerSealRoutes } from "./routes/seal.js"; @@ -72,6 +73,7 @@ export function createSidecarApp(mode: "full" | "writer" = SIDECAR_ROUTE_MODE): // Wallet-execution metrics — placed before auth so operators / scrapers // don't need a token. registerWalletMetricsRoute(app); + registerUploadMetricsRoute(app); app.use(sharedSecretAuthMiddleware); diff --git a/services/server/scripts/sidecar/clients.ts b/services/server/scripts/sidecar/clients.ts index 40aada973..79c2ae861 100644 --- a/services/server/scripts/sidecar/clients.ts +++ b/services/server/scripts/sidecar/clients.ts @@ -23,6 +23,7 @@ import { UPLOAD_RELAY_TIP_TIMEOUT_MS, WALRUS_CLIENT_MAX_AGE_MS, WALRUS_DIRECT_UPLOAD, + WALRUS_STORAGE_NODE_TIMEOUT_MS, WALRUS_PACKAGE_ID, WALRUS_STAKING_POOL_ID, WALRUS_SYSTEM_OBJECT_ID, @@ -70,6 +71,9 @@ function createWalrusClient(): WalrusClient { !WALRUS_DIRECT_UPLOAD && !!WALRUS_UPLOAD_RELAY_URL && WALRUS_UPLOAD_RELAY_URL !== "none"; const baseConfig = { suiClient: suiClient as any, + // Applies to both branches: the relay still reads slivers back from + // storage nodes, and the direct branch has nothing else to set it. + storageNodeClientOptions: { timeout: WALRUS_STORAGE_NODE_TIMEOUT_MS }, ...(useRelay ? { uploadRelay: { diff --git a/services/server/scripts/sidecar/concurrency.ts b/services/server/scripts/sidecar/concurrency.ts index a2a2b7289..9355a70b2 100644 --- a/services/server/scripts/sidecar/concurrency.ts +++ b/services/server/scripts/sidecar/concurrency.ts @@ -126,7 +126,8 @@ export async function acquireWalrusUploadSlots( queuedWalrusUploads = Math.max(0, queuedWalrusUploads - 1); activeWalrusUploads += 1; - const waitMs = Date.now() - startedAt; + const acquiredAt = Date.now(); + const waitMs = acquiredAt - startedAt; if (waitMs >= 1_000) { console.warn(`[walrus/upload] [${traceId}] limiter_acquired ${JSON.stringify({ jobId, @@ -136,15 +137,37 @@ export async function acquireWalrusUploadSlots( })}`); } + // The underlying AsyncSemaphore release is not one-shot: a second call + // can free capacity a successor still holds. Keep the returned callback + // idempotent so route `finally` + error-path cleanup cannot over-release. + let released = false; return () => { + if (released) return; + released = true; activeWalrusUploads = Math.max(0, activeWalrusUploads - 1); releaseGlobal?.(); releaseWallet?.(); + console.log(`[walrus/upload] [${traceId}] limiter_released ${JSON.stringify({ + jobId, + keyIndex, + waitMs, + heldMs: Date.now() - acquiredAt, + counts: getUploadCounts(), + limits: walrusUploadLimitSnapshot(keyIndex), + })}`); }; } catch (err) { queuedWalrusUploads = Math.max(0, queuedWalrusUploads - 1); releaseGlobal?.(); releaseWallet?.(); + console.warn(`[walrus/upload] [${traceId}] limiter_acquire_failed ${JSON.stringify({ + jobId, + keyIndex, + waitMs: Date.now() - startedAt, + error: err instanceof Error ? err.message : String(err), + counts: getUploadCounts(), + limits: walrusUploadLimitSnapshot(keyIndex), + })}`); throw err; } } diff --git a/services/server/scripts/sidecar/config.ts b/services/server/scripts/sidecar/config.ts index a755c7097..6e7bfea9e 100644 --- a/services/server/scripts/sidecar/config.ts +++ b/services/server/scripts/sidecar/config.ts @@ -234,6 +234,27 @@ export const WALRUS_UPLOAD_ACQUIRE_TIMEOUT_MS = parsePositiveIntEnv( 1_000, 180_000 ); +/** Per-request timeout for talking to a Walrus storage node. + * + * Applies to both the direct and relay paths (`storageNodeClientOptions`). + * `uploadRelay.timeout` (default 120s, `WALRUS_RELAY_TIMEOUT_MS`) is a + * different knob: one HTTP request to the upload relay, not a storage + * node. Before this env existed, neither path set + * `storageNodeClientOptions`, so both used StorageNodeClient's 30s + * default. The default here stays 30s so no environment shifts silently; + * what changes is that both paths now have a knob. + * + * Per NODE request, not per write: raising it also lengthens how long a + * single unresponsive node stalls a shard. Keep it under + * WALRUS_UPLOAD_ACQUIRE_TIMEOUT_MS. + */ +export const WALRUS_STORAGE_NODE_TIMEOUT_MS = parsePositiveIntEnv( + "WALRUS_STORAGE_NODE_TIMEOUT_MS", + 30_000, + 1_000, + 180_000 +); + export const WALRUS_UPLOAD_EFFECTS_RETRY_DELAYS_MS = [2_000, 5_000, 10_000, 20_000, 40_000] as const; export const DURABLE_UPLOAD_PROTOCOL_VERSION = 3; diff --git a/services/server/scripts/sidecar/routes/health.ts b/services/server/scripts/sidecar/routes/health.ts index 9b5edc9e6..82a118ca9 100644 --- a/services/server/scripts/sidecar/routes/health.ts +++ b/services/server/scripts/sidecar/routes/health.ts @@ -2,8 +2,9 @@ * Unauthenticated observability endpoints. * * All are registered BEFORE the shared-secret middleware (see app.ts): - * /health is local liveness, /ready validates upload execution identity, and - * /metrics/wallet exposes aggregate metrics to unauthenticated scrapers. + * /health is local liveness, /ready validates upload execution identity, + * /metrics/uploads serves the upload-queue counters, and /metrics/wallet + * exposes aggregate metrics to unauthenticated scrapers. * Per-wallet addresses and balances are served separately behind sidecar auth. */ @@ -138,6 +139,22 @@ export function registerHealthRoute(app: Express, requireProvenance = true): voi }); } +// In-memory upload limiter counters for the relayer's saturation probe; no I/O. +export function registerUploadMetricsRoute(app: Express): void { + app.get("/metrics/uploads", (_req: Request, res: ExpressResponse) => { + const uploads = getUploadCounts(); + res.json({ + activeWalrusUploads: uploads.active, + queuedWalrusUploads: uploads.queued, + walrusUploadLimits: { + globalCapacity: WALRUS_UPLOAD_MAX_CONCURRENCY, + perWalletCapacity: WALRUS_UPLOAD_PER_WALLET_CONCURRENCY, + acquireTimeoutMs: WALRUS_UPLOAD_ACQUIRE_TIMEOUT_MS, + }, + }); + }); +} + // Wallet-execution metrics (observability). // // `walletObjectLockEquivocationTotal` is the canary for concurrent uploads diff --git a/services/server/scripts/sidecar/routes/walrus-upload-journal.ts b/services/server/scripts/sidecar/routes/walrus-upload-journal.ts index b2a7aa4ff..cdef6a03d 100644 --- a/services/server/scripts/sidecar/routes/walrus-upload-journal.ts +++ b/services/server/scripts/sidecar/routes/walrus-upload-journal.ts @@ -4,6 +4,7 @@ * advance and return one checkpointable WriteBlobStep. */ +import { createHash } from "node:crypto"; import express, { type Express } from "express"; import type { WriteBlobStep, @@ -26,6 +27,7 @@ import { JSON_LIMIT_WALRUS_UPLOAD, MAX_WALRUS_EPOCHS, SERVER_SUI_PRIVATE_KEYS, + SUI_CHAIN_IDENTIFIER, SUI_NETWORK, SUI_TYPE, WALRUS_PACKAGE_ID, @@ -256,6 +258,51 @@ function parsePreparedRegisterTransaction(raw: unknown): PreparedRegisterTransac }; } +/** + * The expiration a direct-signed register needs: one epoch of address-balance + * withdrawal. + * + * Paying gas from the address balance puts a `FundsWithdrawal` input in the + * transaction, and Sui admits that withdrawal only inside a `ValidDuring` + * window — the same rule `assertAddressBalanceRegisterTransaction` re-checks + * one line later. Neither `flow.register()` nor `Transaction.build()` sets one, + * so every direct-signed register failed its own assertion with + * "registerTransaction must use a ValidDuring address-balance expiration" and + * no blob was ever certified on this path. + * + * One epoch wide rather than a range: `validatePreparedRegisterTransaction` + * hands `maxEpoch` back as the journal's expiry guard, so a window outliving + * the reservation would keep replaying an entry Sui has already retired. + * + * The nonce is derived from the transaction kind rather than drawn at random, + * so re-preparing the same register — same blob, same epochs, same attributes — + * rebuilds byte-identical and the journal stays idempotent, while two different + * registers still reserve under different nonces. + */ +export function addressBalanceExpiration(epoch: bigint, transactionKind: Uint8Array) { + const bounded = String(epoch); + return { + ValidDuring: { + minEpoch: bounded, + maxEpoch: bounded, + minTimestamp: null, + maxTimestamp: null, + chain: SUI_CHAIN_IDENTIFIER, + nonce: createHash("sha256").update(transactionKind).digest().readUInt32BE(0), + }, + } as const; +} + +async function bindAddressBalanceExpiration(transaction: Transaction): Promise { + const transactionKind = await transaction.build({ + client: suiClient as any, + onlyTransactionKind: true, + }); + transaction.setExpiration( + addressBalanceExpiration(await currentSuiEpoch(), transactionKind), + ); +} + export async function prepareRegisterTransaction( transaction: Transaction, signer: Ed25519Keypair, @@ -301,6 +348,7 @@ export async function prepareRegisterTransaction( // Fail-closed sponsorship already returned above. Remaining path is the // explicit phase-1 / unconfigured-Enoki direct sign. transaction.setGasPayment([]); + await bindAddressBalanceExpiration(transaction); const bytes = await transaction.build({ client: suiClient }); assertAddressBalanceRegisterTransaction(TransactionDataBuilder.fromBytes(bytes)); const signed = await signer.signTransaction(bytes); @@ -368,11 +416,41 @@ export function assertSponsoredRegisterTransaction( assertRegisterTransactionUsesAddressBalanceWal(transactionData); } +/** Render a bound the way the guard tests it, so `null` and `undefined` — which + * the guard treats differently but a template string renders identically — stay + * distinguishable in a log line. */ +function describeExpirationBound(value: unknown): string { + if (value === null) return "null"; + if (value === undefined) return "undefined"; + return JSON.stringify(value); +} + +/** Describe an expiration precisely enough to act on it from a production log. + * + * The guard below rejects on three separate conditions joined by `||`, so the + * bare sentence it used to throw could not say which one fired. A run of these + * failures on testnet (every upload job, both deployments, 2026-09-17) could not + * be diagnosed from the relayer logs at all: the classifier reported the string, + * and the string named the invariant rather than the value that broke it. */ +function describeExpiration(expiration: TransactionDataBuilder["expiration"]): string { + if (!expiration) return "expiration=none"; + if (expiration.$kind !== "ValidDuring") return `expiration=${expiration.$kind}`; + const { minEpoch, maxEpoch, minTimestamp, maxTimestamp } = expiration.ValidDuring; + return "expiration=ValidDuring" + + ` minEpoch=${describeExpirationBound(minEpoch)}` + + ` maxEpoch=${describeExpirationBound(maxEpoch)}` + + ` minTimestamp=${describeExpirationBound(minTimestamp)}` + + ` maxTimestamp=${describeExpirationBound(maxTimestamp)}`; +} + export function assertAddressBalanceRegisterTransaction( transactionData: TransactionDataBuilder, ): bigint { if (transactionData.gasData.payment?.length !== 0) { - throw new Error("registerTransaction must pay gas from the address balance"); + throw new Error( + "registerTransaction must pay gas from the address balance" + + ` (gasData.payment.length=${String(transactionData.gasData.payment?.length ?? "undefined")})`, + ); } const expiration = transactionData.expiration; @@ -381,7 +459,10 @@ export function assertAddressBalanceRegisterTransaction( || expiration.ValidDuring.minTimestamp !== null || expiration.ValidDuring.maxTimestamp !== null ) { - throw new Error("registerTransaction must use a ValidDuring address-balance expiration"); + throw new Error( + "registerTransaction must use a ValidDuring address-balance expiration" + + ` (${describeExpiration(expiration)})`, + ); } assertRegisterTransactionUsesAddressBalanceWal(transactionData); diff --git a/services/server/src/engine/mod.rs b/services/server/src/engine/mod.rs index 13356696d..355b56acc 100644 --- a/services/server/src/engine/mod.rs +++ b/services/server/src/engine/mod.rs @@ -40,6 +40,7 @@ //! middleware, unchanged. pub mod plaintext; +pub mod stage; pub mod walrus_seal; use async_trait::async_trait; diff --git a/services/server/src/engine/stage.rs b/services/server/src/engine/stage.rs new file mode 100644 index 000000000..b73275092 --- /dev/null +++ b/services/server/src/engine/stage.rs @@ -0,0 +1,292 @@ +//! Which step of a recall is running, so one that runs out of time can say +//! where it was stuck. +//! +//! The handler owns a [`StageMarker`] and runs the recall inside +//! [`run_with_deadline`], which makes that marker the task's current one. +//! Anything the recall awaits on the same task — the engine's `fetch_batch` +//! included — calls [`enter`] as it moves on, so the trait needs no extra +//! parameter. Outside a recall (`analyze` shares `fetch_batch`) `enter` does +//! nothing. + +use std::future::Future; +use std::sync::atomic::{AtomicU8, Ordering}; +use std::sync::Arc; +use std::time::Duration; + +use tokio::time::Instant; + +/// Left for the 504 to reach the caller before its own deadline fires. +const DEADLINE_MARGIN: Duration = Duration::from_millis(1_000); +/// Least a short caller deadline is given: less cannot fit an embed, so +/// cutting there would only turn recalls that might have finished into errors. +const MIN_BUDGET: Duration = Duration::from_millis(2_000); +/// Upper bound on a caller-supplied deadline, so no number can overflow an +/// `Instant`. Ten minutes is far past anything a recall caller waits. +const MAX_DEADLINE_MS: u64 = 600_000; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[repr(u8)] +pub enum RecallStage { + Embed = 0, + VectorSearch = 1, + WalrusDownload = 2, + SealDecrypt = 3, + /// Before the recall started: credential check and rate limiting. Only + /// reported when that alone used up the caller's deadline. + Auth = 4, +} + +impl RecallStage { + /// The id sent to callers in a `RECALL_TIMEOUT` body. + pub fn as_str(self) -> &'static str { + match self { + RecallStage::Embed => "embed", + RecallStage::VectorSearch => "vector_search", + RecallStage::WalrusDownload => "walrus_download", + RecallStage::SealDecrypt => "seal_decrypt", + RecallStage::Auth => "auth", + } + } + + fn from_u8(value: u8) -> Self { + match value { + 0 => RecallStage::Embed, + 1 => RecallStage::VectorSearch, + 2 => RecallStage::WalrusDownload, + 3 => RecallStage::SealDecrypt, + _ => RecallStage::Auth, + } + } +} + +/// The stage one recall is in. Shared, because the handler has to read it +/// after the timed-out work has been dropped. +#[derive(Clone, Default)] +pub struct StageMarker(Arc); + +impl StageMarker { + pub fn set(&self, stage: RecallStage) { + self.0.store(stage as u8, Ordering::Relaxed); + } + + pub fn get(&self) -> RecallStage { + RecallStage::from_u8(self.0.load(Ordering::Relaxed)) + } +} + +tokio::task_local! { + static CURRENT: StageMarker; +} + +/// Mark `stage` as running for the recall this task is serving. +pub fn enter(stage: RecallStage) { + let _ = CURRENT.try_with(|marker| marker.set(stage)); +} + +/// What a caller's deadline leaves for the recall. +#[derive(Debug, PartialEq, Eq)] +pub enum Budget { + /// No deadline sent: run to completion. + Unbounded, + /// Stop after this long. + Run(Duration), + /// The deadline went on auth and rate limiting. Any work now would + /// answer a caller that has already given up. + Exhausted, +} + +/// What a caller that waits `deadline_ms` leaves for the recall, `already` +/// having passed since the request arrived. +pub fn budget_for(deadline_ms: Option, already: Duration) -> Budget { + let Some(deadline_ms) = deadline_ms else { + return Budget::Unbounded; + }; + // The floor is for a short deadline, never for time already spent. + let usable = Duration::from_millis(deadline_ms.min(MAX_DEADLINE_MS)) + .saturating_sub(DEADLINE_MARGIN) + .max(MIN_BUDGET); + match usable.checked_sub(already) { + Some(left) if !left.is_zero() => Budget::Run(left), + _ => Budget::Exhausted, + } +} + +#[derive(Debug)] +pub struct StageTimedOut { + pub stage: RecallStage, + pub elapsed: Duration, +} + +/// Run `fut` as the recall `marker` tracks, giving up at `deadline`. +pub async fn run_with_deadline( + marker: &StageMarker, + deadline: Option, + fut: F, +) -> Result { + let started = Instant::now(); + let tracked = CURRENT.scope(marker.clone(), fut); + let Some(deadline) = deadline else { + return Ok(tracked.await); + }; + tokio::time::timeout_at(deadline, tracked) + .await + .map_err(|_| StageTimedOut { + stage: marker.get(), + elapsed: started.elapsed(), + }) +} + +/// Logs the stage a recall was in if it is dropped before finishing, which +/// is what happens when the caller hangs up. Covers every caller, including +/// those that send no deadline. +pub struct HangUpGuard { + marker: StageMarker, + owner: String, + started: Instant, + armed: bool, +} + +impl HangUpGuard { + pub fn new(marker: StageMarker, owner: String) -> Self { + Self { + marker, + owner, + started: Instant::now(), + armed: true, + } + } + + /// The recall finished (or answered its own timeout); nothing to report. + pub fn disarm(mut self) { + self.armed = false; + } +} + +impl Drop for HangUpGuard { + fn drop(&mut self) { + // A panic unwinding through the handler is not the caller hanging + // up, and reports itself. + if self.armed && !std::thread::panicking() { + tracing::warn!( + owner = %self.owner, + stage = self.marker.get().as_str(), + elapsed_ms = self.started.elapsed().as_millis() as u64, + "recall abandoned by the caller before it finished" + ); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test(start_paused = true)] + async fn a_missed_deadline_names_the_stage_that_was_running() { + let marker = StageMarker::default(); + let deadline = Instant::now() + Duration::from_secs(14); + let outcome = run_with_deadline(&marker, Some(deadline), async { + enter(RecallStage::Embed); + tokio::time::sleep(Duration::from_secs(1)).await; + enter(RecallStage::WalrusDownload); + tokio::time::sleep(Duration::from_secs(60)).await; + }) + .await; + + let timed_out = outcome.expect_err("a 60s download must not fit a 14s budget"); + assert_eq!(timed_out.stage, RecallStage::WalrusDownload); + assert_eq!(timed_out.elapsed, Duration::from_secs(14)); + } + + #[tokio::test(start_paused = true)] + async fn without_a_budget_the_recall_runs_to_completion() { + // Callers that send no deadline (Python, older SDKs) must see + // exactly today's behaviour: however slow, the recall finishes. + let marker = StageMarker::default(); + let outcome = run_with_deadline(&marker, None, async { + tokio::time::sleep(Duration::from_secs(600)).await; + 7 + }) + .await; + + assert_eq!(outcome.unwrap(), 7); + } + + #[tokio::test] + async fn entering_a_stage_outside_a_recall_does_nothing() { + // `analyze` shares `fetch_batch`, which marks the decrypt step. + enter(RecallStage::SealDecrypt); + } + + #[test] + fn stage_ids_are_the_wire_names() { + assert_eq!(RecallStage::Embed.as_str(), "embed"); + assert_eq!(RecallStage::VectorSearch.as_str(), "vector_search"); + assert_eq!(RecallStage::WalrusDownload.as_str(), "walrus_download"); + assert_eq!(RecallStage::SealDecrypt.as_str(), "seal_decrypt"); + assert_eq!(RecallStage::Auth.as_str(), "auth"); + } + + #[test] + fn a_fresh_marker_starts_at_embed_and_follows_set() { + let marker = StageMarker::default(); + assert_eq!(marker.get(), RecallStage::Embed); + marker.set(RecallStage::SealDecrypt); + assert_eq!(marker.get(), RecallStage::SealDecrypt); + } + + #[test] + fn the_budget_leaves_the_caller_a_second() { + let none = Duration::ZERO; + assert_eq!(budget_for(None, none), Budget::Unbounded); + assert_eq!( + budget_for(Some(15_000), none), + Budget::Run(Duration::from_millis(14_000)) + ); + // Capped, so no caller-supplied number can overflow an `Instant`. + assert_eq!( + budget_for(Some(u64::MAX), none), + Budget::Run(Duration::from_millis(599_000)) + ); + } + + #[test] + fn a_short_deadline_still_gets_the_floor() { + // Less than this cannot fit an embed. + assert_eq!( + budget_for(Some(1_500), Duration::ZERO), + Budget::Run(Duration::from_millis(2_000)) + ); + } + + #[test] + fn time_spent_before_the_handler_comes_out_of_the_budget() { + // Auth on a cold delegate key can take seconds; the caller's clock + // was running through all of it. + assert_eq!( + budget_for(Some(15_000), Duration::from_millis(3_000)), + Budget::Run(Duration::from_millis(11_000)) + ); + // No floor once time is spent: 2s here would outlive the caller. + assert_eq!( + budget_for(Some(15_000), Duration::from_millis(13_500)), + Budget::Run(Duration::from_millis(500)) + ); + } + + #[test] + fn a_deadline_spent_before_the_recall_starts_is_exhausted() { + assert_eq!( + budget_for(Some(15_000), Duration::from_millis(14_000)), + Budget::Exhausted + ); + assert_eq!( + budget_for(Some(15_000), Duration::from_millis(20_000)), + Budget::Exhausted + ); + assert_eq!( + budget_for(Some(1_500), Duration::from_millis(3_000)), + Budget::Exhausted + ); + } +} diff --git a/services/server/src/engine/walrus_seal.rs b/services/server/src/engine/walrus_seal.rs index 54c0e25a6..b5046f006 100644 --- a/services/server/src/engine/walrus_seal.rs +++ b/services/server/src/engine/walrus_seal.rs @@ -240,15 +240,21 @@ impl MemoryEngine for WalrusSealEngine { importance: f32, agent_public_key: Option<&str>, ) -> Result { - // Pick the next Sui key slot (round-robin) so concurrent stores - // don't serialise on one signer. - let key_index = self.key_pool.next_index().ok_or_else(|| { + // Pick the least-loaded Sui key slot so concurrent stores don't + // serialise on one signer. Round-robin alone could land on a wallet + // that is mid-upload while another sits idle; the per-wallet limit is + // 1, so that costs the full upload of whatever is ahead. + let key_index = self.key_pool.least_loaded_index().ok_or_else(|| { AppError::Internal( "No Sui keys configured (set SERVER_SUI_PRIVATE_KEYS or SERVER_SUI_PRIVATE_KEY)" .into(), ) })?; + // Hold the slot for the duration of the upload, so a store running + // concurrently sees this wallet as busy and picks another. + let _wallet_slot = self.key_pool.begin_attempt(key_index); + // Upload the prepared ciphertext to Walrus via the relay sidecar // (pool key pays gas). `defer_transfer = false` — the blob is // transferred to `owner` immediately, same as the inlined @@ -444,6 +450,7 @@ impl MemoryEngine for WalrusSealEngine { ); // Step 2: batch-decrypt the ciphertexts in chunks. + super::stage::enter(super::stage::RecallStage::SealDecrypt); let seal_start = std::time::Instant::now(); let batch_input: Vec<(String, Vec)> = fetched .iter() diff --git a/services/server/src/jobs.rs b/services/server/src/jobs.rs index bdaacd94a..7e42d1180 100644 --- a/services/server/src/jobs.rs +++ b/services/server/src/jobs.rs @@ -202,6 +202,7 @@ async fn update_remember_job_after_wallet_error( remember_job_id: Option<&str>, error: &WalletJobError, msg: &str, + attempt_info: Option, ) { let Some(jid) = remember_job_id else { return; @@ -210,13 +211,17 @@ async fn update_remember_job_after_wallet_error( // Aborting errors (Permanent or ObjectLockedUntilEpoch) get no further // retries, so the row is terminal — mark it failed rather than leaving it // stuck on `running` forever. Retryable errors stay `running` for the next - // attempt. The error_msg carries the lock detail; the object-lock case - // also fires its own distinct Slack alert. - let status = if error.aborts_retries() { - "failed" - } else { - "running" - }; + // attempt — UNLESS this was already the final attempt. Exhausted retries + // used to stay `running` with an error_msg until the 10-minute stale + // sweeper force-failed them, so clients polling GET /api/remember/:job_id + // reported "still uploading" for minutes after every wallet attempt had + // already died (dev 2026-09-17 upload-slot investigation). Pass + // `attempt_info` only from the attempt that actually failed the upload; + // lock-contention Defer callers pass None so a loser cannot mark failed + // while a winner is still working. + let exhausted = attempt_info.is_some_and(|info| info.retries_exhausted(error)); + let terminal = error.aborts_retries() || exhausted; + let status = if terminal { "failed" } else { "running" }; // Terminal means no later attempt will ever insert the row, so the bytes // this job reserved at admission must go back to the owner now rather than @@ -226,7 +231,7 @@ async fn update_remember_job_after_wallet_error( // // Safe to run even when a concurrent attempt already won and released: // release is a delete by id, so a second call is a no-op. - if error.aborts_retries() { + if terminal { crate::storage::db::release_storage_reservations_with_pool(pool, &[jid.to_string()]).await; } @@ -400,6 +405,31 @@ pub fn backoff_duration(attempt: u32) -> std::time::Duration { std::time::Duration::from_secs(2u64.pow(attempt)) } +/// How long a failed upload attempt should pause before Apalis re-queues it, +/// or `None` when no pause is warranted. +/// +/// Apalis attaches no retry/backoff layer (see `WalletJobError`'s doc +/// comment), so a retriable error is re-polled almost immediately and the +/// whole attempt budget burns inside one upstream rate-limit window. Observed +/// in production: a single job took attempts 2, 3, 4 and 5 against Walrus +/// `503 Too Many Requests` within one second, rotating through four wallets +/// that never had a chance to land, and died as "exhausted retries" about a +/// second after its first failure. The upstream limit is time-based, so +/// rotating wallets cannot help — only waiting can. +/// +/// Returns `None` for an aborting error (retrying it is pointless) and for the +/// final attempt (nothing is coming, so the sleep would only delay the +/// failure the caller is already reporting). +fn upload_retry_backoff( + classified: &WalletJobError, + attempt_info: WalletJobAttemptInfo, +) -> Option { + if classified.aborts_retries() || attempt_info.current >= attempt_info.max { + return None; + } + Some(backoff_duration(attempt_info.current as u32)) +} + pub(crate) fn wallet_job_request( job: WalletJob, ) -> Request { @@ -415,14 +445,15 @@ pub(crate) struct WalletJobAttemptInfo { } impl WalletJobAttemptInfo { + fn retries_exhausted(&self, error: &WalletJobError) -> bool { + !error.aborts_retries() && self.current >= self.max + } + fn exhausted_by(&self, error: &WalletJobError) -> bool { if matches!(error, WalletJobError::WalrusBalanceLow(_)) { return false; } - // Only retryable (non-aborting) errors can "exhaust" the budget. An - // aborting error — Permanent or ObjectLockedUntilEpoch — stops retries - // immediately, so it never produces a misleading "exhausted" alert. - !error.aborts_retries() && self.current >= self.max + self.retries_exhausted(error) } } @@ -521,6 +552,13 @@ pub(crate) async fn execute_wallet_job( .into_apalis_error()); } }; + // Mark this wallet busy for the rest of the attempt, so a + // concurrently-enqueued job picks an idle wallet instead of + // queueing behind this upload. Held by guard rather than paired + // calls because every return below — and there are many — has to + // release it. + let _wallet_slot = state.key_pool.begin_attempt(wallet_index); + if wallet_index != enqueued_wallet_index || attempt_info.current > 1 { tracing::info!( "[wallet-job:upload] selected wallet for attempt: enqueued={} executing={} attempt={}/{}", @@ -590,6 +628,19 @@ pub(crate) async fn execute_wallet_job( policy_package_id, end_epoch, } => { + // Mark the wallet busy for this transaction too. `least_loaded_index` + // answers "is this key signing right now", and only the upload arm + // was telling it — so a metadata+transfer, which signs on the very + // same wallet, read as idle. A concurrently-enqueued upload would + // then pick that key precisely because it looked free, and queue + // behind the transaction anyway. That is the failure join-shortest- + // queue exists to avoid, and it showed up as the pool converging on + // whichever key was mid-transfer. + // + // `enqueued_wallet_index` rather than a fresh pick: this operation + // must run on the key that already owns the blob object. + let _wallet_slot = state.key_pool.begin_attempt(enqueued_wallet_index); + let result = execute_set_metadata_and_transfer( state, enqueued_wallet_index, @@ -691,6 +742,7 @@ pub(crate) async fn execute_wallet_job( remember_job_id.as_deref(), &err, &msg, + Some(attempt_info), ) .await; tracing::error!( @@ -874,8 +926,14 @@ async fn insert_vector_and_mark_remember_done( { let msg = format!("insert_vector failed: {}", e); let classified = WalletJobError::classify_sidecar_error(&msg); - update_remember_job_after_wallet_error(state.db.pool(), remember_job_id, &classified, &msg) - .await; + update_remember_job_after_wallet_error( + state.db.pool(), + remember_job_id, + &classified, + &msg, + None, + ) + .await; tracing::error!( "[wallet-job:upload] job_id={} {} classification={} retryable={}", remember_job_id.unwrap_or("-"), @@ -1494,11 +1552,14 @@ async fn execute_upload_and_transfer( attempt_info.max, err, ); + // No attempt_info: a Defer loser must not mark the row failed + // on the final attempt while the lock holder is still working. update_remember_job_after_wallet_error( state.db.pool(), Some(jid.as_str()), &err, err.message(), + None, ) .await; tokio::time::sleep(backoff_duration(attempt_info.current as u32)).await; @@ -1699,6 +1760,7 @@ async fn execute_upload_and_transfer_locked( remember_job_id.as_deref(), &classified, &msg, + None, ) .await; tracing::error!( @@ -1745,11 +1807,23 @@ async fn execute_upload_and_transfer_locked( // reclassifying its display text would incorrectly make it // retryable and leave the polling row running. let msg = err.message().to_string(); + maybe_alert_walrus_upload_exhausted( + state, + &err, + attempt_info, + Some(jid.as_str()), + &owner, + &namespace, + wallet_index, + &msg, + ) + .await; update_remember_job_after_wallet_error( state.db.pool(), Some(jid.as_str()), &err, &msg, + Some(attempt_info), ) .await; tracing::error!( @@ -1759,6 +1833,22 @@ async fn execute_upload_and_transfer_locked( err.kind(), !err.aborts_retries() ); + // The durable upload has its own exit, so it needs its own + // spacing — this is the path a real retry actually took. + // Observed on dev: `durable Walrus upload request failed` + // classified transient at attempt 1/5, with the next attempt + // starting in the same second because only the legacy exit + // below had been given a backoff. + if let Some(delay) = upload_retry_backoff(&err, attempt_info) { + tracing::info!( + "[wallet-job:upload] job_id={} backing off {:?} before attempt {}/{}", + jid, + delay, + attempt_info.current + 1, + attempt_info.max, + ); + tokio::time::sleep(delay).await; + } Err(err) } }; @@ -1883,6 +1973,7 @@ async fn execute_upload_and_transfer_locked( remember_job_id.as_deref(), &classified, &msg, + None, ) .await; @@ -1997,6 +2088,7 @@ async fn execute_upload_and_transfer_locked( remember_job_id.as_deref(), &classified, &msg, + Some(attempt_info), ) .await; tracing::error!( @@ -2006,6 +2098,20 @@ async fn execute_upload_and_transfer_locked( classified.kind(), !classified.aborts_retries() ); + // The wallet slot stays held across this sleep. That is + // deliberate: a backing-off job is still this wallet's turn, and + // releasing it would invite another job onto a wallet that is + // about to retry anyway. + if let Some(delay) = upload_retry_backoff(&classified, attempt_info) { + tracing::info!( + "[wallet-job:upload] job_id={} backing off {:?} before attempt {}/{}", + remember_job_id.as_deref().unwrap_or("-"), + delay, + attempt_info.current + 1, + attempt_info.max, + ); + tokio::time::sleep(delay).await; + } return Err(classified); } }; @@ -2518,6 +2624,58 @@ impl WalletJobError { lower.contains("timed out waiting for") && lower.contains("upload slot") } + /// True if `msg` is one of the sidecar's register-transaction journal + /// assertions (`validatePreparedRegisterTransaction` and friends in + /// scripts/sidecar/routes/walrus-upload-journal.ts). + /// + /// These describe the SHAPE of a transaction the sidecar already built — + /// wrong gas source, missing WAL withdrawal, sender/gas-owner mismatch, + /// non-canonical bytes, digest mismatch. Replaying the journal rebuilds + /// the same shape, so every retry re-fails identically; the job burns its + /// whole retry budget before dying. Classify Permanent so it dies on the + /// first attempt and the caller is told to send the fact again. + /// + /// Anchored on the assertion text rather than the transport's + /// `NO_SIDE_EFFECT` code on purpose: that code means only "nothing reached + /// the chain" and is also returned for pre-submission infra blips + /// (`causeCode: SHARED_SERVICE_UNAVAILABLE` in retry/rpc.ts), which must + /// stay retryable. + pub fn is_register_transaction_shape_error(msg: &str) -> bool { + let lower = msg.to_ascii_lowercase(); + Self::REGISTER_TRANSACTION_SHAPE_ASSERTIONS + .iter() + .any(|assertion| lower.contains(assertion)) + } + + /// The sidecar's register-transaction assertion sentences, lowercased and + /// carrying the `registerTransaction` token they start with. + /// + /// Matched WHOLE, not as a `registertransaction`-anywhere guard plus a + /// phrase-anywhere test. The text being classified is a wrapper + /// (`durable Walrus upload failed (503 …): {json}`) that can carry a + /// nested cause, so those two conditions can be satisfied by unrelated + /// halves of one message: a transient RPC or relay failure that mentions + /// the register step and, somewhere else entirely, a broad phrase like + /// `digest mismatch` or `sender does not match`. Read loosely it becomes + /// `Permanent`, `aborts_retries()` is true, the row is failed on attempt + /// 1, and a write the next attempt would have landed is lost. + /// + /// The `sponsored registerTransaction …` variants are covered by the + /// shorter forms here, which they contain. + const REGISTER_TRANSACTION_SHAPE_ASSERTIONS: &'static [&'static str] = &[ + "registertransaction must use a validduring address-balance expiration", + "registertransaction must pay gas from the address balance", + "registertransaction has no wal address-balance withdrawal", + "registertransaction resolved wal from an owned coin", + "registertransaction resolved the relay tip from an owned sui coin", + "registertransaction must use a distinct gas owner", + "registertransaction sender does not match", + "registertransaction gas owner does not match", + "registertransaction.transactionbytes is not canonical base64", + "registertransaction digest mismatch", + "registertransaction contains invalid transactiondata", + ]; + /// True if `msg` is a pool-wallet WAL shortfall. Deliberately the /// substring half of `parse_wal_balance_alert_info` without its /// `available < WAL_BALANCE_LOW_THRESHOLD_MIST` gate: that threshold @@ -2552,6 +2710,13 @@ impl WalletJobError { if parse_wal_balance_alert_info(msg).is_some() { return WalletJobError::WalrusBalanceLow(msg.to_string()); } + // Register-transaction shape assertions are deterministic for the same + // journal — see is_register_transaction_shape_error. Checked early so a + // shape rejection cannot fall through to the Transient catch-all at the + // end and spend the job's retry budget re-failing identically. + if Self::is_register_transaction_shape_error(msg) { + return WalletJobError::Permanent(msg.to_string()); + } // Sidecar upload limiter saturated — see UploadSlotCongestion docs. if Self::is_upload_slot_congestion_error(msg) { return WalletJobError::UploadSlotCongestion(msg.to_string()); @@ -2863,7 +3028,7 @@ mod tests { is_walrus_package_version_mismatch, load_upload_journal, lock_outcome, mark_remember_job_failed, parse_locked_object_info, parse_wal_balance_alert_info, persist_upload_journal, persist_uploaded_state, recovery_seal_persistence, - update_remember_job_after_wallet_error, upload_resume_disposition, + update_remember_job_after_wallet_error, upload_resume_disposition, upload_retry_backoff, wallet_index_for_upload_attempt, wallet_job_request, JobUploadLock, LockOutcome, UploadResume, WalletJob, WalletJobAttemptInfo, WalletJobError, WalletOperation, MAX_ATTEMPTS, MAX_CONGESTION_REQUEUES, @@ -2882,8 +3047,125 @@ SequenceNumber(884613305), o#B61aVqEgDskxru255FTdzua2RxbbnhDMFxmQ8SCxvj3n) alrea different transaction: TransactionDigest(8bjFgRyXRRYwrzQapgEjpHnGhdfNDY7d6xA82BtHrp3F) \ { k#80127c70.., k#81626d03.. } with 6842 stake]."; + /// The exact production error string from the dev-relayer write outage of + /// 2026-09-17 (job d67d1fc2…, trace 12b3e920…). Every remember on dev failed + /// with this shape and, classified Transient, burned its retry budget before + /// dying — 0 blobs certified over the whole window. + const PROD_REGISTER_SHAPE_ERROR: &str = + "durable Walrus upload failed (503 Service Unavailable): \ +{\"error\":\"registerTransaction must use a ValidDuring address-balance expiration\",\ +\"code\":\"NO_SIDE_EFFECT\",\"traceId\":\"12b3e920-b94b-4100-bb35-fc0f0a1804e1\"}"; + static DB_SETUP_LOCK: OnceLock> = OnceLock::new(); + #[test] + fn a_spent_budget_is_terminal_for_every_retryable_error() { + let spent = WalletJobAttemptInfo { + current: MAX_ATTEMPTS as usize, + max: MAX_ATTEMPTS as usize, + }; + let low = WalletJobError::WalrusBalanceLow("wallet 0 WAL balance low".into()); + + assert!(!spent.exhausted_by(&low)); + assert!(spent.retries_exhausted(&low)); + + for err in [ + WalletJobError::Transient("durable Walrus upload failed (503)".into()), + WalletJobError::UploadSlotCongestion("timed out waiting for upload slot".into()), + ] { + assert!(spent.retries_exhausted(&err), "{}", err.kind()); + } + + let mid = WalletJobAttemptInfo { current: 1, max: MAX_ATTEMPTS as usize }; + assert!(!mid.retries_exhausted(&low)); + let permanent = WalletJobError::Permanent("registerTransaction digest mismatch".into()); + assert!(!spent.retries_exhausted(&permanent) && permanent.aborts_retries()); + } + + #[test] + fn register_shape_rejection_is_permanent() { + assert!(matches!( + WalletJobError::classify_sidecar_error(PROD_REGISTER_SHAPE_ERROR), + WalletJobError::Permanent(_) + )); + } + + #[test] + fn every_register_shape_assertion_is_permanent() { + for msg in [ + "registerTransaction must pay gas from the address balance", + "registerTransaction has no WAL address-balance withdrawal", + "registerTransaction resolved WAL from an owned coin", + "registerTransaction resolved the relay tip from an owned SUI coin", + "sponsored registerTransaction must use a distinct gas owner", + "sponsored registerTransaction sender does not match the wallet", + "registerTransaction gas owner does not match the journaled wallet", + "registerTransaction.transactionBytes is not canonical base64", + "registerTransaction digest mismatch: expected abc, got def", + "registerTransaction contains invalid TransactionData: bad bytes", + ] { + assert!( + matches!( + WalletJobError::classify_sidecar_error(msg), + WalletJobError::Permanent(_) + ), + "expected Permanent for {msg}" + ); + } + } + + /// `NO_SIDE_EFFECT` alone must stay retryable: retry/rpc.ts returns it for + /// any pre-submission failure, including shared-infra blips that the next + /// attempt succeeds through. + #[test] + fn no_side_effect_without_a_shape_assertion_stays_transient() { + assert!(matches!( + WalletJobError::classify_sidecar_error( + "durable Walrus upload failed (503 Service Unavailable): \ +{\"error\":\"fetch failed\",\"code\":\"NO_SIDE_EFFECT\",\ +\"causeCode\":\"SHARED_SERVICE_UNAVAILABLE\"}" + ), + WalletJobError::Transient(_) + )); + } + + /// The shape check keys on the assertion text, so an unrelated message that + /// merely mentions a sender mismatch must not be swallowed by it. + #[test] + fn unrelated_sender_mismatch_is_not_a_shape_rejection() { + assert!(!WalletJobError::is_register_transaction_shape_error( + "sponsor failed: sender does not match the wallet" + )); + } + + /// The assertion must follow the `registerTransaction` token, not merely + /// share a message with it. The classified text is a wrapper that can + /// carry a nested cause, so a transient failure naming the register step + /// in one clause and a broad phrase in another must stay retryable — + /// read as Permanent it dies on attempt 1 and the fact is lost. + #[test] + fn a_broad_phrase_elsewhere_in_the_message_is_not_a_shape_rejection() { + for msg in [ + "durable Walrus upload failed (503 Service Unavailable): \ +{\"error\":\"timed out submitting registerTransaction\",\"code\":\"NO_SIDE_EFFECT\",\ +\"cause\":\"checkpoint digest mismatch on the fullnode\"}", + "registerTransaction step: upstream RPC error, sender does not match \ +the checkpoint it replied about", + ] { + assert!( + !WalletJobError::is_register_transaction_shape_error(msg), + "expected retryable for {msg}" + ); + assert!( + matches!( + WalletJobError::classify_sidecar_error(msg), + WalletJobError::Transient(_) + ), + "expected Transient for {msg}" + ); + } + } + fn test_database_url() -> String { std::env::var("DATABASE_URL") .unwrap_or_else(|_| "postgresql://memwal:memwal_secret@localhost:5432/memwal".into()) @@ -3033,6 +3315,47 @@ different transaction: TransactionDigest(8bjFgRyXRRYwrzQapgEjpHnGhdfNDY7d6xA82Bt assert_eq!(backoff_duration(5), std::time::Duration::from_secs(32)); } + #[test] + fn upload_retry_backs_off_between_attempts() { + // The production failure this exists for: attempts 2..5 against Walrus + // 503 "Too Many Requests" inside one second, four wallet rotations + // that never had a chance to land. Each attempt must now be spaced. + let rate_limited = WalletJobError::Transient( + "durable Walrus upload failed (503 Service Unavailable): {\"error\":\"Too Many Requests\"}" + .into(), + ); + for attempt in 1..5usize { + assert_eq!( + upload_retry_backoff( + &rate_limited, + WalletJobAttemptInfo { + current: attempt, + max: 5 + } + ), + Some(backoff_duration(attempt as u32)), + "attempt {attempt} should wait before the next one", + ); + } + } + + #[test] + fn upload_retry_does_not_back_off_when_nothing_follows() { + let transient = WalletJobError::Transient("upstream hiccup".into()); + // Final attempt: sleeping only delays the failure already being reported. + assert_eq!( + upload_retry_backoff(&transient, WalletJobAttemptInfo { current: 5, max: 5 }), + None + ); + // Aborting errors never retry, so spacing them buys nothing. + let aborting = WalletJobError::GasPoolExhausted("balance::split ENotEnough".into()); + assert!(aborting.aborts_retries()); + assert_eq!( + upload_retry_backoff(&aborting, WalletJobAttemptInfo { current: 1, max: 5 }), + None + ); + } + #[test] fn congestion_backoff_is_minutes_scale_and_capped() { assert_eq!(congestion_backoff_secs(0), 30); @@ -3953,6 +4276,7 @@ different transaction: TransactionDigest(8bjFgRyXRRYwrzQapgEjpHnGhdfNDY7d6xA82Bt Some(job_id.as_str()), &WalletJobError::Transient("another attempt of upload job is in progress".into()), "another attempt of upload job is in progress", + None, ) .await; @@ -3975,6 +4299,76 @@ different transaction: TransactionDigest(8bjFgRyXRRYwrzQapgEjpHnGhdfNDY7d6xA82Bt .await; } + #[tokio::test] + async fn exhausted_transient_upload_marks_the_row_failed_immediately() { + // Dev 2026-09-17: after attempt 5/5 of a Walrus 503 timeout the row + // stayed `running` until the 10-minute stale sweeper. Clients polling + // the job then reported "still uploading" long after every wallet + // attempt was spent. The final attempt must mark failed itself. + let pool = test_pool().await; + let job_id = format!("remember-job-{}", uuid::Uuid::new_v4()); + insert_job_with_status(&pool, &job_id, "running", None).await; + + let err = WalletJobError::Transient( + "Internal Error: durable Walrus upload failed (503 Service Unavailable)".into(), + ); + update_remember_job_after_wallet_error( + &pool, + Some(job_id.as_str()), + &err, + err.message(), + Some(WalletJobAttemptInfo { + current: MAX_ATTEMPTS as usize, + max: MAX_ATTEMPTS as usize, + }), + ) + .await; + + let row: (String, Option) = + sqlx::query_as("SELECT status, error_msg FROM remember_jobs WHERE id = $1") + .bind(&job_id) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(row.0, "failed"); + assert!( + row.1.as_deref().unwrap_or("").contains("503"), + "error_msg must keep the upload failure: {:?}", + row.1 + ); + + let _ = sqlx::query("DELETE FROM remember_jobs WHERE id = $1") + .bind(&job_id) + .execute(&pool) + .await; + + // An earlier attempt must still leave the row running for the next try. + let mid_id = format!("remember-job-{}", uuid::Uuid::new_v4()); + insert_job_with_status(&pool, &mid_id, "running", None).await; + update_remember_job_after_wallet_error( + &pool, + Some(mid_id.as_str()), + &err, + err.message(), + Some(WalletJobAttemptInfo { + current: (MAX_ATTEMPTS as usize) - 1, + max: MAX_ATTEMPTS as usize, + }), + ) + .await; + let mid: (String,) = sqlx::query_as("SELECT status FROM remember_jobs WHERE id = $1") + .bind(&mid_id) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(mid.0, "running"); + + let _ = sqlx::query("DELETE FROM remember_jobs WHERE id = $1") + .bind(&mid_id) + .execute(&pool) + .await; + } + // RC-4: an uploaded-but-pending resume must route to the TRANSFER recovery op // (carrying the stored object id), NOT to an index/done finalize — otherwise a // never-transferred blob is prematurely marked done. diff --git a/services/server/src/main.rs b/services/server/src/main.rs index 292bdd6c1..4de897dca 100644 --- a/services/server/src/main.rs +++ b/services/server/src/main.rs @@ -14,6 +14,7 @@ mod routes; mod security_delete_auth; mod security_delete_error; mod services; +mod sidecar_saturation; mod storage; mod sui; mod types; @@ -708,10 +709,9 @@ async fn init_apalis_pool( // `lock_timeout` is intentionally omitted. A leaked session lock_timeout // on a transaction-mode pooler backend is what aborted sqlx migrate // (15s wait → panic). Migrate now uses the direct host, but we still - // don't put lock_timeout on pooled backends. statement_timeout and - // idle_in_transaction_session_timeout bound queries without aborting - // lock waits; if they leak onto another pooler client they are still - // a bound, not a migrate-killer. + // don't put lock_timeout on pooled backends. statement_timeout is later, + // not exempt: it covers the lock wait too, so it still cancels a blocked + // statement, just at 300s instead of 15s. let statement_timeout = format!("{}ms", startup_timeout.as_millis().min(300_000)); let pool_future = PgPoolOptions::new() .max_connections(10) @@ -1431,10 +1431,10 @@ async fn main() { // Sidecar upload-queue saturation monitor. The watchdog above only // checks that /health answers; during the 2026-06-10 congestion incident // it stayed green while 120 uploads queued and jobs burned their retry - // budgets. This monitor reads the queue counters that /health already - // exposes and alerts ops while there is still time to act (add wallets / - // throttle the burst) — before queued requests outlive the sidecar's - // 120s acquire timeout and start failing. + // budgets. This monitor reads the sidecar's upload-queue counters and + // alerts ops while there is still time to act (add wallets / throttle the + // burst) — before queued requests outlive the sidecar's 120s acquire + // timeout and start failing. let saturation_threshold = parse_env_u64("SIDECAR_QUEUE_SATURATION_THRESHOLD", 20, 1, 10_000); let saturation_consecutive = parse_env_u32("SIDECAR_QUEUE_SATURATION_CONSECUTIVE", 4, 1, 100); let saturation_interval_secs = @@ -1447,13 +1447,16 @@ async fn main() { ); { let monitor_client = state.http_client.clone(); - let monitor_url = health_url.clone(); + let monitor_url = format!("{}{}", sidecar_url, sidecar_saturation::UPLOAD_METRICS_PATH); let monitor_alerts = Arc::clone(&state.alerts); let monitor_network = config.sui_network.clone(); tokio::spawn(async move { let mut interval = tokio::time::interval(std::time::Duration::from_secs(saturation_interval_secs)); - let mut consecutive_saturated = 0u32; + let mut tracker = sidecar_saturation::SaturationTracker::new( + saturation_threshold, + saturation_consecutive, + ); loop { interval.tick().await; let body = match monitor_client @@ -1465,57 +1468,79 @@ async fn main() { Ok(resp) if resp.status().is_success() => { match resp.json::().await { Ok(v) => v, - Err(_) => continue, + Err(err) => { + tracing::error!( + " sidecar: upload metrics body is not JSON, saturation alert is blind: {}", + err + ); + continue; + } } } - // Liveness problems are the watchdog's job; only the - // healthy-but-saturated case belongs here. - _ => continue, + // The sidecar answered, so the watchdog stays green; a + // non-2xx here means the metrics route itself is broken. + Ok(resp) => { + tracing::error!( + " sidecar: {} returned status={}, saturation alert is blind", + sidecar_saturation::UPLOAD_METRICS_PATH, + resp.status() + ); + continue; + } + // An unreachable sidecar is the watchdog's job. + Err(_) => continue, }; - let queued = body["queuedWalrusUploads"].as_u64().unwrap_or(0); - let active = body["activeWalrusUploads"].as_u64().unwrap_or(0); - let global_capacity = body["walrusUploadLimits"]["globalCapacity"] - .as_u64() - .unwrap_or(0); - - if queued > saturation_threshold { - consecutive_saturated = consecutive_saturated.saturating_add(1); - tracing::warn!( - " sidecar: upload queue saturated queued={} active={} capacity={} consecutive={}/{}", - queued, - active, - global_capacity, - consecutive_saturated, - saturation_consecutive, - ); - } else { - if consecutive_saturated >= saturation_consecutive { + let sample = match sidecar_saturation::parse_upload_metrics(&body) { + Ok(sample) => sample, + Err(field) => { + tracing::error!( + " sidecar: upload metrics missing {}, saturation alert is blind", + field + ); + continue; + } + }; + + match tracker.observe(sample.queued) { + sidecar_saturation::QueueCheck::Clear => {} + sidecar_saturation::QueueCheck::Drained => { tracing::info!( " sidecar: upload queue drained (queued={} <= threshold {})", - queued, + sample.queued, saturation_threshold, ); } - consecutive_saturated = 0; - } - - // Alert once per crossing; the AlertManager dedup window - // handles re-alerting if the backlog persists. - if consecutive_saturated >= saturation_consecutive { - let alert = crate::alerts::WalrusUploadQueueSaturatedAlert { - sui_network: monitor_network.clone(), - queued, - active, - global_capacity, - threshold: saturation_threshold, - consecutive_checks: consecutive_saturated, - }; - if let Err(err) = monitor_alerts - .notify_walrus_upload_queue_saturated(alert) - .await - { - tracing::warn!(" sidecar: saturation alert delivery failed: {}", err); + sidecar_saturation::QueueCheck::Saturated { consecutive, alert } => { + tracing::warn!( + " sidecar: upload queue saturated queued={} active={} capacity={} consecutive={}/{}", + sample.queued, + sample.active, + sample.global_capacity, + consecutive, + saturation_consecutive, + ); + // Alert on every saturated check past the streak; the + // AlertManager dedup window handles re-alerting. + if alert { + let alert = crate::alerts::WalrusUploadQueueSaturatedAlert { + sui_network: monitor_network.clone(), + queued: sample.queued, + active: sample.active, + global_capacity: sample.global_capacity, + threshold: saturation_threshold, + consecutive_checks: consecutive, + }; + if let Err(err) = monitor_alerts + .notify_walrus_upload_queue_saturated(alert) + .await + { + tracing::warn!( + " sidecar: saturation alert delivery failed: {}", + err + ); + } + } } } } @@ -1981,6 +2006,11 @@ async fn main() { // Mode-blind; owner-scoped via AuthInfo. .route("/api/forget", post(routes::forget)) .route("/api/stats", post(routes::stats)) + // Identity echo — tells an authenticated caller which account its + // delegate key resolves to. Must stay inside this authed group; the + // account_id it returns is deliberately withheld from the public + // /api/accounts/{owner}/exists route. See routes::whoami. + .route("/api/whoami", get(routes::whoami)) // Router::layer runs middleware bottom-to-top (last added runs first). // Keep auth outer so AuthInfo is in request extensions before rate limiting reads it. .layer(middleware::from_fn_with_state( diff --git a/services/server/src/observability.rs b/services/server/src/observability.rs index 0f9707d32..f69553095 100644 --- a/services/server/src/observability.rs +++ b/services/server/src/observability.rs @@ -43,6 +43,9 @@ const DEFAULT_OTEL_SERVICE_NAME: &str = "memwal-relayer"; pub struct RequestContext { request_id: String, route: String, + /// When the request reached the relayer, before auth and rate limiting, + /// so a caller-supplied deadline can be measured from arrival. + started: Instant, } tokio::task_local! { @@ -460,6 +463,7 @@ pub async fn request_context_middleware(mut request: Request, next: Next) -> Res request.extensions_mut().insert(RequestContext { request_id: request_id.clone(), route: route.clone(), + started, }); if let Ok(value) = HeaderValue::from_str(&request_id) { request @@ -486,6 +490,7 @@ pub async fn request_context_middleware(mut request: Request, next: Next) -> Res let context = RequestContext { request_id: request_id.clone(), route: route.clone(), + started, }; REQUEST_CONTEXT @@ -545,6 +550,11 @@ pub fn current_request_id() -> Option { REQUEST_CONTEXT.try_with(|ctx| ctx.request_id.clone()).ok() } +/// When the current request reached the relayer, if this task serves one. +pub fn current_request_started() -> Option { + REQUEST_CONTEXT.try_with(|ctx| ctx.started).ok() +} + pub fn current_context() -> Option { REQUEST_CONTEXT.try_with(Clone::clone).ok() } diff --git a/services/server/src/routes/accounts.rs b/services/server/src/routes/accounts.rs index 4c0af8b5f..4daad78ad 100644 --- a/services/server/src/routes/accounts.rs +++ b/services/server/src/routes/accounts.rs @@ -6,7 +6,7 @@ //! Console to hold a delegate key or run its own onchain scan. use axum::extract::{Path, State}; -use axum::Json; +use axum::{Extension, Json}; use std::sync::Arc; use crate::routes::sponsor::validate_sui_address; @@ -74,6 +74,46 @@ pub async fn account_exists( Ok(Json(AccountExistsResponse { exists })) } +/// GET /api/whoami +/// +/// Returns the account identity the caller's delegate key resolves to. +/// +/// Authenticated, and that is what makes returning `account_id` here +/// acceptable where `account_exists` deliberately withholds it: the auth +/// middleware has already proven the caller holds a delegate key registered +/// against this account, so this only ever tells a caller about itself. Do +/// not move this into the unauthenticated router group. +/// +/// Motivating case (WALM-332): a login interrupted between the browser's +/// on-chain `add_delegate_key` and the localhost callback leaves the client +/// holding a valid delegate key but none of the surrounding metadata, so it +/// cannot write a usable `credentials.json`. Everything needed to rebuild one +/// is already resolved during authentication — `account_id` and `owner` from +/// the registry scan, `package_id` from config — so this endpoint hands back +/// what the middleware already computed rather than doing new work. +pub async fn whoami( + State(state): State>, + Extension(auth): Extension, +) -> Result, AppError> { + Ok(Json(whoami_response(auth, state.config.package_id.clone()))) +} + +/// The field mapping, factored out of the handler so it is testable without a +/// live `AppState`/DB (same reason as `account_exists` above — this codebase +/// has no axum-handler test harness). +/// +/// Worth isolating rather than inlining: `account_id` and `owner` are both +/// 0x-prefixed 32-byte hex, so transposing them is invisible to the type +/// checker and would produce credentials that authenticate as the wrong +/// identity. The test below pins the mapping. +fn whoami_response(auth: AuthInfo, package_id: String) -> WhoamiResponse { + WhoamiResponse { + account_id: auth.account_id, + owner: auth.owner, + package_id, + } +} + #[cfg(test)] mod tests { use super::*; @@ -103,4 +143,82 @@ mod tests { // Already-lowercase input must be unaffected (idempotent). assert_eq!(normalized.to_ascii_lowercase(), normalized); } + + // ── GET /api/whoami (WALM-332 recovery) ────────────────────── + + fn auth_fixture() -> AuthInfo { + AuthInfo { + public_key: "aa".repeat(32), + owner: format!("0x{}", "b".repeat(64)), + account_id: format!("0x{}", "a".repeat(64)), + delegate_key: None, + seal_session: None, + } + } + + /// `account_id` and `owner` are indistinguishable by type — both + /// 0x-prefixed 32-byte hex — so a transposition here would compile, pass + /// every other test, and hand a recovering client credentials that + /// authenticate as the wrong identity. Pin the mapping explicitly. + #[test] + fn whoami_maps_account_and_owner_without_transposing_them() { + let auth = auth_fixture(); + let package_id = format!("0x{}", "c".repeat(64)); + + let resp = whoami_response(auth, package_id.clone()); + + assert_eq!( + resp.account_id, + format!("0x{}", "a".repeat(64)), + "account_id must come from AuthInfo::account_id, not owner" + ); + assert_eq!( + resp.owner, + format!("0x{}", "b".repeat(64)), + "owner must come from AuthInfo::owner, not account_id" + ); + assert_eq!(resp.package_id, package_id, "package_id comes from config"); + assert_ne!(resp.account_id, resp.owner, "fixture must distinguish them"); + } + + /// The recovering MCP client cannot send `x-account-id` — not knowing it + /// is the whole reason it is calling this endpoint — so it signs the + /// canonical message with an empty account field and omits the header. + /// The server defaults the hint to `""` for exactly this case. + /// + /// This pins the literal string both sides must build. It is duplicated + /// verbatim in `packages/mcp/test/login-recovery-signing.test.mjs`; if the + /// canonical format in `auth.rs` ever changes, both fail together rather + /// than recovery silently breaking in production. + #[test] + fn whoami_recovery_request_canonical_message_is_stable() { + let timestamp = "1700000000"; + let method = "GET"; + let path = "/api/whoami"; + // sha256 of an empty body — a GET carries none, but the server hashes + // the empty body all the same. + let body_hash = "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"; + let nonce = "550e8400-e29b-41d4-a716-446655440000"; + let account_id_for_sig = String::new(); + + let message = format!( + "{}.{}.{}.{}.{}.{}", + timestamp, method, path, body_hash, nonce, account_id_for_sig + ); + + let expected = concat!( + "1700000000.GET./api/whoami.", + "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855.", + "550e8400-e29b-41d4-a716-446655440000." + ); + assert_eq!(message, expected); + // Six fields → five separators. Nothing else in the message contains a + // dot: the nonce is hyphen-separated and the body hash is bare hex. + assert_eq!(message.matches('.').count(), 5); + assert!( + message.ends_with('.'), + "the empty account id leaves a trailing separator — the client must \ + reproduce this exactly, not trim it" + ); + } } diff --git a/services/server/src/routes/admin.rs b/services/server/src/routes/admin.rs index 501e7b03e..9177c977c 100644 --- a/services/server/src/routes/admin.rs +++ b/services/server/src/routes/admin.rs @@ -150,11 +150,26 @@ pub async fn health(State(state): State>) -> Json ask: ASK_SYSTEM_PROMPT_VERSION.to_string(), }, write_ready: write_ready(&state).await, - writes: writes_health_status(state.config.writes_paused), + writes: writes_health_status( + state.config.writes_paused, + durable_writes_degraded_probe(&state).await, + ), }) } const WRITE_READY_CACHE_TTL: std::time::Duration = std::time::Duration::from_secs(2); +/// How far back `/health` looks when judging whether durable writes land. +/// Long enough to span a few upload attempts with their backoff, short +/// enough that recovery shows up without an operator waiting. +const DURABLE_WRITE_WINDOW: std::time::Duration = std::time::Duration::from_secs(15 * 60); +/// Failures inside the window before the write path is called degraded. +/// Writes fail individually all the time; one or two is not an outage. +const DURABLE_WRITE_FAILURE_THRESHOLD: i64 = 3; +/// Cached far longer than `write_ready` so `/health` does not hit +/// `remember_jobs` on every load-balancer tick. 022's partial +/// `updated_at` index is the steady-state path; the 1s bound is a +/// backstop, not the expected scan. +const DURABLE_WRITE_CACHE_TTL: std::time::Duration = std::time::Duration::from_secs(30); const WRITE_READY_PROBE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(1); /// Neon refuses `smgrextend` once cluster size is at the cap; treat less /// than 1MB remaining as not writable so `/health` trips before the next @@ -316,6 +331,71 @@ fn postgres_can_accept_writes(used_bytes: i64, max_bytes: i64) -> bool { static WRITE_READY_CACHE: std::sync::Mutex> = std::sync::Mutex::new(None); +static DURABLE_WRITE_CACHE: std::sync::Mutex> = + std::sync::Mutex::new(None); + +/// Whether recent durable writes say the write path is degraded. +/// +/// Silence is not failure: a window that finished no writes leaves this +/// false, so a quiet deployment never reports itself broken. Nor is a +/// single loss -- only a window that failed at least +/// `DURABLE_WRITE_FAILURE_THRESHOLD` times AND landed nothing at all +/// counts, which is what a downstream outage looks like and what a run of +/// unlucky individual writes does not. +fn durable_writes_degraded(failed: i64, succeeded: i64) -> bool { + if succeeded > 0 { + return false; + } + failed >= DURABLE_WRITE_FAILURE_THRESHOLD +} + +/// Reads recent write outcomes behind `DURABLE_WRITE_CACHE_TTL`. +/// +/// Fails open, like the Postgres probe: a database that cannot answer this +/// is not evidence that Walrus is down, and `/health` must not invent an +/// outage out of its own query failing. Cached so `/health` does not hit +/// `remember_jobs` on every load-balancer tick. 022's partial `updated_at` +/// index keeps the 1s `WRITE_READY_PROBE_TIMEOUT` a backstop rather than +/// the steady state. +async fn durable_writes_degraded_probe(state: &std::sync::Arc) -> bool { + { + let cache = DURABLE_WRITE_CACHE + .lock() + .unwrap_or_else(|e| e.into_inner()); + if let Some((at, degraded)) = *cache { + if at.elapsed() < DURABLE_WRITE_CACHE_TTL { + return degraded; + } + } + } + + let degraded = match tokio::time::timeout( + WRITE_READY_PROBE_TIMEOUT, + state.db.recent_write_outcomes(DURABLE_WRITE_WINDOW), + ) + .await + { + Ok(Ok((failed, succeeded))) => durable_writes_degraded(failed, succeeded), + Ok(Err(err)) => { + tracing::warn!( + error = %err, + "durable-write outcome probe failed; leaving the write path reported healthy" + ); + false + } + Err(_) => { + tracing::warn!( + "durable-write outcome probe timed out; leaving the write path reported healthy" + ); + false + } + }; + if let Ok(mut cache) = DURABLE_WRITE_CACHE.lock() { + *cache = Some((std::time::Instant::now(), degraded)); + } + degraded +} + /// GET /version pub async fn version() -> Json { Json(crate::compatibility::version_response()) @@ -1306,6 +1386,21 @@ mod tests { ); } + #[test] + fn durable_writes_degraded_only_when_nothing_lands() { + let threshold = super::DURABLE_WRITE_FAILURE_THRESHOLD; + // No finished writes at all: a quiet window, not an outage. + assert!(!super::durable_writes_degraded(0, 0)); + // Below the threshold: individual writes do fail. + assert!(!super::durable_writes_degraded(threshold - 1, 0)); + // At and past it, with nothing landing: this is the outage shape. + assert!(super::durable_writes_degraded(threshold, 0)); + assert!(super::durable_writes_degraded(threshold * 100, 0)); + // A single success means the path works, however many failed + // alongside it -- Walrus is storing blobs, these writes lost. + assert!(!super::durable_writes_degraded(threshold * 100, 1)); + } + #[test] fn postgres_can_accept_writes_false_at_or_within_1mb_of_cap() { let max = 3072 * 1024 * 1024; @@ -1328,6 +1423,16 @@ mod tests { ); } + #[test] + fn durable_write_probe_shares_the_one_second_bound() { + // A hung recent_write_outcomes scan must not stall /health past this; + // the probe fails open, same as the Postgres size check. + assert_eq!( + super::WRITE_READY_PROBE_TIMEOUT, + std::time::Duration::from_secs(1) + ); + } + #[test] fn missing_pg_cluster_size_falls_back_instead_of_fail_open() { let missing = sqlx::Error::Database(Box::new(FakePgError { diff --git a/services/server/src/routes/analyze.rs b/services/server/src/routes/analyze.rs index d8aa1c8d9..b96ad82a3 100644 --- a/services/server/src/routes/analyze.rs +++ b/services/server/src/routes/analyze.rs @@ -865,7 +865,7 @@ pub async fn analyze( } // Pick next wallet slot (round-robin) and enqueue UploadAndTransfer - let Some(wallet_index) = state.key_pool.next_index() else { + let Some(wallet_index) = state.key_pool.least_loaded_index() else { rate_limit::release_storage_quota(&state, &all_ids[idx..]).await; return Err(AppError::Internal("No Sui keys configured".into())); }; diff --git a/services/server/src/routes/mod.rs b/services/server/src/routes/mod.rs index abd037568..c9f17f349 100644 --- a/services/server/src/routes/mod.rs +++ b/services/server/src/routes/mod.rs @@ -38,7 +38,7 @@ mod sponsor; // Re-export every handler so `main.rs` keeps using `routes::` // without having to know which submodule each handler lives in. -pub use accounts::account_exists; +pub use accounts::{account_exists, whoami}; pub use admin::{ask, embed, forget, get_config, health, restore, stats, version}; pub use analyze::analyze; pub use memory_read::{list_owner_agents, list_owner_memories, list_owner_namespaces}; @@ -249,6 +249,21 @@ pub(super) fn select_hits_for_sort( hits } +/// The scoring weights `/api/recall` hands to the ranker. An explicit `sort` +/// must not be re-ranked, so it yields the default (no-op) weights; malformed +/// weights are rejected either way. +pub(super) fn resolve_scoring_weights( + sort: Option, + requested: Option, +) -> Result { + let requested = requested.unwrap_or_default(); + requested.validate()?; + if sort.is_some() { + return Ok(crate::types::ScoringWeights::default()); + } + Ok(requested) +} + /// Project the ranker's output onto the `/api/recall` and `/api/ask` wire /// shape, preserving ranked order. /// @@ -461,6 +476,130 @@ mod recall_sort_tests { } } +/// Each test replays the recall handler's ordering (`select_hits_for_sort`, +/// then the ranker with the resolved weights) over hand-scored rows. +#[cfg(test)] +mod recall_sort_precedence_tests { + use crate::engine::HydratedMemory; + use crate::services::ranker::{CompositeRanker, Ranker}; + use crate::types::{AppError, RecallSort, ScoringWeights, SearchHit}; + + fn ts(rfc3339: &str) -> chrono::DateTime { + chrono::DateTime::parse_from_rfc3339(rfc3339) + .unwrap() + .with_timezone(&chrono::Utc) + } + + fn now() -> chrono::DateTime { + ts("2026-07-06T00:00:00Z") + } + + /// `distance` ascending is the order pgvector hands back. + fn search_hit(blob_id: &str, distance: f64, created_at: &str) -> SearchHit { + SearchHit { + blob_id: blob_id.to_string(), + distance, + created_at: ts(created_at), + importance: 0.5, + } + } + + fn weights(recency: f64) -> ScoringWeights { + ScoringWeights { + semantic: 1.0, + recency, + recency_half_life_days: 30.0, + importance: 0.0, + } + } + + fn recall_order( + hits: Vec, + sort: Option, + requested: ScoringWeights, + ) -> Vec { + let weights = super::resolve_scoring_weights(sort, Some(requested)).unwrap(); + let limit = hits.len(); + let hydrated: Vec = + super::select_hits_for_sort(hits, sort.unwrap_or_default(), limit) + .into_iter() + .map(|h| HydratedMemory { + blob_id: h.blob_id, + text: String::new(), + distance: h.distance, + created_at: Some(h.created_at), + importance: Some(h.importance), + }) + .collect(); + CompositeRanker + .rank(hydrated, &weights, now()) + .into_iter() + .map(|r| r.memory.blob_id) + .collect() + } + + /// Composite scores at semantic 1.0 / recency 0.3 / half-life 30d: + /// old-but-close = 0.90 + 0.3 * 2^(-30/30) = 1.05 + /// newest-but-far = 0.10 + 0.3 * 2^(0/30) = 0.40 + fn ticket_repro_hits() -> Vec { + vec![ + search_hit("old-but-close", 0.10, "2026-06-06T00:00:00Z"), + search_hit("newest-but-far", 0.90, "2026-07-06T00:00:00Z"), + ] + } + + /// Composite scores at semantic 1.0 / recency 0.8 / half-life 30d: + /// close-old = 0.90 + 0.8 * 2^(-60/30) = 1.10 + /// far-new = 0.80 + 0.8 * 2^(0/30) = 1.60 + /// so the ranker, when it runs, reverses the cosine order. + fn weights_reverse_cosine_hits() -> Vec { + vec![ + search_hit("close-old", 0.10, "2026-05-07T00:00:00Z"), + search_hit("far-new", 0.20, "2026-07-06T00:00:00Z"), + ] + } + + #[test] + fn explicit_recent_is_not_reordered_by_scoring_weights() { + assert_eq!( + recall_order(ticket_repro_hits(), Some(RecallSort::Recent), weights(0.3)), + ["newest-but-far", "old-but-close"] + ); + } + + #[test] + fn explicit_relevance_is_not_reordered_by_scoring_weights() { + assert_eq!( + recall_order( + weights_reverse_cosine_hits(), + Some(RecallSort::Relevance), + weights(0.8) + ), + ["close-old", "far-new"] + ); + } + + #[test] + fn omitted_sort_lets_scoring_weights_rerank() { + assert_eq!( + recall_order(weights_reverse_cosine_hits(), None, weights(0.8)), + ["far-new", "close-old"] + ); + } + + /// Malformed weights are a 400 even when an explicit `sort` would ignore + /// them — `sort` does not quietly excuse a bad request. + #[test] + fn malformed_weights_are_rejected_even_when_sort_is_explicit() { + let mut malformed = weights(0.3); + malformed.semantic = f64::NAN; + assert!(matches!( + super::resolve_scoring_weights(Some(RecallSort::Recent), Some(malformed)), + Err(AppError::BadRequest(_)) + )); + } +} + #[cfg(test)] mod recall_result_mapping_tests { use crate::engine::HydratedMemory; diff --git a/services/server/src/routes/oauth.rs b/services/server/src/routes/oauth.rs index b564e590b..4fdc7a586 100644 --- a/services/server/src/routes/oauth.rs +++ b/services/server/src/routes/oauth.rs @@ -24,7 +24,7 @@ use crate::oauth::{self, OAuthError}; use crate::storage::db::oauth_rows::{ OAuthClientRow, OAuthCodeRow, OAuthDelegateRow, OAuthGrantRow, OAuthSessionRow, OAuthTokenRow, }; -use crate::storage::sui::verify_delegate_key_onchain; +use crate::storage::sui::{verify_delegate_key_onchain, GET_OBJECT_ATTEMPTS}; use crate::types::AppState; /// Read the configured `McpOAuthConfig`, or fail with 404 if not configured. @@ -690,6 +690,7 @@ pub async fn session_complete( &account_id, &public_key_bytes, &state.config.package_id, + GET_OBJECT_ATTEMPTS, ) .await .map_err(|e| { @@ -955,9 +956,10 @@ pub async fn token( .as_deref() .ok_or_else(|| OAuthError::invalid_request("code_verifier is required"))?; + let code_sha256 = oauth::hash_token(code); let code_row = state .db - .consume_oauth_code(&client.client_id, &oauth::hash_token(code)) + .fetch_oauth_code(&client.client_id, &code_sha256) .await? .ok_or_else(|| { OAuthError::invalid_grant( @@ -984,6 +986,16 @@ pub async fn token( } } + let code_row = state + .db + .consume_oauth_code(&client.client_id, &code_sha256) + .await? + .ok_or_else(|| { + OAuthError::invalid_grant( + "authorization code is invalid, expired, or already used", + ) + })?; + let grant_id = oauth::random_token(oauth::GRANT_ID_PREFIX); state .db @@ -1140,4 +1152,62 @@ mod tests { .await .unwrap(); } + + #[tokio::test] + async fn fetch_oauth_code_reads_without_consuming() { + let url = std::env::var("DATABASE_URL") + .unwrap_or_else(|_| "postgresql://memwal:memwal_secret@localhost:5432/memwal".into()); + let pool = sqlx::postgres::PgPoolOptions::new() + .max_connections(1) + .acquire_timeout(std::time::Duration::from_secs(5)) + .connect(&url) + .await + .unwrap(); + sqlx::raw_sql(include_str!("../../migrations/011_mcp_oauth.sql")) + .execute(&pool) + .await + .unwrap(); + let db = crate::storage::db::VectorDb::from_pool(pool); + + let client_id = format!("mcpc_{}", uuid::Uuid::new_v4().simple()); + let code_sha256 = format!("{:064x}", uuid::Uuid::new_v4().as_u128()); + db.insert_oauth_code(&super::OAuthCodeRow { + code_sha256: code_sha256.clone(), + client_id: client_id.clone(), + redirect_uri: "https://claude.ai/api/mcp/auth_callback".into(), + scope: "memwal".into(), + resource: "https://relayer.example/mcp".into(), + code_challenge: "challenge".into(), + code_challenge_method: "S256".into(), + delegate_ref: "mwd_ref".into(), + account_id: format!("0x{}", "11".repeat(32)), + owner_address: format!("0x{}", "22".repeat(32)), + expires_at: chrono::Utc::now() + chrono::Duration::minutes(5), + }) + .await + .unwrap(); + + let read = db + .fetch_oauth_code(&client_id, &code_sha256) + .await + .unwrap() + .expect("fetch returns the code"); + assert_eq!(read.code_sha256, code_sha256); + let again = db.fetch_oauth_code(&client_id, &code_sha256).await.unwrap(); + assert!(again.is_some(), "fetch must leave the code redeemable"); + + let consumed = db + .consume_oauth_code(&client_id, &code_sha256) + .await + .unwrap() + .expect("consume returns the code"); + assert_eq!(consumed.code_sha256, code_sha256); + assert!( + db.consume_oauth_code(&client_id, &code_sha256) + .await + .unwrap() + .is_none(), + "the code must stay single-use" + ); + } } diff --git a/services/server/src/routes/recall.rs b/services/server/src/routes/recall.rs index 018c1faf9..6bbef7752 100644 --- a/services/server/src/routes/recall.rs +++ b/services/server/src/routes/recall.rs @@ -13,6 +13,7 @@ use redis::AsyncCommands; use sha2::Digest; use std::sync::Arc; +use crate::engine::stage::{self, Budget, HangUpGuard, RecallStage, StageMarker}; use crate::types::*; // ============================================================ @@ -162,22 +163,107 @@ pub async fn recall( // Validate scoring_weights up front — fail fast on malformed input // (NaN, out-of-range, sub-floor half-life) BEFORE we spend an embed + // vector search + Walrus + SEAL round-trip just to 400 at the end. - let weights = body.scoring_weights.clone().unwrap_or_default(); - weights.validate()?; + // An explicit `sort` suppresses them; see `resolve_scoring_weights`. + let weights = super::resolve_scoring_weights(body.sort, body.scoring_weights.clone())?; + let sort = body.sort.unwrap_or_default(); // Owner is derived from delegate key via onchain verification (auth middleware) let owner = &auth.owner; let namespace = &body.namespace; + tracing::info!( query_len = body.query.len(), owner = %owner, namespace = %namespace, ranker_active = weights.is_ranker_active(), + scoring_weights_ignored = body.sort.is_some() + && body + .scoring_weights + .as_ref() + .is_some_and(ScoringWeights::is_ranker_active), "recall request" ); + // Stop just short of the caller's deadline, so the 504 can say which + // stage was running. The caller's clock started before auth did. + let already = crate::observability::current_request_started() + .map(|arrived| arrived.elapsed()) + .unwrap_or_default(); + let deadline = match stage::budget_for(body.deadline_ms, already) { + Budget::Unbounded => None, + Budget::Run(left) => Some(tokio::time::Instant::now() + left), + Budget::Exhausted => { + return Err(recall_timed_out( + owner, + namespace, + RecallStage::Auth, + already, + &body, + )); + } + }; + let marker = StageMarker::default(); + let guard = HangUpGuard::new(marker.clone(), owner.clone()); + let outcome = stage::run_with_deadline( + &marker, + deadline, + recall_pipeline(&state, &auth, &body, &weights, sort), + ) + .await; + guard.disarm(); + let response = match outcome { + Ok(result) => result?, + Err(timed_out) => { + let elapsed = already + timed_out.elapsed; + return Err(recall_timed_out( + owner, + namespace, + timed_out.stage, + elapsed, + &body, + )); + } + }; + Ok(Json(response)) +} + +fn recall_timed_out( + owner: &str, + namespace: &str, + stage: RecallStage, + elapsed: std::time::Duration, + body: &RecallRequest, +) -> AppError { + let elapsed_ms = elapsed.as_millis() as u64; + tracing::warn!( + owner = %owner, + namespace = %namespace, + stage = stage.as_str(), + elapsed_ms, + deadline_ms = body.deadline_ms, + "recall timed out before the caller's deadline" + ); + AppError::RecallTimeout { + stage: stage.as_str(), + elapsed_ms, + } +} + +/// Everything `recall` does after validation, as one future the handler can +/// put a deadline on. Each `stage::enter` names the step a timeout reports. +async fn recall_pipeline( + state: &Arc, + auth: &AuthInfo, + body: &RecallRequest, + weights: &ScoringWeights, + sort: RecallSort, +) -> Result { + let owner = &auth.owner; + let namespace = &body.namespace; + + stage::enter(RecallStage::Embed); let t0 = std::time::Instant::now(); - let query_vector = generate_recall_embedding_cached(&state, &body.query).await?; + let query_vector = generate_recall_embedding_cached(state, &body.query).await?; let embed_ms = t0.elapsed().as_millis(); // Cap limit to prevent unbounded DB scans / memory use. @@ -187,7 +273,8 @@ pub async fn recall( // row is frequently a mediocre semantic match and would otherwise fall // outside the cosine top-`limit` entirely. `Relevance` fetches exactly // `limit`, so the default path issues the identical query it always has. - let candidate_limit = body.sort.candidate_limit(limit); + let candidate_limit = sort.candidate_limit(limit); + stage::enter(RecallStage::VectorSearch); let t1 = std::time::Instant::now(); let hits = state .db @@ -199,7 +286,7 @@ pub async fn recall( // only distance + created_at, both already on the row, so the over-fetch // costs one wider SQL query instead of 5x the Walrus downloads and SEAL // decrypts. - let hits = super::select_hits_for_sort(hits, body.sort, limit); + let hits = super::select_hits_for_sort(hits, sort, limit); let hit_count = hits.len(); if hits.is_empty() { @@ -207,11 +294,11 @@ pub async fn recall( "recall complete: 0 results (no vector hits) for owner={}", owner ); - return Ok(Json(RecallResponse { + return Ok(RecallResponse { results: vec![], total: 0, dropped_count: 0, - })); + }); } // Hydrate the hits through the storage engine: blob cache -> Walrus @@ -220,6 +307,7 @@ pub async fn recall( // engine owns the // cache/decrypt-batch internals and derives the SEAL credential from // `auth`; per-blob timing breakdowns are visible in its tracing spans. + stage::enter(RecallStage::WalrusDownload); let t2 = std::time::Instant::now(); let hit_refs: Vec<(String, f64)> = hits .iter() @@ -227,7 +315,7 @@ pub async fn recall( .collect(); let (mut hydrated, dropped_count, timings) = state .engine - .fetch_batch(owner, namespace, &hit_refs, &auth) + .fetch_batch(owner, namespace, &hit_refs, auth) .await?; let fetch_ms = t2.elapsed().as_millis(); @@ -260,7 +348,7 @@ pub async fn recall( // this is a no-op and preserves the pgvector cosine order exactly — // pinned by the `default_weights_preserve_input_order` and // `recency_zero_is_short_circuit_no_reorder` tests in services::ranker. - let ranked = state.ranker.rank(hydrated, &weights, chrono::Utc::now()); + let ranked = state.ranker.rank(hydrated, weights, chrono::Utc::now()); let results: Vec = super::recall_results_from_ranked(ranked); let total = results.len(); @@ -293,11 +381,11 @@ pub async fn recall( t0.elapsed().as_millis() ); - Ok(Json(RecallResponse { + Ok(RecallResponse { results, total, dropped_count, - })) + }) } /// POST /api/recall/manual @@ -404,6 +492,33 @@ mod tests { // ── RecallResponse dropped_count serialization ─────────────── + /// The failure report added for accepted-then-failed writes reads the same + /// `remember_jobs.error_msg` the job-status endpoints read, and reaches the + /// same untrusted caller — so it has to be sanitized the same way. It was + /// not, which put the relayer's own wallet address and balance shortfall + /// into every recall response for 24 hours after an infra failure. + /// + /// `infra_wal_balance_failure_hides_relayer_wallet_address` in + /// routes::remember pins this for `GET /api/remember/:job_id`; this pins + /// the same guarantee for the recall path. + #[test] + fn failed_write_report_hides_relayer_wallet_address() { + let raw = "walrus upload failed: Insufficient balance of \ +0x356a26eb9e012a68958082340d4c4116e7f55615ef27affcff209cf0ae544f59::wal::WAL for owner \ +0x8d3c1f0a9b2e4d6c7a5f8e1b0d4c9a2f3e6b7d8c1a0f9e2b3c4d5a6f7e8b9c0d. Required: 64367730, \ +Available: 10708877"; + + let out = crate::routes::remember::sanitize_job_error_for_client("failed", Some(raw.to_string())) + .expect("a failed job keeps an error"); + + // The operator's hot wallet and its shortfall are not the tenant's + // business, and reading them as "top this address up" is the exact + // confusion INFRA_JOB_ERROR_MESSAGE exists to prevent. + assert!(!out.contains("0x8d3c1f0a"), "wallet address leaked: {}", out); + assert!(!out.contains("Available"), "balance leaked: {}", out); + assert!(!out.contains("10708877"), "shortfall leaked: {}", out); + } + #[test] fn recall_response_includes_dropped_count_when_nonzero() { let resp = crate::types::RecallResponse { diff --git a/services/server/src/routes/remember.rs b/services/server/src/routes/remember.rs index 62a99dd31..4085479e0 100644 --- a/services/server/src/routes/remember.rs +++ b/services/server/src/routes/remember.rs @@ -115,7 +115,10 @@ fn redact_hex_addresses(msg: &str) -> String { /// current `status`. Infrastructure failures collapse to fixed copy, chosen by /// whether the job has stopped retrying. Everything else keeps its text with /// addresses redacted. The DB row is untouched. -fn sanitize_job_error_for_client(status: &str, error_msg: Option) -> Option { +pub(crate) fn sanitize_job_error_for_client( + status: &str, + error_msg: Option, +) -> Option { let msg = error_msg?; if crate::jobs::WalletJobError::is_infrastructure_funding_error(&msg) { return Some(if status == "failed" { @@ -226,7 +229,7 @@ fn spawn_prepare_remember_job( ) .await?; - let wallet_index = state.key_pool.next_index().ok_or_else(|| { + let wallet_index = state.key_pool.least_loaded_index().ok_or_else(|| { AppError::Internal( "No Sui keys configured (set SERVER_SUI_PRIVATE_KEYS or SERVER_SUI_PRIVATE_KEY)" .into(), @@ -410,7 +413,7 @@ fn spawn_prepare_bulk_remember_job( for (job_id, namespace, vector, encrypted) in prepared { let wallet_index = state .key_pool - .next_index() + .least_loaded_index() .ok_or_else(|| AppError::Internal("No Sui keys configured".into()))?; let encrypted_b64 = base64::engine::general_purpose::STANDARD.encode(&encrypted); @@ -818,12 +821,35 @@ pub async fn remember( namespace_owned, auth.public_key.clone(), ); + return Ok(( + StatusCode::ACCEPTED, + Json(RememberAcceptedResponse { + job_id: existing_id, + status: "pending".to_string(), + }), + )); } + + // Losing the claim now means a concurrent retry took it, not + // that the TTL blocked us — `failed` rows are re-claimable + // immediately. Report whatever that winner left behind rather + // than asserting "pending" on its behalf: answering with a + // state we did not reach is what told callers a dead job was + // queued. + let actual: String = sqlx::query_scalar( + "SELECT status FROM remember_jobs WHERE id = $1", + ) + .bind(&existing_id) + .fetch_optional(state.db.pool()) + .await + .map_err(|e| AppError::Internal(format!("Failed to re-read job status: {}", e)))? + .unwrap_or_else(|| existing_status.clone()); + return Ok(( StatusCode::ACCEPTED, Json(RememberAcceptedResponse { job_id: existing_id, - status: "pending".to_string(), + status: actual, }), )); } @@ -1020,6 +1046,25 @@ fn should_spawn_after_reset(rows_affected: u64) -> bool { rows_affected == 1 } +/// How long a preparation claim fences other claimants. +/// +/// The TTL exists to stop a second request stealing a claim from a preparation +/// that is still running. It must NOT apply to a job that already reached +/// `failed`: that preparation is over — it either errored on its own or the +/// stale sweeper failed it and cleared its token — so there is no live task to +/// protect, and waiting out the TTL only blocks the retry the caller was just +/// told to make. +/// +/// That was not theoretical. `memwal_remember` tells an agent to send a failed +/// fact again; the derived idempotency key collapses the retry onto the failed +/// row; the claim was refused because it was less than 60s old; and the route +/// answered 202 ACCEPTED anyway. The caller was told the write was queued while +/// nothing whatsoever was running. +/// +/// Letting a `failed` row be re-claimed immediately is safe because fencing is +/// done by the TOKEN, not the clock: a new claim rotates `prepare_claim_token`, +/// and any straggler's own UPDATE is `WHERE ... prepare_claim_token = `, +/// so it matches zero rows and returns before `enqueue_wallet_job`. const PREPARE_CLAIM_TTL_SECS: i64 = 60; async fn claim_remember_preparation( @@ -1028,7 +1073,7 @@ async fn claim_remember_preparation( ) -> Result, AppError> { let token = uuid::Uuid::new_v4().to_string(); let claimed: Option = sqlx::query_scalar( - "UPDATE remember_jobs SET prepare_claimed_at = NOW(), prepare_claim_token = $3, status = CASE WHEN status = 'failed' AND blob_id IS NULL THEN 'pending' ELSE status END, error_msg = CASE WHEN blob_id IS NULL THEN NULL ELSE error_msg END, updated_at = NOW() WHERE id = $1 AND blob_id IS NULL AND status IN ('pending', 'failed') AND (prepare_claimed_at IS NULL OR prepare_claimed_at < NOW() - make_interval(secs => $2)) RETURNING prepare_claim_token", + "UPDATE remember_jobs SET prepare_claimed_at = NOW(), prepare_claim_token = $3, status = CASE WHEN status = 'failed' AND blob_id IS NULL THEN 'pending' ELSE status END, error_msg = CASE WHEN blob_id IS NULL THEN NULL ELSE error_msg END, updated_at = NOW() WHERE id = $1 AND blob_id IS NULL AND status IN ('pending', 'failed') AND (prepare_claimed_at IS NULL OR prepare_claimed_at < NOW() - make_interval(secs => $2) OR status = 'failed') RETURNING prepare_claim_token", ) .bind(job_id) .bind(PREPARE_CLAIM_TTL_SECS) @@ -1292,6 +1337,7 @@ pub async fn remember_bulk( other => other, })?; } + validate_idempotency_key(body.idempotency_key.as_deref())?; let owner = &auth.owner; tracing::info!( @@ -1303,31 +1349,68 @@ pub async fn remember_bulk( let mut job_ids: Vec = Vec::with_capacity(body.items.len()); let mut pending_items: Vec = Vec::with_capacity(body.items.len()); - for item in body.items { - let job_id = uuid::Uuid::new_v4().to_string(); + for (i, item) in body.items.into_iter().enumerate() { + let mut job_id = uuid::Uuid::new_v4().to_string(); + let item_key = body + .idempotency_key + .as_deref() + .map(|key| format!("bulk:{}:{}", key, i)); + let fingerprint = request_fingerprint(&item.text, &item.namespace); - if let Err(e) = sqlx::query( + let inserted = match sqlx::query( // `pending` (not `running`) so a fresh job takes the plain Upload // path; only a retry of an in-flight job (worker-set `running`) // triggers the crash-window reconcile. See the single-remember insert. - "INSERT INTO remember_jobs (id, owner, namespace, status) VALUES ($1, $2, $3, 'pending')", + "INSERT INTO remember_jobs (id, owner, namespace, status, idempotency_key, request_fingerprint) VALUES ($1, $2, $3, 'pending', $4, $5) + ON CONFLICT (owner, idempotency_key) WHERE idempotency_key IS NOT NULL DO NOTHING", ) .bind(&job_id) .bind(owner) .bind(&item.namespace) + .bind(item_key.as_deref()) + .bind(item_key.as_ref().map(|_| fingerprint.as_str())) .execute(state.db.pool()) .await { - crate::alerts::maybe_alert_sqlx_postgres_storage_exhausted( - &state.alerts, - &state.config.sui_network, - &e, - ) - .await; - return Err(AppError::Internal(format!( - "Failed to create bulk job row: {}", - e - ))); + Ok(inserted) => inserted, + Err(e) => { + crate::alerts::maybe_alert_sqlx_postgres_storage_exhausted( + &state.alerts, + &state.config.sui_network, + &e, + ) + .await; + return Err(AppError::Internal(format!( + "Failed to create bulk job row: {}", + e + ))); + } + }; + + if inserted.rows_affected() == 0 { + let existing = match item_key.as_deref() { + Some(key) => find_remember_job_by_key(state.db.pool(), owner, key).await?, + None => None, + }; + let (existing_id, _, _, existing_fingerprint) = existing.ok_or_else(|| { + AppError::Internal("Failed to create bulk job row: insert affected no rows".into()) + })?; + if existing_fingerprint + .as_deref() + .is_some_and(|stored| stored != fingerprint) + { + return Err(AppError::Conflict( + "idempotency_key was already used for a request with different content".into(), + )); + } + if claim_remember_preparation(state.db.pool(), &existing_id) + .await? + .is_none() + { + job_ids.push(existing_id); + continue; + } + job_id = existing_id; } pending_items.push(PendingBulkRememberItem { @@ -1340,13 +1423,15 @@ pub async fn remember_bulk( let total = job_ids.len(); - spawn_prepare_bulk_remember_job( - Arc::clone(&state), - owner.clone(), - auth.account_id.clone(), - auth.public_key.clone(), - pending_items, - ); + if !pending_items.is_empty() { + spawn_prepare_bulk_remember_job( + Arc::clone(&state), + owner.clone(), + auth.account_id.clone(), + auth.public_key.clone(), + pending_items, + ); + } tracing::info!("remember_bulk accepted: {} items owner={}", total, owner,); @@ -1581,6 +1666,69 @@ mod tests { assert_eq!(status, "done"); } + #[tokio::test] + async fn a_failed_job_is_reclaimable_immediately() { + let pool = idem_test_pool().await; + let job_id = uuid::Uuid::new_v4().to_string(); + sqlx::query( + "INSERT INTO remember_jobs (id, owner, namespace, status, prepare_claimed_at, prepare_claim_token) + VALUES ($1, '0xowner', 'ns', 'failed', NOW(), 'stale-token')", + ) + .bind(&job_id) + .execute(&pool) + .await + .unwrap(); + + // Claimed one second ago — well inside PREPARE_CLAIM_TTL_SECS. + let claimed = claim_remember_preparation(&pool, &job_id).await.unwrap(); + assert!( + claimed.is_some(), + "a failed job must be re-claimable without waiting out the TTL", + ); + + // And the retry is actually live, not merely reported as such. + let status: String = sqlx::query_scalar("SELECT status FROM remember_jobs WHERE id = $1") + .bind(&job_id) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(status, "pending"); + + // Rotating the token is what fences the previous attempt, so the old + // one can no longer write its preparation or reach the wallet queue. + let straggler = sqlx::query( + "UPDATE remember_jobs SET preparation_encrypted_b64 = 'late' + WHERE id = $1 AND prepare_claim_token = $2", + ) + .bind(&job_id) + .bind("stale-token") + .execute(&pool) + .await + .unwrap(); + assert_eq!(straggler.rows_affected(), 0, "the old claim must be fenced out"); + } + + /// The TTL still does its real job: a claim on a job that is genuinely + /// mid-preparation (`pending`, claimed just now) must not be stolen. + #[tokio::test] + async fn a_live_pending_claim_is_still_fenced() { + let pool = idem_test_pool().await; + let job_id = uuid::Uuid::new_v4().to_string(); + sqlx::query( + "INSERT INTO remember_jobs (id, owner, namespace, status, prepare_claimed_at, prepare_claim_token) + VALUES ($1, '0xowner', 'ns', 'pending', NOW(), 'live-token')", + ) + .bind(&job_id) + .execute(&pool) + .await + .unwrap(); + + assert!( + claim_remember_preparation(&pool, &job_id).await.unwrap().is_none(), + "a preparation still running must keep its claim", + ); + } + #[tokio::test] async fn initial_preparation_claim_is_fenced_after_reclaim() { let pool = idem_test_pool().await; @@ -1848,6 +1996,111 @@ mod tests { .await; } + #[tokio::test] + async fn bulk_derived_key_collapses_a_retried_item() { + let pool = idem_test_pool().await; + let owner = format!("0xowner-{}", uuid::Uuid::new_v4()); + let item_key = "bulk:batch-1:0"; + let fingerprint = request_fingerprint("bulk item text", "ns"); + let job_id = format!("remember-job-{}", uuid::Uuid::new_v4()); + + let insert = "INSERT INTO remember_jobs (id, owner, namespace, status, idempotency_key, request_fingerprint) VALUES ($1, $2, 'ns', 'pending', $3, $4) + ON CONFLICT (owner, idempotency_key) WHERE idempotency_key IS NOT NULL DO NOTHING"; + + sqlx::query(insert) + .bind(&job_id) + .bind(&owner) + .bind(item_key) + .bind(&fingerprint) + .execute(&pool) + .await + .unwrap(); + + let retry = sqlx::query(insert) + .bind(format!("remember-job-{}", uuid::Uuid::new_v4())) + .bind(&owner) + .bind(item_key) + .bind(&fingerprint) + .execute(&pool) + .await + .unwrap(); + assert_eq!( + retry.rows_affected(), + 0, + "retried bulk item must not mint a second job" + ); + + let found = find_remember_job_by_key(&pool, &owner, item_key) + .await + .unwrap() + .unwrap(); + assert_eq!(found.0, job_id); + assert_eq!(found.3.as_deref(), Some(fingerprint.as_str())); + + let _ = sqlx::query("DELETE FROM remember_jobs WHERE owner = $1") + .bind(&owner) + .execute(&pool) + .await; + } + + #[tokio::test] + async fn bulk_retry_redrives_an_orphaned_pending_item() { + let pool = idem_test_pool().await; + let owner = format!("0xowner-{}", uuid::Uuid::new_v4()); + let item_key = "bulk:batch-orphan:0"; + let fingerprint = request_fingerprint("bulk item text", "ns"); + let job_id = format!("remember-job-{}", uuid::Uuid::new_v4()); + + sqlx::query( + "INSERT INTO remember_jobs (id, owner, namespace, status, idempotency_key, request_fingerprint) VALUES ($1, $2, 'ns', 'pending', $3, $4)", + ) + .bind(&job_id) + .bind(&owner) + .bind(item_key) + .bind(&fingerprint) + .execute(&pool) + .await + .unwrap(); + + let claimed_at: Option> = + sqlx::query_scalar("SELECT prepare_claimed_at FROM remember_jobs WHERE id = $1") + .bind(&job_id) + .fetch_one(&pool) + .await + .unwrap(); + assert!( + claimed_at.is_none(), + "fixture must start in the orphan state" + ); + + let existing = find_remember_job_by_key(&pool, &owner, item_key) + .await + .unwrap() + .unwrap(); + assert_eq!(existing.0, job_id); + + let claim = claim_remember_preparation(&pool, &existing.0) + .await + .unwrap(); + assert!( + claim.is_some(), + "an orphaned pending bulk item must be re-driven, not collapsed" + ); + + assert!( + claim_remember_preparation(&pool, &existing.0) + .await + .unwrap() + .is_none(), + "an item already being prepared must collapse onto the existing job" + ); + + let _ = sqlx::query("DELETE FROM remember_jobs WHERE owner = $1") + .bind(&owner) + .execute(&pool) + .await; + } + #[tokio::test] async fn keyless_writes_are_never_deduped() { let pool = idem_test_pool().await; diff --git a/services/server/src/services/embedder.rs b/services/server/src/services/embedder.rs index 90252e6a3..f7956a51c 100644 --- a/services/server/src/services/embedder.rs +++ b/services/server/src/services/embedder.rs @@ -27,8 +27,26 @@ pub const EMBEDDING_MODEL: &str = "openai/text-embedding-3-small"; pub const EMBEDDING_DIMS: usize = 1536; /// 16384 = 8192 tokens × 2 chars/token under cl100k; do not use admin 64KiB. +/// +/// A limit of [`EMBEDDING_MODEL`]'s context window, so it applies only when a +/// provider key is set. The key-less fallback hashes locally with no limit. pub(crate) const MAX_EMBED_INPUT_BYTES: usize = 16384; +/// Whether a provider 400 complains about input length; any other 400 is a +/// server-side problem, not the caller's. +fn is_context_length_error(body: &str) -> bool { + let body = body.to_ascii_lowercase(); + [ + "context_length_exceeded", + "maximum context length", + "reduce the length", + "too many tokens", + "string too long", + ] + .iter() + .any(|needle| body.contains(needle)) +} + fn reject_oversized_embed_input(text: &str) -> Result<(), AppError> { if text.len() > MAX_EMBED_INPUT_BYTES { return Err(AppError::BadRequest(format!( @@ -67,114 +85,130 @@ impl OpenAiEmbedder { impl Embedder for OpenAiEmbedder { #[tracing::instrument(name = "embedder.embed", skip_all, fields(text_len = text.len()))] async fn embed(&self, text: &str) -> Result, AppError> { - reject_oversized_embed_input(text)?; - match &self.config.openai_api_key { - Some(api_key) => { - // Real embedding via OpenRouter/OpenAI-compatible API - let url = format!("{}/embeddings", self.config.openai_api_base); - - let started = std::time::Instant::now(); - let resp = self - .http_client - .post(&url) - .header("Authorization", format!("Bearer {}", api_key)) - .header("Content-Type", "application/json") - .json(&EmbeddingApiRequest { - model: EMBEDDING_MODEL.to_string(), - input: text.to_string(), - }) - .send() - .await - .map_err(|e| { - crate::observability::observe_external( - "openai", - "embeddings", - "transport_error", - started.elapsed(), - ); - AppError::Internal(format!("Embedding API request failed: {}", e)) - })?; - let status_label = resp.status().as_u16().to_string(); - crate::observability::observe_external( - "openai", - "embeddings", - &status_label, - started.elapsed(), - ); + embed_text( + &self.http_client, + self.config.openai_api_key.as_deref(), + &self.config.openai_api_base, + text, + ) + .await + } +} - if !resp.status().is_success() { - let status = resp.status(); - let body = resp.text().await.unwrap_or_default(); - if crate::services::extractor::is_upstream_status_transient(status) { - return Err(AppError::UpstreamUnavailable(format!( - "Embedding API upstream error ({}): {}", - status, body - ))); - } - if status == reqwest::StatusCode::BAD_REQUEST { - tracing::warn!(%status, body, "embedding API returned 400"); - return Err(AppError::BadRequest( - "embedding input exceeds the model context limit".into(), - )); - } - return Err(AppError::Internal(format!( - "Embedding API error ({}): {}", - status, body - ))); - } +/// The embedding call itself, taking only what it reads from `Config` so the +/// tests can drive it against a local stand-in provider. +async fn embed_text( + http_client: &reqwest::Client, + api_key: Option<&str>, + api_base: &str, + text: &str, +) -> Result, AppError> { + match api_key { + Some(api_key) => { + reject_oversized_embed_input(text)?; + // Real embedding via OpenRouter/OpenAI-compatible API + let url = format!("{}/embeddings", api_base); - // same pattern as the extractor — capture body - // as text first so we can (1) treat transport-level - // failures as transient, and (2) detect OpenRouter - // error envelopes wrapped in HTTP 200. Both route to - // `AppError::UpstreamUnavailable` (HTTP 503) so the - // SDK / harness retry policy can recover. See - // `extractor::parse_openrouter_error_envelope`. - let body = resp.text().await.map_err(|e| { - AppError::UpstreamUnavailable(format!( - "Failed to read embedding response body: {}", - e - )) + let started = std::time::Instant::now(); + let resp = http_client + .post(&url) + .header("Authorization", format!("Bearer {}", api_key)) + .header("Content-Type", "application/json") + .json(&EmbeddingApiRequest { + model: EMBEDDING_MODEL.to_string(), + input: text.to_string(), + }) + .send() + .await + .map_err(|e| { + crate::observability::observe_external( + "openai", + "embeddings", + "transport_error", + started.elapsed(), + ); + AppError::Internal(format!("Embedding API request failed: {}", e)) })?; + let status_label = resp.status().as_u16().to_string(); + crate::observability::observe_external( + "openai", + "embeddings", + &status_label, + started.elapsed(), + ); - if let Some(envelope) = - crate::services::extractor::parse_openrouter_error_envelope(&body) - { + if !resp.status().is_success() { + let status = resp.status(); + let body = resp.text().await.unwrap_or_default(); + if crate::services::extractor::is_upstream_status_transient(status) { return Err(AppError::UpstreamUnavailable(format!( - "OpenRouter upstream error (code={}): {}", - envelope.code, envelope.message + "Embedding API upstream error ({}): {}", + status, body ))); } + if status == reqwest::StatusCode::BAD_REQUEST && is_context_length_error(&body) { + tracing::warn!(%status, body, "embedding API rejected the input length"); + return Err(AppError::BadRequest( + "embedding input exceeds the model context limit".into(), + )); + } + return Err(AppError::Internal(format!( + "Embedding API error ({}): {}", + status, body + ))); + } - let api_resp: EmbeddingApiResponse = serde_json::from_str(&body).map_err(|e| { - AppError::Internal(format!("Failed to parse embedding response: {}", e)) - })?; + // same pattern as the extractor — capture body + // as text first so we can (1) treat transport-level + // failures as transient, and (2) detect OpenRouter + // error envelopes wrapped in HTTP 200. Both route to + // `AppError::UpstreamUnavailable` (HTTP 503) so the + // SDK / harness retry policy can recover. See + // `extractor::parse_openrouter_error_envelope`. + let body = resp.text().await.map_err(|e| { + AppError::UpstreamUnavailable(format!( + "Failed to read embedding response body: {}", + e + )) + })?; - let vector = api_resp - .data - .into_iter() - .next() - .ok_or_else(|| AppError::Internal("Embedding API returned no data".into()))? - .embedding; - Ok(vector) - } - None => { - // Mock embedding (deterministic hash-based) — for keyless dev - tracing::warn!(" → Using MOCK embedding (no OPENAI_API_KEY set)"); - use sha2::Digest; - let hash = sha2::Sha256::digest(text.as_bytes()); - let mock_vector: Vec = hash - .iter() - .cycle() - .take(EMBEDDING_DIMS) - .enumerate() - .map(|(i, &b)| { - let val = (b as f32 / 255.0) * 2.0 - 1.0; - val * (1.0 + (i as f32 * 0.001).sin()) - }) - .collect(); - Ok(mock_vector) + if let Some(envelope) = + crate::services::extractor::parse_openrouter_error_envelope(&body) + { + return Err(AppError::UpstreamUnavailable(format!( + "OpenRouter upstream error (code={}): {}", + envelope.code, envelope.message + ))); } + + let api_resp: EmbeddingApiResponse = serde_json::from_str(&body).map_err(|e| { + AppError::Internal(format!("Failed to parse embedding response: {}", e)) + })?; + + let vector = api_resp + .data + .into_iter() + .next() + .ok_or_else(|| AppError::Internal("Embedding API returned no data".into()))? + .embedding; + Ok(vector) + } + None => { + // Mock embedding (deterministic hash-based) — for keyless dev + tracing::warn!(" → Using MOCK embedding (no OPENAI_API_KEY set)"); + use sha2::Digest; + let hash = sha2::Sha256::digest(text.as_bytes()); + let mock_vector: Vec = hash + .iter() + .cycle() + .take(EMBEDDING_DIMS) + .enumerate() + .map(|(i, &b)| { + let val = (b as f32 / 255.0) * 2.0 - 1.0; + val * (1.0 + (i as f32 * 0.001).sin()) + }) + .collect(); + Ok(mock_vector) } } } @@ -251,4 +285,92 @@ mod tests { "input is over the embedding input limit", ); } + + // ── embed_text: size limit and upstream 400 mapping ────────────── + + /// Stand-in embeddings provider answering every `POST /embeddings` with + /// `status` and `body`. Returns its base URL and a call counter. + async fn stub_provider( + status: u16, + body: &'static str, + ) -> (String, std::sync::Arc) { + use std::sync::atomic::{AtomicUsize, Ordering}; + let calls = std::sync::Arc::new(AtomicUsize::new(0)); + let counter = std::sync::Arc::clone(&calls); + let app = axum::Router::new().route( + "/embeddings", + axum::routing::post(move || { + let counter = std::sync::Arc::clone(&counter); + async move { + counter.fetch_add(1, Ordering::SeqCst); + (axum::http::StatusCode::from_u16(status).unwrap(), body) + } + }), + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + (format!("http://{addr}"), calls) + } + + #[tokio::test] + async fn without_a_key_input_over_the_model_limit_still_embeds() { + let text = "a".repeat(100 * 1024); + let vector = super::embed_text( + &reqwest::Client::new(), + None, + "http://unused.invalid", + &text, + ) + .await + .expect("key-less embedding has no input limit"); + assert_eq!(vector.len(), super::EMBEDDING_DIMS); + } + + #[tokio::test] + async fn with_a_key_oversized_input_is_rejected_before_calling_the_provider() { + let (base, calls) = stub_provider(200, r#"{"data":[{"embedding":[0.0]}]}"#).await; + let text = "q".repeat(super::MAX_EMBED_INPUT_BYTES + 1); + let result = super::embed_text(&reqwest::Client::new(), Some("key"), &base, &text).await; + assert!( + matches!(result, Err(crate::types::AppError::BadRequest(_))), + "got {result:?}" + ); + assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn a_provider_400_about_input_length_is_the_callers_bad_request() { + for body in [ + r#"{"error":{"code":"context_length_exceeded","message":"..."}}"#, + r#"{"error":{"message":"This model's maximum context length is 8192 tokens"}}"#, + r#"{"error":{"message":"Please reduce the length of the messages."}}"#, + r#"{"error":{"message":"Too many tokens in input"}}"#, + r#"{"error":{"message":"String too long. Expected a string with maximum length 8192"}}"#, + ] { + let (base, _) = stub_provider(400, body).await; + let result = super::embed_text(&reqwest::Client::new(), Some("key"), &base, "q").await; + assert!( + matches!(result, Err(crate::types::AppError::BadRequest(_))), + "{body} → {result:?}" + ); + } + } + + #[tokio::test] + async fn any_other_provider_400_is_an_internal_error() { + for body in [ + r#"{"error":{"message":"The model `text-embedding-9` does not exist"}}"#, + r#"{"error":{"code":"invalid_api_key","message":"Incorrect API key provided"}}"#, + r#"{"error":{"message":"Unrecognized request argument supplied: dimensions"}}"#, + "", + ] { + let (base, _) = stub_provider(400, body).await; + let result = super::embed_text(&reqwest::Client::new(), Some("key"), &base, "q").await; + assert!( + matches!(result, Err(crate::types::AppError::Internal(_))), + "{body} → {result:?}" + ); + } + } } diff --git a/services/server/src/sidecar_saturation.rs b/services/server/src/sidecar_saturation.rs new file mode 100644 index 000000000..636bf1c45 --- /dev/null +++ b/services/server/src/sidecar_saturation.rs @@ -0,0 +1,176 @@ +//! Decision logic for the sidecar upload-queue saturation monitor. +//! +//! The polling loop lives in `main.rs`; parsing and the consecutive-check +//! state live here so the alert decision can be tested without a sidecar. + +use serde_json::Value; + +/// Sidecar route serving the upload-queue counters (`registerUploadMetricsRoute` +/// in `scripts/sidecar/routes/health.ts`). +pub const UPLOAD_METRICS_PATH: &str = "/metrics/uploads"; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct UploadQueueSample { + pub queued: u64, + pub active: u64, + pub global_capacity: u64, +} + +/// Reads the counters from a `/metrics/uploads` body. A missing or +/// non-integer field is `Err` with that field's JSON pointer; it must never be +/// read as 0. +pub fn parse_upload_metrics(body: &Value) -> Result { + let field = + |pointer: &'static str| body.pointer(pointer).and_then(Value::as_u64).ok_or(pointer); + Ok(UploadQueueSample { + queued: field("/queuedWalrusUploads")?, + active: field("/activeWalrusUploads")?, + global_capacity: field("/walrusUploadLimits/globalCapacity")?, + }) +} + +#[derive(Debug, PartialEq, Eq)] +pub enum QueueCheck { + /// At or below the threshold, with no alerting streak to end. + Clear, + /// Back at or below the threshold after a streak that alerted. + Drained, + /// Above the threshold; `alert` once the streak is long enough. + Saturated { consecutive: u32, alert: bool }, +} + +pub struct SaturationTracker { + threshold: u64, + alert_after: u32, + consecutive: u32, +} + +impl SaturationTracker { + pub fn new(threshold: u64, alert_after: u32) -> Self { + Self { + threshold, + alert_after, + consecutive: 0, + } + } + + pub fn observe(&mut self, queued: u64) -> QueueCheck { + if queued > self.threshold { + self.consecutive = self.consecutive.saturating_add(1); + return QueueCheck::Saturated { + consecutive: self.consecutive, + alert: self.consecutive >= self.alert_after, + }; + } + let alerted = self.consecutive >= self.alert_after; + self.consecutive = 0; + if alerted { + QueueCheck::Drained + } else { + QueueCheck::Clear + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn parses_the_upload_metrics_body() { + let body = json!({ + "activeWalrusUploads": 5, + "queuedWalrusUploads": 118, + "walrusUploadLimits": { + "globalCapacity": 5, + "perWalletCapacity": 1, + "acquireTimeoutMs": 120000 + } + }); + assert_eq!( + parse_upload_metrics(&body), + Ok(UploadQueueSample { + queued: 118, + active: 5, + global_capacity: 5, + }) + ); + } + + #[test] + fn a_liveness_body_is_an_error_not_an_empty_queue() { + let body = json!({ "status": "ok", "uptimeMs": 86_400_000 }); + assert_eq!(parse_upload_metrics(&body), Err("/queuedWalrusUploads")); + } + + #[test] + fn a_missing_nested_capacity_is_an_error() { + let body = json!({ "activeWalrusUploads": 5, "queuedWalrusUploads": 118 }); + assert_eq!( + parse_upload_metrics(&body), + Err("/walrusUploadLimits/globalCapacity") + ); + } + + #[test] + fn a_non_integer_counter_is_an_error() { + let body = json!({ + "activeWalrusUploads": 5, + "queuedWalrusUploads": "118", + "walrusUploadLimits": { "globalCapacity": 5 } + }); + assert_eq!(parse_upload_metrics(&body), Err("/queuedWalrusUploads")); + } + + #[test] + fn alerts_once_the_queue_stays_above_threshold_for_the_configured_checks() { + let mut tracker = SaturationTracker::new(20, 3); + let checks: Vec = [21, 40, 120, 118].map(|q| tracker.observe(q)).into(); + assert_eq!( + checks, + vec![ + QueueCheck::Saturated { + consecutive: 1, + alert: false + }, + QueueCheck::Saturated { + consecutive: 2, + alert: false + }, + QueueCheck::Saturated { + consecutive: 3, + alert: true + }, + QueueCheck::Saturated { + consecutive: 4, + alert: true + }, + ] + ); + } + + #[test] + fn a_queue_at_the_threshold_resets_the_streak() { + let mut tracker = SaturationTracker::new(20, 3); + tracker.observe(21); + tracker.observe(21); + assert_eq!(tracker.observe(20), QueueCheck::Clear); + assert_eq!( + tracker.observe(21), + QueueCheck::Saturated { + consecutive: 1, + alert: false + } + ); + } + + #[test] + fn draining_is_reported_only_after_a_streak_that_alerted() { + let mut tracker = SaturationTracker::new(20, 2); + tracker.observe(50); + tracker.observe(50); + assert_eq!(tracker.observe(0), QueueCheck::Drained); + assert_eq!(tracker.observe(0), QueueCheck::Clear); + } +} diff --git a/services/server/src/storage/db.rs b/services/server/src/storage/db.rs index 90e8cd18c..2382e5e1d 100644 --- a/services/server/src/storage/db.rs +++ b/services/server/src/storage/db.rs @@ -2,7 +2,7 @@ use std::sync::Arc; use pgvector::Vector; use sqlx::postgres::PgPoolOptions; -use sqlx::PgPool; +use sqlx::{PgConnection, PgPool}; use crate::alerts::AlertManager; use crate::types::{AppError, SearchHit}; @@ -11,11 +11,168 @@ use crate::types::{AppError, SearchHit}; /// background sweep. Keep a single constant so the two cannot drift. pub const TOMBSTONE_RETENTION: chrono::Duration = chrono::Duration::days(30); +/// One migration file, paired with the name the pipeline reports it by. +type Migration = (&'static str, &'static str); + +/// Pairs a migration's file name with its embedded contents so the two +/// cannot drift apart. +macro_rules! migration { + ($file:literal) => { + ($file, include_str!(concat!("../../migrations/", $file))) + }; +} + +/// The migration pipeline, split at the two points where `VectorDb::new` +/// has to run Rust in between files. +/// +/// This stays an explicit, hand-ordered list rather than a +/// `sqlx::migrate!` directory scan, for three reasons that all still +/// hold: +/// +/// 1. `014_storage_reservations.sql` must run *before* +/// `014_memory_read_api_columns.sql` — the reverse of their +/// alphabetical order. +/// 2. Two files share the version number `014`, which `sqlx::migrate!` +/// rejects outright. +/// 3. `backfill_updated_at` and `recover_invalid_concurrent_indexes` are +/// Rust steps that have to land between specific files. +/// +/// What is no longer manual is *completeness*: every `.sql` file in +/// `services/server/migrations` must appear in one of these three +/// slices, and `every_migration_file_is_wired_into_the_pipeline` fails +/// the test suite if one does not. The original 021 ADD reached +/// origin/dev as a file but never entered this list, which is why the +/// check exists. This 021 is the DROP of that column, for environments +/// that ran a PR-branch build that did wire the ADD. +const MIGRATIONS_BEFORE_BACKFILL: &[Migration] = &[ + migration!("001_init.sql"), + migration!("002_add_namespace.sql"), + migration!("003_rate_limiter.sql"), + migration!("004_delegate_key_cache_expires.sql"), + migration!("005_remember_jobs.sql"), + // composite index on (owner, status, updated_at DESC) for bulk poll + migration!("006_bulk_remember.sql"), + // collapse per-wallet Apalis queues to a single `wallet_jobs` queue. + // Equivocation locks are no longer a practical concern on Sui (per + // Will Bradley, Mysten, 2026-05-12); concurrent workers on one wallet + // + retry handling is sufficient. + migration!("007_collapse_wallet_queues.sql"), + // nullable `plaintext` column for benchmark-mode storage + // (PlaintextEngine). NULL for all production rows — additive. + // Renumbered from 007 -> 008 during rebase onto dev to avoid collision + // with the wallet-queue collapse migration. + migration!("008_benchmark_plaintext.sql"), + // importance signal column on vector_entries. + migration!("009_importance_signal.sql"), + // Permanent restore-failure negative cache (GH #501 / WALM-299). + migration!("010_restore_failed_blobs.sql"), + // MCP OAuth 2.1 (Claude custom connectors): client registry, + // server-custodied delegate keys, and authorization state. + migration!("011_mcp_oauth.sql"), + // Durable idempotency, preparation, and paid-upload recovery state. + migration!("012_remember_write_idempotency.sql"), + // Build the owner/key uniqueness constraint without blocking writes. + migration!("013_remember_write_idempotency_index.sql"), + // per-owner storage quota reservations. Makes quota admission atomic + // with the eventual insert (GH #532 / WALM-359). + migration!("014_storage_reservations.sql"), + // owner-scoped read API: updated_at cursor column + agent_id/package_id. + // Split across 014-019 (see each file's header, and + // backfill_updated_at's / recover_invalid_concurrent_indexes's doc + // comments below) to avoid holding ACCESS EXCLUSIVE across the + // full-table backfill or the index build. + migration!("014_memory_read_api_columns.sql"), +]; + +/// Applied after `backfill_updated_at`: 015 validates NOT NULL and will +/// error if any `updated_at` row is still NULL. +const MIGRATIONS_AFTER_BACKFILL: &[Migration] = + &[migration!("015_memory_read_api_updated_at_not_null.sql")]; + +/// Applied after `recover_invalid_concurrent_indexes`, which must precede +/// every `CREATE INDEX CONCURRENTLY IF NOT EXISTS` in this slice — 016, +/// 018 and 022. Each of those would otherwise silently no-op forever +/// against a permanently INVALID index left behind by an interrupted +/// build. Every such index is registered in +/// `CONCURRENTLY_BUILT_INDEXES`; a new one added to this slice without +/// being registered fails +/// `every_concurrently_built_index_is_registered_for_recovery`. +const MIGRATIONS_AFTER_INDEX_RECOVERY: &[Migration] = &[ + // keyset-pagination index for the memories listing endpoint. + // Must stay in its own file/transaction — see 016's header comment. + migration!("016_memory_read_api_index.sql"), + // per-memory expiry columns. + migration!("017_memory_expiry_columns.sql"), + // index on expiry_synced_at so the periodic expiry refresh sweep + // doesn't full-scan vector_entries every tick. Must stay in its own + // file/transaction — see 018's header comment. + migration!("018_memory_expiry_synced_at_index.sql"), + // Finalizes updated_at NOT NULL cheaply using the validated CHECK + // constraint 015 set up — see 019's header. + migration!("019_memory_read_api_updated_at_set_not_null.sql"), + migration!("020_read_api_followups.sql"), + // Drops failure_reported_at if a preview deploy created it. IF EXISTS, + // so a database that never had the column (CI, origin/dev-only) is fine. + migration!("021_drop_failed_write_report_ack.sql"), + // Partial index for /health recent_write_outcomes. CONCURRENTLY, own + // file — see 022's header. + migration!("022_remember_jobs_recent_outcomes.sql"), +]; + +/// Every migration the pipeline applies, in the order it applies them. +#[cfg(test)] +fn all_migrations() -> impl Iterator { + MIGRATIONS_BEFORE_BACKFILL + .iter() + .chain(MIGRATIONS_AFTER_BACKFILL) + .chain(MIGRATIONS_AFTER_INDEX_RECOVERY) +} + +/// Applies one slice of the pipeline, naming the file that failed. +async fn run_migrations(pool: &PgPool, migrations: &[Migration]) -> Result<(), AppError> { + for (name, sql) in migrations { + sqlx::raw_sql(sql) + .execute(pool) + .await + .map_err(|e| AppError::Internal(format!("Failed to run migration {}: {}", name, e)))?; + } + Ok(()) +} + +/// Same as `run_migrations`, but on a held session — used so recover + +/// `CREATE INDEX CONCURRENTLY` share the advisory lock that serializes +/// replica boots. +async fn run_migrations_on( + conn: &mut PgConnection, + migrations: &[Migration], +) -> Result<(), AppError> { + for (name, sql) in migrations { + sqlx::raw_sql(sql) + .execute(&mut *conn) + .await + .map_err(|e| AppError::Internal(format!("Failed to run migration {}: {}", name, e)))?; + } + Ok(()) +} + pub struct VectorDb { pool: PgPool, storage_alerts: Option<(Arc, String)>, } +/// Serialises `VectorDb::new()` across the test binary. +/// +/// The migration chain is idempotent per statement but NOT safe to run +/// concurrently: `CREATE TABLE IF NOT EXISTS` is not atomic in Postgres, so two +/// tests that build a db at the same moment race inside migration 011 and one +/// loses with `duplicate key value violates unique constraint +/// "pg_type_typname_nsp_index"` on `mcp_oauth_clients`. Tests then fail on an +/// unreachable database rather than on anything they assert. `jobs::tests` +/// already guards its pool this way; these modules did not, which left the race +/// latent until a test was added. +#[cfg(test)] +static DB_SETUP_LOCK: std::sync::OnceLock> = std::sync::OnceLock::new(); + impl VectorDb { pub fn with_storage_alerts(self, alerts: Arc, sui_network: String) -> Self { Self { @@ -53,6 +210,151 @@ mod tests { static VECTOR_SCHEMA_SETUP_LOCK: OnceLock> = OnceLock::new(); + /// Looks a migration up in the pipeline by file name. + /// + /// Test helpers below build deliberately partial schemas — only the + /// tables a given module touches — so they cannot just replay the + /// whole pipeline. Going through this lookup still keeps them from + /// drifting: a renamed or deleted migration panics here by name + /// instead of failing later as a missing column. + fn migration_sql(name: &str) -> &'static str { + super::all_migrations() + .find(|(n, _)| *n == name) + .map(|(_, sql)| *sql) + .unwrap_or_else(|| panic!("migration {name} is not wired into the pipeline")) + } + + /// Every `.sql` file in `services/server/migrations` must be wired + /// into the pipeline. + #[test] + fn every_migration_file_is_wired_into_the_pipeline() { + use std::collections::BTreeSet; + + let dir = concat!(env!("CARGO_MANIFEST_DIR"), "/migrations"); + let on_disk: BTreeSet = std::fs::read_dir(dir) + .expect("migrations directory should be readable") + .map(|entry| { + entry + .expect("migrations directory entry should be readable") + .file_name() + .to_string_lossy() + .into_owned() + }) + .filter(|name| name.ends_with(".sql")) + .collect(); + + let wired: BTreeSet = super::all_migrations() + .map(|(name, _)| (*name).to_owned()) + .collect(); + + let unwired: Vec<&String> = on_disk.difference(&wired).collect(); + assert!( + unwired.is_empty(), + "migration file(s) exist on disk but are not applied by VectorDb::new: {unwired:?}. \ + Add them to MIGRATIONS_BEFORE_BACKFILL / _AFTER_BACKFILL / \ + _AFTER_INDEX_RECOVERY in the position the pipeline needs." + ); + + let missing: Vec<&String> = wired.difference(&on_disk).collect(); + assert!( + missing.is_empty(), + "pipeline references migration file(s) that no longer exist: {missing:?}" + ); + } + + /// Every index a migration in `MIGRATIONS_AFTER_INDEX_RECOVERY` builds + /// CONCURRENTLY must be registered in `CONCURRENTLY_BUILT_INDEXES`. + /// + /// Forgetting one fails silently, which is why it needs a test rather + /// than a convention: an interrupted build leaves the index INVALID, + /// `IF NOT EXISTS` then no-ops its migration on every later boot, and + /// the query it served quietly falls back to a scan. 022 is the case + /// that motivates the check -- it keeps `/health`'s write-outcome + /// probe cheap, so losing it disables a silent-failure detector, + /// silently. + #[test] + fn every_concurrently_built_index_is_registered_for_recovery() { + use std::collections::BTreeSet; + + // `CREATE [UNIQUE] INDEX CONCURRENTLY [IF NOT EXISTS] `, + // whether the name sits with the keywords or the `ON` clause wraps + // to the next line. + fn index_name(statement: &str) -> Option { + let tokens: Vec<&str> = statement.split_whitespace().collect(); + let at = tokens.iter().position(|t| *t == "concurrently")?; + let mut rest = &tokens[at + 1..]; + for keyword in ["if", "not", "exists"] { + if rest.first() == Some(&keyword) { + rest = &rest[1..]; + } + } + rest.first() + .map(|name| { + name.trim_matches(|c: char| !c.is_alphanumeric() && c != '_') + .to_owned() + }) + .filter(|name| !name.is_empty()) + } + + let mut built: BTreeSet = BTreeSet::new(); + for (file, sql) in super::MIGRATIONS_AFTER_INDEX_RECOVERY { + for statement in sql.split(';') { + // Drop comment lines so a file header that merely mentions + // CREATE INDEX CONCURRENTLY is not read as one. + let code = statement + .lines() + .filter(|line| !line.trim_start().starts_with("--")) + .collect::>() + .join(" ") + .to_ascii_lowercase(); + if !(code.contains("create") + && code.contains("index") + && code.contains("concurrently")) + { + continue; + } + built.insert( + index_name(&code) + .unwrap_or_else(|| panic!("could not read an index name out of {file}")), + ); + } + } + + let registered: BTreeSet = super::CONCURRENTLY_BUILT_INDEXES + .iter() + .map(|(name, _)| (*name).to_owned()) + .collect(); + + let unregistered: Vec<&String> = built.difference(®istered).collect(); + assert!( + unregistered.is_empty(), + "index(es) built CONCURRENTLY but not registered for invalid-index \ + recovery: {unregistered:?}. Add them to CONCURRENTLY_BUILT_INDEXES, or an \ + interrupted build leaves them INVALID and their migration no-ops forever." + ); + + let stale: Vec<&String> = registered.difference(&built).collect(); + assert!( + stale.is_empty(), + "CONCURRENTLY_BUILT_INDEXES names index(es) that no migration in \ + MIGRATIONS_AFTER_INDEX_RECOVERY builds: {stale:?}" + ); + } + + #[test] + fn drop_invalid_concurrent_index_uses_if_exists() { + let sql = super::drop_invalid_concurrent_index_sql("remember_jobs_recent_outcomes_idx"); + assert!(sql.to_ascii_uppercase().contains("CONCURRENTLY"), "{sql}"); + assert!(sql.to_ascii_uppercase().contains("IF EXISTS"), "{sql}"); + assert!(sql.contains("remember_jobs_recent_outcomes_idx"), "{sql}"); + } + + #[test] + fn invalid_index_is_not_dropped_while_a_build_is_in_progress() { + assert!(!super::should_drop_invalid_concurrent_index(true)); + assert!(super::should_drop_invalid_concurrent_index(false)); + } + fn test_database_url() -> Option { std::env::var("DATABASE_URL").ok() } @@ -74,13 +376,13 @@ mod tests { .lock() .await; for migration in [ - include_str!("../../migrations/001_init.sql"), - include_str!("../../migrations/002_add_namespace.sql"), - include_str!("../../migrations/003_rate_limiter.sql"), - include_str!("../../migrations/008_benchmark_plaintext.sql"), - include_str!("../../migrations/009_importance_signal.sql"), - include_str!("../../migrations/010_restore_failed_blobs.sql"), - include_str!("../../migrations/014_memory_read_api_columns.sql"), + migration_sql("001_init.sql"), + migration_sql("002_add_namespace.sql"), + migration_sql("003_rate_limiter.sql"), + migration_sql("008_benchmark_plaintext.sql"), + migration_sql("009_importance_signal.sql"), + migration_sql("010_restore_failed_blobs.sql"), + migration_sql("014_memory_read_api_columns.sql"), ] { sqlx::raw_sql(migration).execute(&pool).await.unwrap(); } @@ -91,23 +393,21 @@ mod tests { // CONCURRENTLY IF NOT EXISTS. super::backfill_updated_at(&pool).await.unwrap(); - sqlx::raw_sql(include_str!( - "../../migrations/015_memory_read_api_updated_at_not_null.sql" - )) - .execute(&pool) - .await - .unwrap(); + sqlx::raw_sql(migration_sql("015_memory_read_api_updated_at_not_null.sql")) + .execute(&pool) + .await + .unwrap(); - super::recover_invalid_pagination_index(&pool) + super::recover_invalid_concurrent_indexes(&pool) .await .unwrap(); for migration in [ - include_str!("../../migrations/016_memory_read_api_index.sql"), - include_str!("../../migrations/017_memory_expiry_columns.sql"), - include_str!("../../migrations/018_memory_expiry_synced_at_index.sql"), - include_str!("../../migrations/019_memory_read_api_updated_at_set_not_null.sql"), - include_str!("../../migrations/020_read_api_followups.sql"), + migration_sql("016_memory_read_api_index.sql"), + migration_sql("017_memory_expiry_columns.sql"), + migration_sql("018_memory_expiry_synced_at_index.sql"), + migration_sql("019_memory_read_api_updated_at_set_not_null.sql"), + migration_sql("020_read_api_followups.sql"), ] { sqlx::raw_sql(migration).execute(&pool).await.unwrap(); } @@ -186,7 +486,7 @@ mod tests { async fn oauth_test_db() -> Option { let db = test_db().await?; - sqlx::raw_sql(include_str!("../../migrations/011_mcp_oauth.sql")) + sqlx::raw_sql(migration_sql("011_mcp_oauth.sql")) .execute(db.pool()) .await .expect("OAuth migration must create tables on a fresh test database"); @@ -930,9 +1230,9 @@ mod tests { async fn remember_jobs_test_db() -> Option { let db = test_db().await?; for migration in [ - include_str!("../../migrations/005_remember_jobs.sql"), - include_str!("../../migrations/012_remember_write_idempotency.sql"), - include_str!("../../migrations/013_remember_write_idempotency_index.sql"), + migration_sql("005_remember_jobs.sql"), + migration_sql("012_remember_write_idempotency.sql"), + migration_sql("013_remember_write_idempotency_index.sql"), ] { sqlx::raw_sql(migration).execute(db.pool()).await.unwrap(); } @@ -1126,11 +1426,60 @@ fn db_status(result: &Result) -> &'static str { } } -/// Name of the keyset-pagination index migration 016 builds. Shared -/// between the invalid-index recovery check and (in spirit) migration -/// 016's own `CREATE INDEX CONCURRENTLY IF NOT EXISTS` -- kept as a -/// constant here so the two names can't drift apart. -const PAGINATION_INDEX_NAME: &str = "idx_vector_entries_owner_updated_id"; +/// Indexes built with `CREATE INDEX CONCURRENTLY IF NOT EXISTS`, paired +/// with the migration that builds each one. +/// +/// `IF NOT EXISTS` matches by index NAME only -- it cannot tell whether +/// an index already carrying that name is usable. So any entry here can +/// be left permanently INVALID by an interrupted build, after which its +/// migration no-ops forever while the planner refuses to use what was +/// left behind. Nothing errors; the query just silently degrades. +/// +/// The interruption that used to matter was the server +/// `statement_timeout` cancelling a large `CREATE INDEX CONCURRENTLY`. +/// The locked rebuild now sets `statement_timeout = 0` for that +/// session so a large `remember_jobs` cannot leave 022 INVALID by +/// hitting the GUC. Recovery still exists for a crash or kill mid-build. +/// +/// Only indexes whose loss is SILENT belong here. Migration 013's +/// `uq_remember_jobs_owner_idempotency_key` is deliberately absent: it is +/// a unique index backing `ON CONFLICT (owner, idempotency_key)`, and an +/// invalid one makes that upsert ERROR outright rather than quietly +/// degrade, so it reports itself. It also runs in +/// `MIGRATIONS_BEFORE_BACKFILL`, ahead of the single recovery pass below. +/// +/// Add an entry whenever a migration in `MIGRATIONS_AFTER_INDEX_RECOVERY` +/// builds an index CONCURRENTLY -- +/// `every_concurrently_built_index_is_registered_for_recovery` fails the +/// suite if one is missed. +const CONCURRENTLY_BUILT_INDEXES: &[(&str, &str)] = &[ + // Keyset pagination for the memories listing endpoint. + ("idx_vector_entries_owner_updated_id", "016"), + // Expiry-refresh sweep ordering (ASC NULLS FIRST). + ("idx_vector_entries_expiry_synced_at", "018"), + // `/health` recent_write_outcomes window scan -- the probe that keeps + // `writes=degraded` alive, so losing it silently disables a + // silent-failure detector. + ("remember_jobs_recent_outcomes_idx", "022"), +]; + +/// Session advisory lock for recover + `CREATE INDEX CONCURRENTLY`. +/// Two replicas booting together (first 022 deploy, restart-all, scale-up) +/// must not interleave recover and build: `indisvalid` is false for the +/// entire live CONCURRENTLY build, so a follower that dropped "the +/// leftover" would drop the leader's in-progress (or just-finished) +/// index. Stable across deploys so rolling replicas share the key. +const CONCURRENT_INDEX_LOCK_KEYS: (i32, i32) = (872_122, 22); +const CONCURRENT_INDEX_LOCK_POLL: std::time::Duration = std::time::Duration::from_secs(2); +const CONCURRENT_INDEX_LOCK_WAIT: std::time::Duration = std::time::Duration::from_secs(20 * 60); + +fn should_drop_invalid_concurrent_index(build_in_progress: bool) -> bool { + !build_in_progress +} + +fn drop_invalid_concurrent_index_sql(name: &str) -> String { + format!("DROP INDEX CONCURRENTLY IF EXISTS {}", name) +} /// Backfill `vector_entries.updated_at` from `created_at` in bounded /// batches. @@ -1195,62 +1544,217 @@ async fn backfill_updated_at(pool: &PgPool) -> Result<(), AppError> { Ok(()) } -/// Detect and recover from an INVALID `idx_vector_entries_owner_updated_id` -/// left behind by an interrupted `CREATE INDEX CONCURRENTLY` build. +/// Drop any INVALID index left behind by an interrupted +/// `CREATE INDEX CONCURRENTLY` build, so the migration that owns it can +/// rebuild it on this same boot. /// -/// Migration 013 runs `CREATE INDEX CONCURRENTLY IF NOT EXISTS`, and -/// `IF NOT EXISTS` matches by index NAME only -- it has no idea whether -/// an existing index with that name is actually usable. A -/// `CONCURRENTLY` build that gets interrupted (process crash, -/// statement timeout, deploy killing the connection mid-build) leaves -/// behind a permanently INVALID index under the target name. From that -/// point on, every future `VectorDb::new()` sees the name already -/// exists, silently no-ops migration 013 forever, and every -/// memories-listing query keyset-paginating on `(owner, updated_at, -/// id)` silently degrades to a sequential scan -- with no error ever -/// surfaced. +/// See `CONCURRENTLY_BUILT_INDEXES` for why `IF NOT EXISTS` cannot +/// recover on its own and which indexes are in scope. /// -/// Called immediately before migration 013 runs. If an INVALID index is -/// found, it is dropped (via `DROP INDEX CONCURRENTLY`, which -- like -/// `CREATE INDEX CONCURRENTLY` -- cannot run inside a transaction -/// block, hence the bare `sqlx::query(..).execute(pool)` with no -/// explicit transaction wrapper) so migration 013's own `CREATE INDEX -/// CONCURRENTLY IF NOT EXISTS` can actually rebuild it. The recovery is -/// logged at `warn` level so it is visible in observability rather than -/// silently happening on every boot. -async fn recover_invalid_pagination_index(pool: &PgPool) -> Result<(), AppError> { - let index_is_invalid: Option = sqlx::query_scalar( - "SELECT pg_index.indisvalid FROM pg_index \ - JOIN pg_class ON pg_class.oid = pg_index.indexrelid \ - WHERE pg_class.relname = $1", - ) - .bind(PAGINATION_INDEX_NAME) - .fetch_optional(pool) - .await - .map_err(|e| { +/// `indisvalid = false` is also true for the entire duration of a *live* +/// `CREATE INDEX CONCURRENTLY`. Dropping that catalog row races two +/// replicas on first deploy of 022: the follower waits out the leader's +/// build, then drops the now-valid index (or errors if the sibling +/// already dropped it, which fails `VectorDb::new`). Skip the drop when +/// `pg_stat_progress_create_index` shows a live build for that name. +/// `DROP INDEX CONCURRENTLY IF EXISTS` so a sibling that already cleaned +/// up does not abort boot. +/// +/// Called immediately before `MIGRATIONS_AFTER_INDEX_RECOVERY` on the +/// same session that holds `CONCURRENT_INDEX_LOCK_KEYS`. +async fn recover_invalid_concurrent_indexes(pool: &PgPool) -> Result<(), AppError> { + let mut conn = pool.acquire().await.map_err(|e| { AppError::Internal(format!( - "Failed to check validity of {}: {}", - PAGINATION_INDEX_NAME, e + "Failed to acquire a connection for concurrent-index recovery: {}", + e )) })?; + recover_invalid_concurrent_indexes_on(&mut conn).await +} + +async fn recover_invalid_concurrent_indexes_on(conn: &mut PgConnection) -> Result<(), AppError> { + for (index, migration) in CONCURRENTLY_BUILT_INDEXES { + let index_is_invalid: Option = sqlx::query_scalar( + "SELECT pg_index.indisvalid FROM pg_index \ + JOIN pg_class ON pg_class.oid = pg_index.indexrelid \ + WHERE pg_class.relname = $1", + ) + .bind(index) + .fetch_optional(&mut *conn) + .await + .map_err(|e| AppError::Internal(format!("Failed to check validity of {}: {}", index, e)))?; + + // None means no index by that name -- nothing built it yet, which + // its own migration handles. Some(true) is a healthy index. + if index_is_invalid != Some(false) { + continue; + } + + if !should_drop_invalid_concurrent_index( + concurrent_index_build_in_progress(conn, index).await?, + ) { + tracing::info!( + index = %index, + migration = %migration, + "INVALID concurrent index is a live CREATE INDEX CONCURRENTLY; \ + not dropping it" + ); + continue; + } - if index_is_invalid == Some(false) { tracing::warn!( - index = PAGINATION_INDEX_NAME, - "found INVALID pagination index, likely left behind by an interrupted \ - CREATE INDEX CONCURRENTLY build -- dropping it so migration 013 can rebuild it" + index = %index, + migration = %migration, + "found INVALID index, likely left behind by an interrupted CREATE INDEX \ + CONCURRENTLY build -- dropping it so its migration can rebuild it" ); - let drop_stmt = format!("DROP INDEX CONCURRENTLY {}", PAGINATION_INDEX_NAME); - sqlx::query(&drop_stmt).execute(pool).await.map_err(|e| { + let drop_stmt = drop_invalid_concurrent_index_sql(index); + sqlx::query(&drop_stmt) + .execute(&mut *conn) + .await + .map_err(|e| { + AppError::Internal(format!("Failed to drop invalid index {}: {}", index, e)) + })?; + } + + Ok(()) +} + +async fn concurrent_index_build_in_progress( + conn: &mut PgConnection, + index: &str, +) -> Result { + match sqlx::query_scalar::<_, bool>( + "SELECT EXISTS ( + SELECT 1 + FROM pg_stat_progress_create_index p + JOIN pg_class c ON c.oid = p.index_relid + WHERE c.relname = $1 + )", + ) + .bind(index) + .fetch_one(&mut *conn) + .await + { + Ok(in_progress) => Ok(in_progress), + Err(e) => { + // Regular roles may not see other backends' progress rows. + // The session advisory lock is the fence between replica boots; + // treat an unreadable view as "no live build we can see". + tracing::warn!( + error = %e, + index, + "could not read pg_stat_progress_create_index; \ + not treating the INVALID index as a live build" + ); + Ok(false) + } + } +} + +/// Recover leftover INVALID indexes, then run 016/018/022, on one session. +/// +/// The lock serializes replica boots so a follower cannot mistake the +/// leader's in-progress CONCURRENTLY build for a crashed leftover. It is +/// polled with `pg_try_advisory_lock` rather than waited on: a blocking +/// `pg_advisory_lock` waiter holds a snapshot for the whole wait, and +/// `CREATE INDEX CONCURRENTLY` waits for exactly such snapshots to finish, +/// so the two deadlock. `statement_timeout = 0` for this session so a large +/// `remember_jobs` cannot leave 022 INVALID by hitting the server GUC. +async fn recover_and_rebuild_concurrent_indexes(pool: &PgPool) -> Result<(), AppError> { + let mut conn = pool.acquire().await.map_err(|e| { + AppError::Internal(format!( + "Failed to acquire a connection for concurrent-index rebuild: {}", + e + )) + })?; + + let previous: String = sqlx::query_scalar("SHOW statement_timeout") + .fetch_one(&mut *conn) + .await + .map_err(|e| AppError::Internal(format!("Failed to read statement_timeout: {}", e)))?; + + sqlx::query("SET statement_timeout = 0") + .execute(&mut *conn) + .await + .map_err(|e| { AppError::Internal(format!( - "Failed to drop invalid index {}: {}", - PAGINATION_INDEX_NAME, e + "Failed to disable statement_timeout for concurrent index builds: {}", + e )) })?; + + let (result, lock_released) = lock_and_rebuild_concurrent_indexes(&mut conn).await; + + let timeout_restored = sqlx::query("SELECT set_config('statement_timeout', $1, false)") + .bind(&previous) + .execute(&mut *conn) + .await + .is_ok(); + + if !timeout_restored || !lock_released { + // Either a leaked `statement_timeout = 0` or a still-held session + // advisory lock. Both ride on the connection, so close it instead of + // returning it to the pool. + tracing::warn!( + timeout_restored, + lock_released, + "concurrent-index rebuild left the session dirty; discarding it \ + rather than returning it to the pool" + ); + let _ = conn.close().await; } - Ok(()) + result +} + +async fn lock_and_rebuild_concurrent_indexes( + conn: &mut PgConnection, +) -> (Result<(), AppError>, bool) { + let mut waited = std::time::Duration::ZERO; + loop { + let acquired = sqlx::query_scalar::<_, bool>("SELECT pg_try_advisory_lock($1, $2)") + .bind(CONCURRENT_INDEX_LOCK_KEYS.0) + .bind(CONCURRENT_INDEX_LOCK_KEYS.1) + .fetch_one(&mut *conn) + .await + .map_err(|e| { + AppError::Internal(format!("Failed to lock concurrent-index rebuild: {}", e)) + }); + match acquired { + Ok(true) => break, + Ok(false) => {} + Err(e) => return (Err(e), true), + } + if waited >= CONCURRENT_INDEX_LOCK_WAIT { + return ( + Err(AppError::Internal(format!( + "Timed out after {}s waiting for the concurrent-index rebuild lock", + CONCURRENT_INDEX_LOCK_WAIT.as_secs() + ))), + true, + ); + } + tokio::time::sleep(CONCURRENT_INDEX_LOCK_POLL).await; + waited += CONCURRENT_INDEX_LOCK_POLL; + } + + let result = async { + recover_invalid_concurrent_indexes_on(conn).await?; + run_migrations_on(conn, MIGRATIONS_AFTER_INDEX_RECOVERY).await?; + Ok(()) + } + .await; + + let unlocked = sqlx::query("SELECT pg_advisory_unlock($1, $2)") + .bind(CONCURRENT_INDEX_LOCK_KEYS.0) + .bind(CONCURRENT_INDEX_LOCK_KEYS.1) + .execute(&mut *conn) + .await + .is_ok(); + + (result, unlocked) } /// Release storage reservations given only a pool handle. @@ -1330,190 +1834,19 @@ impl VectorDb { .await .map_err(|e| AppError::Internal(format!("Failed to connect to database: {}", e)))?; - // Run migrations - let migration_001 = include_str!("../../migrations/001_init.sql"); - sqlx::raw_sql(migration_001) - .execute(&pool) - .await - .map_err(|e| AppError::Internal(format!("Failed to run migration 001: {}", e)))?; - - let migration_002 = include_str!("../../migrations/002_add_namespace.sql"); - sqlx::raw_sql(migration_002) - .execute(&pool) - .await - .map_err(|e| AppError::Internal(format!("Failed to run migration 002: {}", e)))?; - - let migration_003 = include_str!("../../migrations/003_rate_limiter.sql"); - sqlx::raw_sql(migration_003) - .execute(&pool) - .await - .map_err(|e| AppError::Internal(format!("Failed to run migration 003: {}", e)))?; - - let migration_004 = include_str!("../../migrations/004_delegate_key_cache_expires.sql"); - sqlx::raw_sql(migration_004) - .execute(&pool) - .await - .map_err(|e| AppError::Internal(format!("Failed to run migration 004: {}", e)))?; - - let migration_005 = include_str!("../../migrations/005_remember_jobs.sql"); - sqlx::raw_sql(migration_005) - .execute(&pool) - .await - .map_err(|e| AppError::Internal(format!("Failed to run migration 005: {}", e)))?; - - // composite index on (owner, status, updated_at DESC) for bulk poll - let migration_006 = include_str!("../../migrations/006_bulk_remember.sql"); - sqlx::raw_sql(migration_006) - .execute(&pool) - .await - .map_err(|e| AppError::Internal(format!("Failed to run migration 006: {}", e)))?; - - // collapse per-wallet Apalis queues to a single `wallet_jobs` - // queue. Equivocation locks are no longer a practical concern on Sui - // (per Will Bradley, Mysten, 2026-05-12); concurrent workers on one - // wallet + retry handling is sufficient. - let migration_007 = include_str!("../../migrations/007_collapse_wallet_queues.sql"); - sqlx::raw_sql(migration_007) - .execute(&pool) - .await - .map_err(|e| AppError::Internal(format!("Failed to run migration 007: {}", e)))?; - - // nullable `plaintext` column for benchmark-mode storage - // (PlaintextEngine). NULL for all production rows — additive. - // Renumbered from 007 → 008 during rebase onto dev to avoid collision - // with the wallet-queue collapse migration. - let migration_008 = include_str!("../../migrations/008_benchmark_plaintext.sql"); - sqlx::raw_sql(migration_008) - .execute(&pool) - .await - .map_err(|e| AppError::Internal(format!("Failed to run migration 008: {}", e)))?; - - // importance signal column on vector_entries. - let migration_009 = include_str!("../../migrations/009_importance_signal.sql"); - sqlx::raw_sql(migration_009) - .execute(&pool) - .await - .map_err(|e| AppError::Internal(format!("Failed to run migration 009: {}", e)))?; - - // Permanent restore-failure negative cache (GH #501 / WALM-299). - let migration_010 = include_str!("../../migrations/010_restore_failed_blobs.sql"); - sqlx::raw_sql(migration_010) - .execute(&pool) - .await - .map_err(|e| AppError::Internal(format!("Failed to run migration 010: {}", e)))?; - - // MCP OAuth 2.1 (Claude custom connectors): client registry, - // server-custodied delegate keys, and authorization state. - let migration_011 = include_str!("../../migrations/011_mcp_oauth.sql"); - sqlx::raw_sql(migration_011) - .execute(&pool) - .await - .map_err(|e| AppError::Internal(format!("Failed to run migration 011: {}", e)))?; - - // Durable idempotency, preparation, and paid-upload recovery state. - let migration_012 = include_str!("../../migrations/012_remember_write_idempotency.sql"); - sqlx::raw_sql(migration_012) - .execute(&pool) - .await - .map_err(|e| AppError::Internal(format!("Failed to run migration 012: {}", e)))?; - - // Build the owner/key uniqueness constraint without blocking writes. - let migration_013 = - include_str!("../../migrations/013_remember_write_idempotency_index.sql"); - sqlx::raw_sql(migration_013) - .execute(&pool) - .await - .map_err(|e| AppError::Internal(format!("Failed to run migration 013: {}", e)))?; - - // per-owner storage quota reservations. Makes quota admission atomic - // with the eventual insert (GH #532 / WALM-359). - let migration_014_reservations = - include_str!("../../migrations/014_storage_reservations.sql"); - sqlx::raw_sql(migration_014_reservations) - .execute(&pool) - .await - .map_err(|e| { - AppError::Internal(format!( - "Failed to run migration 014 (storage reservations): {}", - e - )) - })?; - - // owner-scoped read API: updated_at cursor column + agent_id/package_id. - // Split across 014-019 (see each file's header, and - // backfill_updated_at's / recover_invalid_pagination_index's doc - // comments above) to avoid holding ACCESS EXCLUSIVE across the - // full-table backfill or index build. - let migration_014_read_api = - include_str!("../../migrations/014_memory_read_api_columns.sql"); - sqlx::raw_sql(migration_014_read_api) - .execute(&pool) - .await - .map_err(|e| { - AppError::Internal(format!( - "Failed to run migration 014 (read API columns): {}", - e - )) - })?; + // Run migrations. The ordering, and the two Rust steps woven + // between these slices, are load-bearing — see + // MIGRATIONS_BEFORE_BACKFILL's comment. + run_migrations(&pool, MIGRATIONS_BEFORE_BACKFILL).await?; // Backfill runs as batched Rust code, not a migration file, since // Postgres can't COMMIT mid-loop inside a plain migration // statement — see backfill_updated_at()'s doc comment. backfill_updated_at(&pool).await?; - // Requires the backfill above to have already completed — this - // validates NOT NULL and will error if any updated_at row is - // still NULL. - let migration_015 = - include_str!("../../migrations/015_memory_read_api_updated_at_not_null.sql"); - sqlx::raw_sql(migration_015) - .execute(&pool) - .await - .map_err(|e| AppError::Internal(format!("Failed to run migration 015: {}", e)))?; - - // Must run before migration 016's CREATE INDEX CONCURRENTLY IF NOT - // EXISTS, which would otherwise silently no-op forever against a - // permanently INVALID index from an interrupted build. - recover_invalid_pagination_index(&pool).await?; - - // keyset-pagination index for the memories listing endpoint. - // Must stay in its own file/transaction — see 016's header comment. - let migration_016 = include_str!("../../migrations/016_memory_read_api_index.sql"); - sqlx::raw_sql(migration_016) - .execute(&pool) - .await - .map_err(|e| AppError::Internal(format!("Failed to run migration 016: {}", e)))?; - - // per-memory expiry columns. - let migration_017 = include_str!("../../migrations/017_memory_expiry_columns.sql"); - sqlx::raw_sql(migration_017) - .execute(&pool) - .await - .map_err(|e| AppError::Internal(format!("Failed to run migration 017: {}", e)))?; - - // index on expiry_synced_at so the periodic expiry refresh sweep - // doesn't full-scan vector_entries every tick. Must stay - // in its own file/transaction — see 018's header comment. - let migration_018 = include_str!("../../migrations/018_memory_expiry_synced_at_index.sql"); - sqlx::raw_sql(migration_018) - .execute(&pool) - .await - .map_err(|e| AppError::Internal(format!("Failed to run migration 018: {}", e)))?; + run_migrations(&pool, MIGRATIONS_AFTER_BACKFILL).await?; - // Finalizes updated_at NOT NULL cheaply using the validated CHECK - // constraint 015 set up — see 019's header. - let migration_019 = - include_str!("../../migrations/019_memory_read_api_updated_at_set_not_null.sql"); - sqlx::raw_sql(migration_019) - .execute(&pool) - .await - .map_err(|e| AppError::Internal(format!("Failed to run migration 019: {}", e)))?; - - let migration_020 = include_str!("../../migrations/020_read_api_followups.sql"); - sqlx::raw_sql(migration_020) - .execute(&pool) - .await - .map_err(|e| AppError::Internal(format!("Failed to run migration 020: {}", e)))?; + recover_and_rebuild_concurrent_indexes(&pool).await?; tracing::info!("database connected and migrations applied"); @@ -1837,6 +2170,73 @@ impl VectorDb { Ok(row) } + /// How durable writes that finished inside `window` turned out, + /// across every owner: `(failed, succeeded)`. + /// + /// `write_ready` is a sidecar probe AND a Postgres size check. Neither + /// can see Walrus refusing every upload, so a total Walrus outage left + /// `/health` reporting a healthy write path while every remember + /// failed minutes after being accepted. This is the missing term. + /// + /// Counted rather than listed, and read behind a cache measured in + /// tens of seconds. Served by `remember_jobs_recent_outcomes_idx` + /// (022). The statement is also cancelled at 1s (`SET LOCAL`) so a + /// sequential scan cannot stall the public `/health` handler; the + /// probe fails open on timeout. + pub async fn recent_write_outcomes( + &self, + window: std::time::Duration, + ) -> Result<(i64, i64), AppError> { + let started = std::time::Instant::now(); + let since = chrono::Utc::now() - chrono::Duration::from_std(window).unwrap_or_default(); + // SET LOCAL needs a transaction; the pool's statement_timeout is + // the startup bound (up to 300s) and would let this scan run that + // long. 1000ms matches WRITE_READY_PROBE_TIMEOUT on /health. + let mut tx = self.pool.begin().await.map_err(|e| { + AppError::Internal(format!("Failed to count recent remember outcomes: {}", e)) + })?; + sqlx::query("SELECT set_config('statement_timeout', '1000ms', true)") + .execute(&mut *tx) + .await + .map_err(|e| { + AppError::Internal(format!("Failed to count recent remember outcomes: {}", e)) + })?; + let outcome = sqlx::query_as::<_, (i64, i64)>( + "SELECT + count(*) FILTER (WHERE status = 'failed'), + count(*) FILTER (WHERE status IN ('done', 'uploaded')) + FROM remember_jobs + WHERE status IN ('failed', 'done', 'uploaded') + AND updated_at >= $1", + ) + .bind(since) + .fetch_one(&mut *tx) + .await; + let _ = tx.rollback().await; + + match outcome { + Ok(counts) => { + crate::observability::observe_db( + "remember_jobs.recent_outcomes", + "ok", + started.elapsed(), + ); + Ok(counts) + } + Err(e) => { + crate::observability::observe_db( + "remember_jobs.recent_outcomes", + "error", + started.elapsed(), + ); + Err(AppError::Internal(format!( + "Failed to count recent remember outcomes: {}", + e + ))) + } + } + } + /// Hard-delete all vector index rows for a given owner + namespace. /// (Walrus blobs themselves persist — Walrus has no delete; this only /// removes the local `vector_entries` rows, so the memories stop being @@ -2103,9 +2503,22 @@ impl VectorDb { Ok(rows) } - /// Mark worker-claimed remember jobs as failed when no worker has updated - /// them within the stale TTL. Pending rows are left alone because they may - /// simply be waiting behind legitimate queue backlog. + /// Mark remember jobs as failed once nothing can still move them: + /// `running`/`uploaded` whose worker stopped updating them, and `pending` + /// rows whose preparation never finished. + /// + /// `prepare_claimed_at IS NOT NULL` is load-bearing — only the single + /// `remember` path claims a preparation slot, so without it the sweep also + /// matches healthy `/api/remember/bulk` and `/api/analyze` rows, which + /// never set `preparation_encrypted_b64` at all. + /// + /// Clearing `prepare_claim_token` fences a slow preparation: its own + /// UPDATE is keyed on that token, so it can no longer reach + /// `enqueue_wallet_job`. Failing is the only option — the row stores + /// ciphertext, never plaintext, so nothing can be retried from. + /// + /// Quota is reclaimed by `release_reservations_for_terminal_jobs`, which + /// `main` runs immediately after this on the same tick. pub async fn fail_stale_remember_jobs( &self, stale_after: std::time::Duration, @@ -2128,7 +2541,45 @@ impl VectorDb { if rows > 0 { tracing::warn!("Marked {} stale remember jobs as failed", rows); } - Ok(rows) + + // Second pass rather than one OR'd predicate: this branch clears the + // preparation claim and carries its own error text, and the two are + // different enough that folding them together would hide which case + // actually fired in the logs. + let orphaned = sqlx::query( + "UPDATE remember_jobs + SET status = 'failed', + error_msg = COALESCE( + error_msg, + 'preparation never completed — the relayer stopped before this write was encrypted, so the fact was never stored and must be sent again' + ), + prepare_claim_token = NULL, + prepare_claimed_at = NULL, + updated_at = NOW() + WHERE status = 'pending' + AND prepare_claimed_at IS NOT NULL + AND preparation_encrypted_b64 IS NULL + AND updated_at < NOW() - ($1 * INTERVAL '1 second')", + ) + .bind(stale_after_secs) + .execute(&self.pool) + .await + .map_err(|e| { + AppError::Internal(format!("Failed to fail orphaned remember preparations: {}", e)) + })?; + + let orphaned_rows = orphaned.rows_affected(); + if orphaned_rows > 0 { + // Distinct wording from the sweep above: this one means writes were + // accepted and silently lost, which is an availability signal about + // the relayer, not a Walrus or wallet problem. + tracing::warn!( + "Marked {} remember jobs as failed whose preparation never completed", + orphaned_rows + ); + } + + Ok(rows + orphaned_rows) } /// Rows whose expiry data has never been synced, or was synced more @@ -2751,6 +3202,24 @@ impl VectorDb { Ok(()) } + pub async fn fetch_oauth_code( + &self, + client_id: &str, + code_sha256: &str, + ) -> Result, AppError> { + sqlx::query_as::<_, oauth_rows::OAuthCodeRow>( + "SELECT code_sha256, client_id, redirect_uri, scope, resource, code_challenge, + code_challenge_method, delegate_ref, account_id, owner_address, expires_at + FROM mcp_oauth_codes + WHERE code_sha256 = $1 AND client_id = $2 AND expires_at > NOW()", + ) + .bind(code_sha256) + .bind(client_id) + .fetch_optional(&self.pool) + .await + .map_err(|e| AppError::Internal(format!("Failed to read oauth code: {}", e))) + } + /// Single-use consume via `DELETE ... RETURNING` — the first successful /// exchange deletes the row; any replay finds nothing. Also filters on /// `client_id` so a code minted for one client can never be redeemed by @@ -3222,6 +3691,10 @@ mod quota_admission_tests { } async fn test_db() -> VectorDb { + let _guard = super::DB_SETUP_LOCK + .get_or_init(|| tokio::sync::Mutex::new(())) + .lock() + .await; VectorDb::new(&test_database_url()) .await .expect("test database must be reachable with pgvector installed") @@ -3629,3 +4102,238 @@ mod quota_admission_tests { cleanup(&db, &owner).await; } } + +#[cfg(test)] +mod stale_sweep_tests { + use super::*; + use std::time::Duration; + + fn test_database_url() -> String { + std::env::var("DATABASE_URL") + .unwrap_or_else(|_| "postgresql://memwal:memwal_secret@localhost:5432/memwal".into()) + } + + async fn test_db() -> VectorDb { + let _guard = super::DB_SETUP_LOCK + .get_or_init(|| tokio::sync::Mutex::new(())) + .lock() + .await; + VectorDb::new(&test_database_url()) + .await + .expect("test database must be reachable with pgvector installed") + } + + /// Unique per test so concurrent runs cannot see each other's rows. + fn unique_owner(tag: &str) -> String { + format!("0xtest-{}-{}", tag, uuid::Uuid::new_v4()) + } + + /// Insert one remember job, aged by `age_secs`, optionally already prepared. + /// + /// `prepare_claimed_at` is stamped only alongside a claim token, because + /// that is the only way a row can reach the database: the single-write + /// path claims and stamps together, while `/api/remember/bulk` and + /// `/api/analyze` insert neither. Stamping it unconditionally would hand + /// every seeded row the one column the orphan sweep keys on, so a helper + /// detail — not the sweep — would decide what the tests below prove. + async fn seed_job( + db: &VectorDb, + owner: &str, + status: &str, + prepared: bool, + claim_token: Option<&str>, + age_secs: i64, + ) -> String { + let id = uuid::Uuid::new_v4().to_string(); + sqlx::query( + "INSERT INTO remember_jobs + (id, owner, namespace, status, preparation_encrypted_b64, + prepare_claim_token, prepare_claimed_at, created_at, updated_at) + VALUES ($1, $2, 'default', $3, $4, $5, + CASE WHEN $5::text IS NULL THEN NULL + ELSE NOW() - ($6 * INTERVAL '1 second') END, + NOW() - ($6 * INTERVAL '1 second'), + NOW() - ($6 * INTERVAL '1 second'))", + ) + .bind(&id) + .bind(owner) + .bind(status) + .bind(if prepared { Some("ZW5jcnlwdGVk") } else { None }) + .bind(claim_token) + .bind(age_secs) + .execute(&db.pool) + .await + .expect("seed remember job"); + id + } + + async fn status_of(db: &VectorDb, id: &str) -> String { + sqlx::query_scalar("SELECT status FROM remember_jobs WHERE id = $1") + .bind(id) + .fetch_one(&db.pool) + .await + .expect("read status") + } + + /// A row committed by the route whose preparation never ran has no task + /// left to resume it: preparation lives in a `tokio::spawn` inside the + /// relayer, so a restart in that window strands it. Before this it sat at + /// `pending` forever and `memwal_remember_status` reported it as still + /// uploading — a write silently lost while the user was told it was coming. + #[tokio::test] + async fn orphaned_preparation_is_failed() { + let db = test_db().await; + let owner = unique_owner("orphan"); + let id = seed_job(&db, &owner, "pending", false, Some("claim-1"), 900).await; + + db.fail_stale_remember_jobs(Duration::from_secs(600)) + .await + .expect("sweep"); + + assert_eq!(status_of(&db, &id).await, "failed"); + let msg: Option = + sqlx::query_scalar("SELECT error_msg FROM remember_jobs WHERE id = $1") + .bind(&id) + .fetch_one(&db.pool) + .await + .unwrap(); + assert!( + msg.unwrap_or_default().contains("never stored"), + "the message has to say the fact is gone, not merely that a job died", + ); + } + + /// The reason `pending` cannot be swept wholesale. A prepared job waits at + /// `pending` until a wallet worker takes it, and with + /// `WALRUS_UPLOAD_PER_WALLET_CONCURRENCY` defaulting to 1 that queue is + /// legitimately minutes deep. Failing these would abandon paid work that + /// was about to run. + #[tokio::test] + async fn prepared_job_waiting_on_the_upload_queue_is_left_alone() { + let db = test_db().await; + let owner = unique_owner("queued"); + let id = seed_job(&db, &owner, "pending", true, Some("claim-1"), 900).await; + + db.fail_stale_remember_jobs(Duration::from_secs(600)) + .await + .expect("sweep"); + + assert_eq!(status_of(&db, &id).await, "pending"); + } + + /// Preparation itself takes a moment (summarize, embed, SEAL encrypt), so + /// a young unprepared row is in-flight, not orphaned. + #[tokio::test] + async fn a_preparation_still_in_flight_is_left_alone() { + let db = test_db().await; + let owner = unique_owner("young"); + let id = seed_job(&db, &owner, "pending", false, Some("claim-1"), 5).await; + + db.fail_stale_remember_jobs(Duration::from_secs(600)) + .await + .expect("sweep"); + + assert_eq!(status_of(&db, &id).await, "pending"); + } + + /// What makes the sweep safe against a preparation that was merely very + /// slow rather than dead. Its own UPDATE is fenced on the claim token, so + /// once the sweeper clears it that statement matches zero rows and the task + /// returns before `enqueue_wallet_job` — it cannot mint a paid blob for a + /// job just declared dead. + #[tokio::test] + async fn clearing_the_claim_fences_a_late_preparation() { + let db = test_db().await; + let owner = unique_owner("fence"); + let id = seed_job(&db, &owner, "pending", false, Some("claim-1"), 900).await; + + db.fail_stale_remember_jobs(Duration::from_secs(600)) + .await + .expect("sweep"); + + // Exactly the statement `spawn_prepare_remember_job` runs when it + // finishes, token and all. + let late = sqlx::query( + "UPDATE remember_jobs SET preparation_encrypted_b64 = $1, updated_at = NOW() + WHERE id = $2 AND ($3::TEXT IS NULL OR prepare_claim_token = $3)", + ) + .bind("ZW5jcnlwdGVk") + .bind(&id) + .bind("claim-1") + .execute(&db.pool) + .await + .expect("late preparation"); + + assert_eq!( + late.rows_affected(), + 0, + "a late preparation must be fenced out, or it would queue a paid write for a dead job", + ); + } + + /// The pre-existing sweep is unchanged. + #[tokio::test] + async fn a_stalled_worker_claim_still_fails() { + let db = test_db().await; + let owner = unique_owner("running"); + let id = seed_job(&db, &owner, "running", true, None, 900).await; + + db.fail_stale_remember_jobs(Duration::from_secs(600)) + .await + .expect("sweep"); + + assert_eq!(status_of(&db, &id).await, "failed"); + } + + /// `/api/remember/bulk` inserts its rows directly and never claims a + /// preparation slot, so `preparation_encrypted_b64` is ALWAYS NULL for a + /// bulk job — healthy or not. Keying the sweep on that column alone would + /// fail every bulk write that waits out the TTL behind a normal upload + /// backlog, destroying paid work that was about to run. + #[tokio::test] + async fn a_queued_bulk_job_is_never_swept() { + let db = test_db().await; + let owner = unique_owner("bulk"); + // Exactly what remember_bulk writes: pending, no claim, no preparation. + let id = seed_job(&db, &owner, "pending", false, None, 900).await; + + db.fail_stale_remember_jobs(Duration::from_secs(600)) + .await + .expect("sweep"); + + assert_eq!( + status_of(&db, &id).await, + "pending", + "a bulk job waiting on the upload queue must survive the sweep", + ); + } + + /// `/api/analyze` inserts the same shape (`prepare_claim_token: None`), so + /// it needs the same protection. + #[tokio::test] + async fn a_queued_analyze_job_is_never_swept() { + let db = test_db().await; + let owner = unique_owner("analyze"); + let id = seed_job(&db, &owner, "pending", false, None, 3600).await; + + db.fail_stale_remember_jobs(Duration::from_secs(600)) + .await + .expect("sweep"); + + assert_eq!(status_of(&db, &id).await, "pending"); + } + + /// A finished write is terminal and the sweeper must never touch it. + #[tokio::test] + async fn a_done_job_is_never_swept() { + let db = test_db().await; + let owner = unique_owner("done"); + let id = seed_job(&db, &owner, "done", true, None, 900).await; + + db.fail_stale_remember_jobs(Duration::from_secs(600)) + .await + .expect("sweep"); + + assert_eq!(status_of(&db, &id).await, "done"); + } +} diff --git a/services/server/src/storage/sui.rs b/services/server/src/storage/sui.rs index e48906014..2601fce7f 100644 --- a/services/server/src/storage/sui.rs +++ b/services/server/src/storage/sui.rs @@ -24,6 +24,7 @@ pub async fn verify_delegate_key_onchain( account_object_id: &str, public_key_bytes: &[u8], expected_type_origin_package_id: &str, + get_object_attempts: u32, ) -> Result { if let Some(grpc_client) = grpc_client { return verify_delegate_key_onchain_grpc( @@ -31,6 +32,7 @@ pub async fn verify_delegate_key_onchain( account_object_id, public_key_bytes, expected_type_origin_package_id, + get_object_attempts, ) .await; } @@ -475,15 +477,14 @@ pub async fn verify_delegate_key_cached( // nothing. The owned key is built only where the map is actually written. let probe: &dyn DelegateAccountKey = &(account_object_id, public_key_bytes); - if let Some(cached) = cache - .entries - .read() - .await - .get(probe) - .filter(|c| c.is_fresh()) - { - return Ok(cached.owner.clone()); - } + let get_object_attempts = { + let entries = cache.entries.read().await; + match entries.get(probe) { + Some(cached) if cached.is_fresh() => return Ok(cached.owner.clone()), + Some(cached) if cached.is_servable_while_unavailable() => 1, + _ => GET_OBJECT_ATTEMPTS, + } + }; // Read before the chain call, compared after it. See `evictions`. let generation_before = cache.evictions.load(std::sync::atomic::Ordering::Acquire); @@ -513,6 +514,7 @@ pub async fn verify_delegate_key_cached( account_object_id, public_key_bytes, expected_type_origin_package_id, + get_object_attempts, ) .await { @@ -870,9 +872,54 @@ fn map_get_object_status(status: tonic::Status) -> OnchainVerifyError { } } +pub const GET_OBJECT_ATTEMPTS: u32 = 3; +const GET_OBJECT_RETRY_BASE_DELAY_MS: u64 = 500; +static GET_OBJECT_RETRY_JITTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); + +async fn with_get_object_retry(attempts: u32, mut call: F) -> Result +where + F: FnMut() -> Fut, + Fut: std::future::Future>, +{ + for attempt in 0..attempts { + let started = std::time::Instant::now(); + let result = + match tokio::time::timeout(crate::sui::DEFAULT_RPC_ATTEMPT_TIMEOUT, call()).await { + Ok(result) => result, + Err(_) => Err(tonic::Status::deadline_exceeded(format!( + "GetObject exceeded {}ms", + crate::sui::DEFAULT_RPC_ATTEMPT_TIMEOUT.as_millis() + ))), + }; + let status_label = match &result { + Ok(_) => "200".to_string(), + Err(status) => status.code().to_string(), + }; + crate::observability::observe_external( + "sui_grpc", + "GetObject", + &status_label, + started.elapsed(), + ); + + match result { + Err(status) + if attempt + 1 < attempts && crate::sui::is_transient_grpc_code(status.code()) => {} + outcome => return outcome, + } + + let base = GET_OBJECT_RETRY_BASE_DELAY_MS.saturating_mul(1_u64 << attempt.min(16)); + let sequence = GET_OBJECT_RETRY_JITTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + let jitter = base.saturating_mul(((sequence * 37 + u64::from(attempt)) % 21) + 90) / 100; + tokio::time::sleep(std::time::Duration::from_millis(jitter)).await; + } + unreachable!("attempts is non-zero") +} + async fn grpc_get_object( - mut client: sui_rpc::Client, + client: sui_rpc::Client, account_object_id: &str, + attempts: u32, ) -> Result { let address = parse_object_id(account_object_id)?; let mut request = sui_rpc::proto::sui::rpc::v2::GetObjectRequest::new(&address); @@ -880,24 +927,16 @@ async fn grpc_get_object( paths: vec!["json".to_string(), "object_type".to_string()], }); - let started = std::time::Instant::now(); - let response = client.ledger_client().get_object(request).await; - let status_label = match &response { - Ok(_) => "200".to_string(), - Err(status) => status.code().to_string(), - }; - crate::observability::observe_external( - "sui_grpc", - "GetObject", - &status_label, - started.elapsed(), - ); - - response - .map_err(map_get_object_status)? - .into_inner() - .object - .ok_or_else(|| OnchainVerifyError::NotFound("gRPC response missing object".into())) + with_get_object_retry(attempts, || { + let mut client = client.clone(); + let request = request.clone(); + async move { client.ledger_client().get_object(request).await } + }) + .await + .map_err(map_get_object_status)? + .into_inner() + .object + .ok_or_else(|| OnchainVerifyError::NotFound("gRPC response missing object".into())) } /// gRPC counterpart of `verify_delegate_key_onchain` above — same checks @@ -913,8 +952,9 @@ async fn verify_delegate_key_onchain_grpc( account_object_id: &str, public_key_bytes: &[u8], expected_type_origin_package_id: &str, + attempts: u32, ) -> Result { - let object = grpc_get_object(client, account_object_id).await?; + let object = grpc_get_object(client, account_object_id, attempts).await?; // #398: verify the Move type before trusting any field (gRPC path). ensure_memwal_account_type( @@ -992,7 +1032,7 @@ async fn list_delegate_keys_onchain_grpc( account_object_id: &str, expected_type_origin_package_id: &str, ) -> Result, OnchainVerifyError> { - let object = grpc_get_object(client, account_object_id).await?; + let object = grpc_get_object(client, account_object_id, GET_OBJECT_ATTEMPTS).await?; ensure_memwal_account_type( object.object_type.as_deref(), @@ -1241,6 +1281,7 @@ pub async fn find_account_by_delegate_key( account_id, public_key_bytes, expected_type_origin_package_id, + GET_OBJECT_ATTEMPTS, ) .await { @@ -1781,6 +1822,45 @@ mod tests { } } + #[tokio::test(start_paused = true)] + async fn get_object_retries_a_transient_unavailable() { + let calls = std::sync::atomic::AtomicU32::new(0); + let object = with_get_object_retry(GET_OBJECT_ATTEMPTS, || { + let attempt = calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + async move { + if attempt == 0 { + Err(tonic::Status::unavailable( + "The service is currently unavailable", + )) + } else { + Ok("object") + } + } + }) + .await + .expect("second attempt succeeds"); + + assert_eq!(object, "object"); + assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 2); + } + + #[tokio::test(start_paused = true)] + async fn get_object_makes_one_attempt_when_a_stale_entry_can_be_served() { + let calls = std::sync::atomic::AtomicU32::new(0); + let result: Result<&str, tonic::Status> = with_get_object_retry(1, || { + calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + async move { + Err(tonic::Status::unavailable( + "The service is currently unavailable", + )) + } + }) + .await; + + assert_eq!(result.unwrap_err().code(), tonic::Code::Unavailable); + assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 1); + } + #[test] fn invalid_object_id_is_not_unavailable() { let err = parse_object_id("not-a-sui-object-id").unwrap_err(); @@ -2930,8 +3010,14 @@ mod tests { let expected_pkg = "0xcf6ad755a1cdff7217865c796778fabe5aa399cb0cf2eba986f4b582047229c6"; let client = sui_rpc::Client::new("https://fullnode.testnet.sui.io").unwrap(); - let result = - verify_delegate_key_onchain_grpc(client, account_id, &wrong_key, expected_pkg).await; + let result = verify_delegate_key_onchain_grpc( + client, + account_id, + &wrong_key, + expected_pkg, + GET_OBJECT_ATTEMPTS, + ) + .await; // The account genuinely exists and gRPC parses it correctly — a // non-matching key must fail with KeyNotFound, not RpcError. Getting @@ -2956,6 +3042,7 @@ mod tests { fake_id, &[0u8; 32], "0xcf6ad755a1cdff7217865c796778fabe5aa399cb0cf2eba986f4b582047229c6", + GET_OBJECT_ATTEMPTS, ) .await; assert!( diff --git a/services/server/src/sui/client.rs b/services/server/src/sui/client.rs index 0d63d2f54..36d365e64 100644 --- a/services/server/src/sui/client.rs +++ b/services/server/src/sui/client.rs @@ -27,7 +27,7 @@ const MAX_RATE_LIMIT_COOLDOWN: Duration = Duration::from_secs(60 * 60); // Must cover one full 10-second provider window: a request arriving just // after the 2,970th admission still gets a bounded chance to enter next window. const INTERACTIVE_GATE_BUDGET: Duration = Duration::from_secs(12); -const DEFAULT_RPC_ATTEMPT_TIMEOUT: Duration = Duration::from_secs(5); +pub const DEFAULT_RPC_ATTEMPT_TIMEOUT: Duration = Duration::from_secs(5); const MAX_RPC_ATTEMPT_TIMEOUT: Duration = Duration::from_secs(60); const DEFAULT_RPC_MAX_IN_FLIGHT: usize = 64; const MAX_RPC_MAX_IN_FLIGHT: usize = 10_000; @@ -365,14 +365,7 @@ impl SuiClient { } SuiErr::RateLimited } - Ok(Err(status)) - if matches!( - status.code(), - tonic::Code::Unavailable - | tonic::Code::DeadlineExceeded - | tonic::Code::Aborted - ) => - { + Ok(Err(status)) if is_transient_grpc_code(status.code()) => { SuiErr::Transport(status.to_string()) } Ok(Err(status)) => return Err(Self::classify_rejection(status)), @@ -459,6 +452,13 @@ impl SuiClient { } } +pub fn is_transient_grpc_code(code: tonic::Code) -> bool { + matches!( + code, + tonic::Code::Unavailable | tonic::Code::DeadlineExceeded | tonic::Code::Aborted + ) +} + fn provider_retry_delay(status: &tonic::Status) -> Duration { if let Some(delay) = metadata_duration(status, "grpc-retry-pushback-ms", true) { return delay.min(MAX_RATE_LIMIT_COOLDOWN); diff --git a/services/server/src/sui/mod.rs b/services/server/src/sui/mod.rs index cd79bb003..f48bb811d 100644 --- a/services/server/src/sui/mod.rs +++ b/services/server/src/sui/mod.rs @@ -9,7 +9,10 @@ pub mod input_freshness; pub mod tx_build; pub mod verifier; -pub use client::{expires_at_from_epoch, RequestPriority, SuiClient, WalrusEpochSchedule}; +pub use client::{ + expires_at_from_epoch, is_transient_grpc_code, RequestPriority, SuiClient, WalrusEpochSchedule, + DEFAULT_RPC_ATTEMPT_TIMEOUT, +}; use async_trait::async_trait; use sui_sdk_types::ExecutionStatus; diff --git a/services/server/src/types.rs b/services/server/src/types.rs index 61c113537..0db58a048 100644 --- a/services/server/src/types.rs +++ b/services/server/src/types.rs @@ -1,3 +1,7 @@ + /// Sanitized by `sanitize_job_error_for_client`, as on every other + /// client-facing job-status path: an infrastructure-funding failure is + /// replaced wholesale (its raw text names the relayer's own wallet and + /// balance), and long hex runs are redacted. use base64::Engine as _; use serde::{Deserialize, Serialize}; use std::sync::Arc; @@ -290,13 +294,30 @@ pub struct AppState { pub struct KeyPool { keys: Vec, cursor: AtomicUsize, + /// Wallet transactions currently executing, per key. + /// + /// The sidecar enforces one upload at a time per wallet (concurrent + /// transactions from one signer can equivocate its owned objects, which + /// then stay locked until the epoch boundary). Blind round-robin does not + /// know that, so it hands a job to a wallet that is mid-upload while other + /// wallets sit idle — observed in production as a job waiting 6.2s for its + /// assigned wallet while the global limiter still had two free slots and + /// an empty queue. + /// + /// Counted per *executing attempt*, incremented and decremented inside one + /// scope by `WalletAttemptGuard`. Nothing crosses the job-queue boundary, + /// so a counter cannot drift: a restart zeroes it, which is accurate, and + /// every early return still releases because `Drop` runs. + inflight: Vec, } impl KeyPool { pub fn new(keys: Vec) -> Self { + let inflight = keys.iter().map(|_| AtomicUsize::new(0)).collect(); Self { keys, cursor: AtomicUsize::new(0), + inflight, } } @@ -318,6 +339,73 @@ impl KeyPool { } } + /// Returns the least-loaded key index, breaking ties in round-robin order. + /// + /// Join-shortest-queue rather than `next_index()`'s blind modulo: a wallet + /// that is mid-upload is skipped in favour of an idle one, which is the + /// whole point — the per-wallet limit is 1, so landing on a busy wallet + /// means queueing behind it even when the pool has capacity. + /// + /// The tie-break matters as much as the minimum. On an idle pool every + /// counter is 0, and a plain `argmin` would return index 0 every time, + /// funnelling all traffic onto one wallet — strictly worse than the + /// round-robin this replaces. Starting the scan at the rotating cursor + /// keeps equally-loaded keys spreading exactly as before. + pub fn least_loaded_index(&self) -> Option { + let len = self.keys.len(); + if len == 0 { + return None; + } + let start = self.cursor.fetch_add(1, Ordering::Relaxed) % len; + let mut best = start; + let mut best_load = self.inflight[start].load(Ordering::Relaxed); + for step in 1..len { + let idx = (start + step) % len; + let load = self.inflight[idx].load(Ordering::Relaxed); + // Strictly less, so the first key scanned — the cursor's own — + // wins any tie and the rotation is preserved. + if load < best_load { + best = idx; + best_load = load; + } + } + Some(best) + } + + /// Marks `index` busy for as long as the returned guard lives. + /// + /// Call this around the wallet transaction itself, not around enqueueing + /// one: the counter is meant to answer "is this wallet signing right now", + /// which is what the per-wallet limit actually serialises on. + pub fn begin_attempt(self: &Arc, index: usize) -> WalletAttemptGuard { + if let Some(slot) = self.inflight.get(index) { + slot.fetch_add(1, Ordering::Relaxed); + } + WalletAttemptGuard { + pool: Arc::clone(self), + index, + } + } + + /// In-flight count per key. Observability only. + pub fn inflight_snapshot(&self) -> Vec { + self.inflight + .iter() + .map(|slot| slot.load(Ordering::Relaxed)) + .collect() + } + + fn end_attempt(&self, index: usize) { + if let Some(slot) = self.inflight.get(index) { + // Saturating: an extra release must not wrap to usize::MAX and + // leave this key looking permanently busiest, which would exclude + // it from selection for the life of the process. + let _ = slot.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |v| { + Some(v.saturating_sub(1)) + }); + } + } + #[allow(dead_code)] pub fn is_empty(&self) -> bool { self.keys.is_empty() @@ -329,6 +417,22 @@ impl KeyPool { } } +/// Releases a wallet's in-flight count when dropped. +/// +/// A guard rather than paired calls because `execute_wallet_job` returns from +/// many points; every one of them has to decrement, and `Drop` is the only way +/// to get that for free — including while unwinding from a panic. +pub struct WalletAttemptGuard { + pool: Arc, + index: usize, +} + +impl Drop for WalletAttemptGuard { + fn drop(&mut self) { + self.pool.end_attempt(self.index); + } +} + // ============================================================ // Config // ============================================================ @@ -928,10 +1032,13 @@ fn env_bool(name: &str) -> bool { .unwrap_or(false) } -/// `/health` `writes` wire value: `"paused"` when `WRITES_PAUSED` is set. -pub(crate) fn writes_health_status(paused: bool) -> String { +/// `/health` `writes` wire value: `"paused"`, `"degraded"`, or `"ok"`. +/// Paused wins if both flags are set — write routes already 503. +pub(crate) fn writes_health_status(paused: bool, degraded: bool) -> String { if paused { "paused".to_string() + } else if degraded { + "degraded".to_string() } else { "ok".to_string() } @@ -1264,6 +1371,8 @@ pub struct RememberBulkItem { pub struct RememberBulkRequest { /// 1–MAX_BULK_ITEMS items to remember in one batched operation. pub items: Vec, + #[serde(default)] + pub idempotency_key: Option, } /// POST /api/remember/bulk — 202 Accepted response. @@ -1414,8 +1523,16 @@ pub struct RecallRequest { pub scoring_weights: Option, /// How to order results. Omitted → [`RecallSort::Relevance`], today's /// behaviour. See [`RecallSort`]. + /// + /// `Option` because an explicit `sort`, `relevance` included, suppresses + /// `scoring_weights`, so omitted and explicit must stay distinct. #[serde(default)] - pub sort: RecallSort, + pub sort: Option, + /// How long the caller waits, in ms. When set, the recall stops just + /// short of it and answers [`AppError::RecallTimeout`] naming the stage + /// it was in; omitted, it runs to completion. + #[serde(default)] + pub deadline_ms: Option, } /// Result ordering mode for `/api/recall`. @@ -1883,6 +2000,24 @@ pub struct AccountExistsResponse { pub exists: bool, } +/// GET /api/whoami — the identity the caller's delegate key resolves to. +/// +/// Unlike `AccountExistsResponse` this *does* carry `account_id`, which is +/// safe here precisely because the route is authenticated: the caller proved +/// possession of a delegate key already registered against this account, so +/// it is being told its own identity, not anyone else's. +/// +/// Exists so a client that holds a delegate key but lost the surrounding +/// metadata can rebuild `credentials.json` (WALM-332). All three fields are +/// required for that: `account_id` and `owner` come from the registry scan, +/// `package_id` from server config. +#[derive(Debug, Serialize)] +pub struct WhoamiResponse { + pub account_id: String, + pub owner: String, + pub package_id: String, +} + /// POST /api/stats — count + stored bytes for a namespace. /// Used by the benchmark harness for verification. Mode-blind. #[derive(Debug, Deserialize)] @@ -1925,9 +2060,18 @@ pub struct HealthResponse { /// fail open so CI `wait-for-relayer` does not hang. `status` stays /// `"ok"` while the relayer process is up. pub write_ready: bool, - /// Write-path admission: `"ok"` or `"paused"`. `"paused"` when - /// `WRITES_PAUSED` is set; write routes then return HTTP 503. - /// Distinct from `write_ready`. `/health` stays HTTP 200. + /// Write-path state: `"ok"`, `"degraded"`, or `"paused"`. + /// + /// `"paused"` when `WRITES_PAUSED` is set; write routes then return + /// HTTP 503. `"degraded"` when recent durable writes have been failing + /// and none have landed -- the relayer still accepts and durably + /// queues a write, but Walrus is not storing it, so a caller should + /// expect the job to fail minutes later rather than queue more. + /// + /// Deliberately separate from `write_ready`, which stays true through + /// a downstream outage: CI's wait-for-relayer gate blocks on + /// `write_ready is True`, so folding this into it would make a Walrus + /// outage hang every deploy. `/health` stays HTTP 200 throughout. pub writes: String, } @@ -2111,6 +2255,12 @@ pub enum AppError { /// Operator write pause (`WRITES_PAUSED`). HTTP 503 with a stable /// client-visible message, distinct from transient upstream failures. WritesPaused(String), + /// A recall about to miss the caller's `deadline_ms`. HTTP 504 with + /// `code: "RECALL_TIMEOUT"` and the stage it was stuck in. + RecallTimeout { + stage: &'static str, + elapsed_ms: u64, + }, } impl std::fmt::Display for AppError { @@ -2123,6 +2273,9 @@ impl std::fmt::Display for AppError { AppError::Forbidden(msg) => write!(f, "Forbidden: {}", msg), AppError::Conflict(msg) => write!(f, "Conflict: {}", msg), AppError::RateLimited(msg) => write!(f, "Rate Limited: {}", msg), + AppError::RecallTimeout { stage, elapsed_ms } => { + write!(f, "Recall Timeout: {} after {}ms", stage, elapsed_ms) + } AppError::QuotaExceeded(msg) => write!(f, "Quota Exceeded: {}", msg), AppError::UpstreamUnavailable(msg) => write!(f, "Upstream Unavailable: {}", msg), AppError::WritesPaused(msg) => write!(f, "Writes Paused: {}", msg), @@ -2160,6 +2313,20 @@ impl axum::response::IntoResponse for AppError { AppError::WritesPaused(msg) => { (axum::http::StatusCode::SERVICE_UNAVAILABLE, msg.clone()) } + AppError::RecallTimeout { stage, elapsed_ms } => { + // Not `{error}` alone: the SDK reads `code` into + // `serverCode` and shows `message`, and callers read + // `stage` to decide what to do next. + let message = format!("Recall timed out after {elapsed_ms}ms during {stage}"); + let body = serde_json::json!({ + "error": message, + "message": message, + "code": "RECALL_TIMEOUT", + "stage": stage, + "elapsed_ms": elapsed_ms, + }); + return (axum::http::StatusCode::GATEWAY_TIMEOUT, axum::Json(body)).into_response(); + } AppError::UpstreamUnavailable(msg) => { // log the upstream details server-side, return // 503 so the SDK / harness will retry per their @@ -2198,6 +2365,7 @@ impl AppError { AppError::QuotaExceeded(_) => "quota_exceeded", AppError::UpstreamUnavailable(_) => "upstream_unavailable", AppError::WritesPaused(_) => "writes_paused", + AppError::RecallTimeout { .. } => "recall_timeout", } } } @@ -2223,6 +2391,126 @@ mod tests { static WALRUS_STORAGE_EPOCHS_ENV_LOCK: Mutex<()> = Mutex::new(()); + fn pool(n: usize) -> Arc { + Arc::new(KeyPool::new((0..n).map(|i| format!("key{i}")).collect())) + } + + #[test] + fn an_idle_pool_still_spreads_round_robin() { + // The regression this guards: with every counter at 0 a plain argmin + // returns index 0 forever, funnelling the whole pool onto one wallet + // — worse than the round-robin it replaces. + let pool = pool(4); + let picked: Vec<_> = (0..8).map(|_| pool.least_loaded_index().unwrap()).collect(); + assert_eq!(picked, vec![0, 1, 2, 3, 0, 1, 2, 3]); + } + + #[test] + fn a_busy_wallet_is_skipped_for_an_idle_one() { + // The production case: a job waited 6.2s for its assigned wallet while + // other wallets sat idle, because round-robin could not see the load. + let pool = pool(4); + let _busy = pool.begin_attempt(1); + // Cursor lands on 1 next, but 1 is busy and 2 is not. + let _ = pool.least_loaded_index(); + assert_eq!(pool.least_loaded_index().unwrap(), 2); + } + + #[test] + fn selection_avoids_every_busy_wallet_until_only_busy_ones_remain() { + let pool = pool(3); + let _a = pool.begin_attempt(0); + let _b = pool.begin_attempt(1); + // Only wallet 2 is free, so every pick goes there regardless of cursor. + for _ in 0..6 { + assert_eq!(pool.least_loaded_index().unwrap(), 2); + } + } + + #[test] + fn a_fully_busy_pool_falls_back_to_spreading_evenly() { + // Saturated is not a special case: equal load means the tie-break + // decides, so behaviour degrades exactly to round-robin. + let pool = pool(3); + let _a = pool.begin_attempt(0); + let _b = pool.begin_attempt(1); + let _c = pool.begin_attempt(2); + let picked: Vec<_> = (0..6).map(|_| pool.least_loaded_index().unwrap()).collect(); + assert_eq!(picked, vec![0, 1, 2, 0, 1, 2]); + } + + #[test] + fn dropping_the_guard_frees_the_wallet() { + let pool = pool(2); + { + let _busy = pool.begin_attempt(0); + assert_eq!(pool.inflight_snapshot(), vec![1, 0]); + } + assert_eq!(pool.inflight_snapshot(), vec![0, 0]); + } + + #[test] + fn the_guard_releases_on_an_early_return() { + // `execute_wallet_job` returns from many points; the guard exists so + // none of them has to remember to decrement. + let pool = pool(2); + fn bail(pool: &Arc) -> Result<(), ()> { + let _slot = pool.begin_attempt(1); + Err(()) + } + assert!(bail(&pool).is_err()); + assert_eq!(pool.inflight_snapshot(), vec![0, 0]); + } + + #[test] + fn nested_attempts_on_one_wallet_count_and_release_independently() { + let pool = pool(2); + let first = pool.begin_attempt(0); + let second = pool.begin_attempt(0); + assert_eq!(pool.inflight_snapshot(), vec![2, 0]); + drop(first); + assert_eq!(pool.inflight_snapshot(), vec![1, 0]); + drop(second); + assert_eq!(pool.inflight_snapshot(), vec![0, 0]); + } + + #[test] + fn releasing_below_zero_saturates_instead_of_wrapping() { + // A wrapped counter would read as usize::MAX and exclude the wallet + // from selection for the life of the process. + let pool = pool(2); + pool.end_attempt(0); + pool.end_attempt(0); + assert_eq!(pool.inflight_snapshot(), vec![0, 0]); + let _busy = pool.begin_attempt(1); + assert_eq!(pool.least_loaded_index().unwrap(), 0); + } + + #[test] + fn an_out_of_range_index_is_ignored_rather_than_panicking() { + let pool = pool(2); + let guard = pool.begin_attempt(99); + assert_eq!(pool.inflight_snapshot(), vec![0, 0]); + drop(guard); + assert_eq!(pool.inflight_snapshot(), vec![0, 0]); + } + + #[test] + fn an_empty_pool_selects_nothing() { + let pool = Arc::new(KeyPool::new(vec![])); + assert_eq!(pool.least_loaded_index(), None); + assert_eq!(pool.next_index(), None); + } + + #[test] + fn next_index_is_unchanged_for_the_retry_path() { + // Retries derive their wallet from the job's own start index, not the + // global cursor, so `next_index` has to keep its old semantics. + let pool = pool(3); + let picked: Vec<_> = (0..6).map(|_| pool.next_index().unwrap()).collect(); + assert_eq!(picked, vec![0, 1, 2, 0, 1, 2]); + } + #[test] fn balance_monitor_interval_has_a_safe_minimum() { assert_eq!(normalized_balance_monitor_interval(1), 30); @@ -3178,6 +3466,57 @@ mod tests { .unwrap(); } + // ── RecallRequest.sort — omitted vs explicit ───────────────────────── + + #[test] + fn recall_sort_keeps_omitted_apart_from_explicit_relevance() { + let parse = |body: &str| serde_json::from_str::(body).unwrap().sort; + assert_eq!(parse(r#"{"query":"q"}"#), None); + assert_eq!( + parse(r#"{"query":"q","sort":"relevance"}"#), + Some(RecallSort::Relevance) + ); + assert_eq!( + parse(r#"{"query":"q","sort":"recent"}"#), + Some(RecallSort::Recent) + ); + } + + // ── RecallRequest.deadline_ms / RECALL_TIMEOUT ─────────────────────── + + #[test] + fn recall_deadline_is_optional() { + let parse = |body: &str| { + serde_json::from_str::(body) + .unwrap() + .deadline_ms + }; + assert_eq!(parse(r#"{"query":"q"}"#), None); + assert_eq!(parse(r#"{"query":"q","deadline_ms":15000}"#), Some(15_000)); + } + + #[tokio::test] + async fn recall_timeout_is_a_504_that_names_the_stage() { + let err = AppError::RecallTimeout { + stage: "walrus_download", + elapsed_ms: 14_001, + }; + let resp = axum::response::IntoResponse::into_response(err); + assert_eq!(resp.status(), axum::http::StatusCode::GATEWAY_TIMEOUT); + let body = axum::body::to_bytes(resp.into_body(), 4096).await.unwrap(); + let body: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(body["code"], "RECALL_TIMEOUT"); + assert_eq!(body["stage"], "walrus_download"); + assert_eq!(body["elapsed_ms"], 14_001); + // `message` is what the TypeScript SDK shows; `error` is what every + // other relayer error carries. + assert_eq!(body["message"], body["error"]); + assert!(body["message"] + .as_str() + .unwrap() + .contains("walrus_download")); + } + // ── ScoringWeights::is_ranker_active() — opt-in predicate ──────────── #[test] @@ -3266,8 +3605,11 @@ mod tests { #[tokio::test] async fn writes_paused_maps_to_503_with_stable_message() { - assert_eq!(writes_health_status(false), "ok"); - assert_eq!(writes_health_status(true), "paused"); + assert_eq!(writes_health_status(false, false), "ok"); + assert_eq!(writes_health_status(false, true), "degraded"); + assert_eq!(writes_health_status(true, false), "paused"); + // An operator pause is the stronger statement and wins. + assert_eq!(writes_health_status(true, true), "paused"); assert!(reject_if_writes_paused(false).is_ok()); let err = reject_if_writes_paused(true).expect_err("paused writes"); assert_eq!(err.kind(), "writes_paused");