trending/ai-memory
1
0
Fork 0
mirror of https://github.com/akitaonrails/ai-memory.git synced 2026-10-02 03:24:46 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge main into #756 (resolve CHANGELOG against the stamped [2.5.1])

Browse Source 
# Conflicts:
#	CHANGELOG.md
This commit is contained in:
AkitaOnRails
2026-10-01 02:39:10 -03:00
parent 10b8746024 695805eb88
commit 02c414e534
382 changed files with 94857 additions and 3436 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/ci.yml
+87 -13
View File
@@ -71,11 +71,40 @@ jobs:
if: matrix.os == 'ubuntu-latest'
run: echo "TAILWIND_BUILD=1" >> "$GITHUB_ENV"
- run: cargo test --workspace --all-targets
# `--all-targets` deliberately excludes doctests, so a doc-comment code
# block that fails to compile (e.g. an indented shell example rustdoc
# reads as Rust) shipped unnoticed. Run them explicitly.
- run: cargo test --workspace --doc
- uses: astral-sh/setup-uv@bec219d24cd3e171d82865faccec33120bb574f4 # v10.1.0
with:
enable-cache: false
- name: Build the external lifecycle relay
env:
CARGO_TARGET_DIR: ${{ github.workspace }}/target
run: cargo build --locked --manifest-path companions/ai-memory-relay/Cargo.toml
- name: External lifecycle delivery against the real server
run: >
uv run --no-project python tests/e2e/external_relay_smoke.py
--ai-memory-bin "${{ github.workspace }}/target/debug/ai-memory"
--relay-bin "${{ github.workspace }}/target/debug/ai-memory-relay"
# The capture policy every generated OpenCode/OMP/Pi/OpenClaw plugin
# embeds is TypeScript; only Node can execute it against the shared
# capture-policy fixture. The test is ignored locally, where Node is not
# a build requirement.
- uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7.0.0
if: matrix.os == 'ubuntu-latest'
with:
node-version: "24"
- name: Generated capture-policy TypeScript runtime evidence
if: matrix.os == 'ubuntu-latest'
# `--workspace` reuses the test binaries built above instead of
# re-resolving features for one package.
run: cargo test --workspace --lib -- --ignored --exact commands::render_shared::tests::generated_capture_policy_v1_node_runtime_evidence
- name: Vendored tailwind.css is current
if: matrix.os == 'ubuntu-latest'
run: git diff --exit-code -- crates/ai-memory-web/static/tailwind.css
# Native Windows coverage lives in `.github/workflows/windows.yml`: it
# Native Windows *test* coverage lives in `.github/workflows/windows.yml`: it
# runs nightly, on demand, and on any pull request carrying the `windows`
# label. It was ~1000s here against ~250s for the
# same tests on Linux, which meant every pull request waited roughly
@@ -83,14 +112,51 @@ jobs:
# and it was `continue-on-error`, so those extra minutes gated nothing.
# Add the `windows` label to a pull request that touches path handling,
# file locking, or git plumbing.
#
# `windows-cross` below is the cheap per-merge complement: it cross-*builds*
# the Windows target from Linux, so a `#[cfg(windows)]` compile or link break
# is caught in this ~8-minute gate instead of at release time. It does NOT run
# tests — cross-compiling proves the code builds and links for Windows, not
# that it behaves correctly there. Wine is deliberately not used to fake a
# runtime: it cannot faithfully reproduce native PowerShell, NTFS
# case-folding, Win32 file-locking/sharing-violation semantics, or verbatim
# (`\\?\`) path handling — exactly the surface windows.yml (and the local
# dockur VM loop, see docs/design-windows-ci.md) exist to validate.
windows-cross:
name: windows cross-build (msvc)
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
# Pin to 1.95 to match rust-toolchain.toml: that file overrides the
# active toolchain inside the repo, so the MSVC target's std must be added
# to 1.95 (not `stable`) or the cross-build fails with "can't find crate
# for `core`".
- uses: dtolnay/rust-toolchain@4cda84d5c5c54efe2404f9d843567869ab1699d4 # 1.95
with:
toolchain: "1.95"
targets: x86_64-pc-windows-msvc
- uses: Swatinem/rust-cache@f0d9c3887740aee45f6153b24b3a6b815192ec16 # v2
with:
key: windows-cross
# cargo-xwin drives clang-cl / lld-link against an auto-downloaded MSVC
# CRT + Windows SDK, so bundled SQLite and vendored libgit2 (both C) build
# for the target without a Windows host.
- name: Install clang/lld toolchain
run: sudo apt-get update && sudo apt-get install -y --no-install-recommends clang llvm lld
- name: Install cargo-xwin
run: cargo install cargo-xwin --locked
# --all-targets so the `#[cfg(windows)]` test binaries compile too — that
# is the class of break (a test that only compiles on Windows) this gate
# is here to catch before it reaches windows.yml or a release.
- run: cargo xwin build --workspace --all-targets --target x86_64-pc-windows-msvc
# The companion importer is deliberately outside the root workspace
# (its own [workspace] in companions/ai-memory-importer/Cargo.toml), so
# none of the jobs above compile it. Gate it here with the same
# fmt/clippy/test trio docs/companion-crates.md prescribes, or a
# toolchain bump / convention change rots it silently.
# Companions have their own workspaces and need separate checks.
companions:
name: companions (ai-memory-importer)
name: companions (${{ matrix.companion }})
strategy:
fail-fast: false
matrix:
companion: [ai-memory-importer, ai-memory-relay]
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
@@ -99,10 +165,10 @@ jobs:
components: rustfmt, clippy
- uses: Swatinem/rust-cache@f0d9c3887740aee45f6153b24b3a6b815192ec16 # v2
with:
workspaces: companions/ai-memory-importer
- run: cargo fmt --check --manifest-path companions/ai-memory-importer/Cargo.toml
- run: cargo clippy --manifest-path companions/ai-memory-importer/Cargo.toml --all-targets -- -D warnings
- run: cargo test --manifest-path companions/ai-memory-importer/Cargo.toml
workspaces: companions/${{ matrix.companion }}
- run: cargo fmt --check --manifest-path companions/${{ matrix.companion }}/Cargo.toml
- run: cargo clippy --locked --manifest-path companions/${{ matrix.companion }}/Cargo.toml --all-targets -- -D warnings
- run: cargo test --locked --manifest-path companions/${{ matrix.companion }}/Cargo.toml
# Build the release binary on its own so a release-only failure
# (LTO crash, codegen issue, dead-code-with-debug-assertions) doesn't
@@ -156,6 +222,7 @@ jobs:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
- run: sudo apt-get update && sudo apt-get install -y rpm
- run: scripts/check-native-packaging.sh
- run: tests/wrapper_upgrade.sh
@@ -266,6 +333,13 @@ jobs:
log-level: warn
command: check
arguments: --all-features
- name: Relay dependency policy
uses: EmbarkStudios/cargo-deny-action@3c6349835b2b7b196a839186cb8b78e02f7b5f25 # v2
with:
log-level: warn
command: check
manifest-path: companions/ai-memory-relay/Cargo.toml
arguments: --all-features
audit:
name: cargo-audit
@@ -279,8 +353,8 @@ jobs:
cargo audit
--ignore RUSTSEC-2025-0141
--ignore RUSTSEC-2024-0320
--ignore RUSTSEC-2026-0194
--ignore RUSTSEC-2026-0195
- run: cargo audit --file companions/ai-memory-relay/Cargo.lock
- run: cargo audit --file companions/ai-memory-importer/Cargo.lock
gitleaks:
name: gitleaks (secret scan)
.github/workflows/macos-app.yml
+22
View File
@@ -0,0 +1,22 @@
name: macos-app
on:
pull_request:
types: [opened, synchronize, reopened, labeled]
workflow_dispatch:
permissions:
contents: read
jobs:
swift-test:
name: swift test (ai-memory-macos)
if: >-
github.event_name == 'workflow_dispatch' ||
contains(github.event.pull_request.labels.*.name, 'macos') ||
contains(github.event.pull_request.labels.*.name, 'full-ci')
runs-on: macos-latest
steps:
- uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
- name: Swift tests
run: swift test --package-path companions/ai-memory-macos
.github/workflows/release.yml
+58 -3
View File
@@ -241,7 +241,12 @@ jobs:
# etc. land at the zip root, mirroring the Linux tarball layout.
Get-ChildItem $stage | Compress-Archive -DestinationPath $zip
$hash = (Get-FileHash $zip -Algorithm SHA256).Hash.ToLower()
"$hash $zip" | Out-File -Encoding ascii "$zip.sha256"
# -NoNewline plus an explicit `n, because Out-File's own line
# terminator is CRLF on Windows. `sha256sum -c` treats the CR as part
# of the filename and fails with "No such file or directory", and the
# release body concatenates every platform's file into one block, so
# the stray byte shows up there too.
"$hash $zip`n" | Out-File -Encoding ascii -NoNewline "$zip.sha256"
- name: Smoke test release zip
shell: pwsh
@@ -264,7 +269,7 @@ jobs:
}
}
$checksum = Get-Content "$zip.sha256" -Raw
if ($checksum -notmatch '^[0-9a-f]{64} ai-memory-windows-x86_64\.zip\r?\n?$') {
if ($checksum -cnotmatch '^[0-9a-f]{64} ai-memory-windows-x86_64\.zip\n$') {
throw "Windows release zip checksum is not sha256sum format"
}
@@ -483,7 +488,7 @@ jobs:
github-release:
name: github release
needs: [binary, macos, windows, validate-version]
needs: [binary, macos, windows, rpm, validate-version]
runs-on: ubuntu-latest
permissions:
contents: write
@@ -533,6 +538,8 @@ jobs:
echo "yay -S ai-memory-bin # prebuilt Linux x86_64/aarch64 binary" >> body.md
echo "yay -S ai-memory # builds from source" >> body.md
echo "" >> body.md
echo "# Fedora: download the matching RPM below, then run sudo dnf install ./ai-memory-<version>-1.<arch>.rpm" >> body.md
echo "" >> body.md
echo "# Docker" >> body.md
echo "docker pull akitaonrails/ai-memory:${{ needs.validate-version.outputs.version }}" >> body.md
echo "" >> body.md
@@ -661,3 +668,51 @@ jobs:
publish_pkg ai-memory "$source_pkgbuild"
publish_pkg ai-memory-bin "$bin_pkgbuild"
EOF
rpm:
name: Fedora RPM (${{ matrix.arch }})
needs: [binary, validate-version]
runs-on: ${{ matrix.runner }}
strategy:
matrix:
include:
- arch: x86_64
runner: ubuntu-24.04
- arch: aarch64
runner: ubuntu-24.04-arm
steps:
- uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
- uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1
with:
name: ai-memory-linux-${{ matrix.arch }}
path: artifacts
- name: Build RPM
env:
VERSION: ${{ needs.validate-version.outputs.version }}
ARCH: ${{ matrix.arch }}
run: |
sudo apt-get update
sudo apt-get install -y cpio rpm
mkdir -p "$HOME/rpmbuild/SOURCES" "$HOME/rpmbuild/SPECS" "$HOME/rpmbuild/RPMS"
cp artifacts/ai-memory-linux-${ARCH}.tar.gz "$HOME/rpmbuild/SOURCES/"
cp packaging/rpm/ai-memory.spec "$HOME/rpmbuild/SPECS/"
sed -i "s/^Version: .*/Version: ${VERSION}/" "$HOME/rpmbuild/SPECS/ai-memory.spec"
rpmbuild -bb "$HOME/rpmbuild/SPECS/ai-memory.spec"
rpm_file="$(find "$HOME/rpmbuild/RPMS" -name '*.rpm' -print -quit)"
test -n "$rpm_file"
rpm -qp --queryformat '%{VERSION} %{ARCH}\n' "$rpm_file" | grep -Fx "$VERSION $ARCH"
rpm -ql "$rpm_file" | grep -Fx /usr/lib/systemd/system/ai-memory.service
root="$(mktemp -d)"
trap 'rm -rf "$root"' EXIT
rpm2cpio "$rpm_file" | (cd "$root" && cpio -idm --quiet)
"$root/usr/bin/ai-memory" --version | grep -Fx "ai-memory $VERSION"
rpm_name="ai-memory-$(rpm -qp --queryformat '%{VERSION}-%{RELEASE}.%{ARCH}' "$rpm_file").rpm"
mv "$rpm_file" "artifacts/$rpm_name"
(cd artifacts && sha256sum "$rpm_name" > "$rpm_name.sha256")
- uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: ai-memory-rpm-${{ matrix.arch }}
path: artifacts/*.rpm*
.github/workflows/windows.yml
+117 -9
View File
@@ -19,17 +19,36 @@ name: windows
#
# * a nightly run catches toolchain and dependency drift that no code
# change would trigger;
# * `workflow_dispatch` and the `windows` label give a pull request that
# touches platform-sensitive code — path handling, file locking, git
# plumbing — a way to opt in before merging;
# * a pull request that touches platform-sensitive code — path handling,
# file locking, git plumbing, the hook bundle — opts in AUTOMATICALLY:
# the leading `changes` job (see below) detects those paths and the
# Windows jobs run without anyone remembering a label;
# * `workflow_dispatch` and the `windows` label remain manual overrides —
# the label forces the Windows jobs on a pull request that falls OUTSIDE
# those paths (e.g. a dependency bump someone wants to vet on Windows);
# * the job no longer carries `continue-on-error`, so failures in those
# scheduled, labelled, and manual runs are red rather than advisory.
# scheduled, labelled, path-triggered, and manual runs are red rather
# than advisory.
#
# If you are changing any of those areas, add the `windows` label to the
# pull request. During feature iteration this workflow does not run on pushes
# A pull request that changes those areas now runs Windows on its own; the
# `windows` label is only needed to force a run for a PR outside them.
# During feature iteration this workflow does not run on pushes
# to main: fast Linux CI gates each merge, and the full Windows suite runs
# nightly, on demand, and MANDATORILY right before a release (dispatch it
# on the release-candidate SHA and wait for green — see AGENTS.md).
#
# The `hooks` job is separate so a Rust failure cannot stop the hook suite from
# reporting, but it follows the same gate as `test` on pull requests (the
# path-change auto-trigger plus the `windows` label override). Both jobs also
# run nightly and on manual dispatch.
#
# The `pull_request` trigger is deliberately kept WITHOUT a `paths:` filter.
# A `paths:` filter on the event would make the whole workflow fire only for
# path-matching PRs — which would also drop the `labeled` event for a PR
# outside those paths and so break the manual `windows`-label override. The
# path gate is instead computed by the leading `changes` job below, and each
# job's `if:` combines it with the label override; all PRs still enter the
# workflow so the `labeled` type keeps working.
on:
pull_request:
types: [opened, synchronize, reopened, labeled]
@@ -40,6 +59,9 @@ on:
permissions:
contents: read
# `dorny/paths-filter` lists a pull request's changed files through the
# GitHub API on `pull_request` events; that read scope is what it needs.
pull-requests: read
concurrency:
# A newer push supersedes an in-flight run for the same ref. Nothing
@@ -53,16 +75,102 @@ env:
RUSTFLAGS: "-D warnings"
jobs:
# Compute whether a pull request touches platform-sensitive code, so the
# Windows jobs auto-run on those PRs without a human remembering the label.
# It runs on every event (a cheap Linux job) so the `needs:` edge below is
# always satisfied; the filter step only fires on `pull_request`, where its
# `platform` output drives the gate. On schedule / manual dispatch the output
# is empty and unused, because those jobs' `if:` short-circuits on
# `github.event_name != 'pull_request'`.
changes:
name: detect platform-sensitive changes
runs-on: ubuntu-latest
outputs:
platform: ${{ steps.filter.outputs.platform }}
steps:
# For `pull_request` events this action reads the changed-file list from
# the GitHub API, so no checkout is needed. Pinned to a full commit SHA
# with a version comment, matching this repo's third-party-action policy.
- uses: dorny/paths-filter@ceb8a2b8f2d89434be7ff52d3de7ec3738c5cc9d # v4.0.3
id: filter
if: github.event_name == 'pull_request'
with:
# Path handling, file locking, git plumbing, and the hook bundle —
# the areas AGENTS.md calls out as needing Windows coverage.
filters: |
platform:
- 'crates/ai-memory-wiki/**'
- 'crates/ai-memory-store/**'
- 'crates/ai-memory-hooks/**'
- 'crates/ai-memory-cli/src/commands/install_hooks.rs'
- 'hooks/**'
- 'tests/hooks/**'
- '.github/workflows/windows.yml'
test:
name: test (windows-latest)
# On a pull request, only when explicitly opted in. The schedule and
# manual dispatch always run.
needs: changes
# On a pull request, run when a platform-sensitive path changed OR the
# `windows` label was applied (the manual override for PRs outside those
# paths). The schedule and manual dispatch always run.
if: >-
github.event_name != 'pull_request' ||
contains(github.event.pull_request.labels.*.name, 'windows')
contains(github.event.pull_request.labels.*.name, 'windows') ||
needs.changes.outputs.platform == 'true'
runs-on: windows-latest
steps:
- uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
- uses: dtolnay/rust-toolchain@4cda84d5c5c54efe2404f9d843567869ab1699d4 # stable
- uses: Swatinem/rust-cache@f0d9c3887740aee45f6153b24b3a6b815192ec16 # v2
- run: cargo test --workspace --all-targets
- name: Check the external lifecycle relay
env:
CARGO_TARGET_DIR: ${{ github.workspace }}/target
run: |
cargo fmt --check --manifest-path companions/ai-memory-relay/Cargo.toml
if ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE }
cargo clippy --locked --manifest-path companions/ai-memory-relay/Cargo.toml --all-targets -- -D warnings
if ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE }
cargo test --locked --manifest-path companions/ai-memory-relay/Cargo.toml
if ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE }
cargo build --locked --manifest-path companions/ai-memory-relay/Cargo.toml
if ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE }
- uses: astral-sh/setup-uv@bec219d24cd3e171d82865faccec33120bb574f4 # v10.1.0
with:
enable-cache: false
- name: External lifecycle delivery against the real server
run: >
uv run --no-project python tests/e2e/external_relay_smoke.py
--ai-memory-bin "${{ github.workspace }}/target/debug/ai-memory.exe"
--relay-bin "${{ github.workspace }}/target/debug/ai-memory-relay.exe"
--timeout-seconds 120
# `ci.yml`'s `hooks-shell` job covers the POSIX bundle across four awks, all
# on Linux. The suite also drives `hooks/lib/ai-memory-hook.ps1`, and that
# half only ever executes where PowerShell is native — so the one platform
# its PowerShell branch is written for was the one platform no job ran it on.
#
# Its own job so a Rust failure cannot stop the shell suite from reporting,
# but gated the same way as `test` above. The project keeps Windows off
# per-PR feedback by default (see the header's cost argument), so a pull
# request touching a platform-sensitive path — `hooks/` and `tests/hooks/`
# among them — opts in automatically via the `changes` job; the `windows`
# label forces it for a PR outside those paths. It also runs nightly and on
# manual dispatch. (The job is cheap — a checkout and a few seconds of `sh` —
# but consistency with the fast-CI-per-merge rule wins over an every-PR
# carve-out.)
hooks:
name: hook bundle (windows-latest)
needs: changes
# Same gate as `test`: platform-path change OR the `windows` label on a
# pull request; schedule and manual dispatch always run.
if: >-
github.event_name != 'pull_request' ||
contains(github.event.pull_request.labels.*.name, 'windows') ||
needs.changes.outputs.platform == 'true'
runs-on: windows-latest
steps:
- uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
- name: hook bundle (Git Bash)
shell: bash
run: sh tests/hooks/test_lib.sh
.gitignore
+3
View File
@@ -1,6 +1,9 @@
# Rust build artefacts.
/target/
/companions/*/target/
/companions/*/.build/
/companions/*/.swiftpm/
/companions/*/dist/
/dist/
**/*.rs.bk
**/*.rs.orig
.gitleaks.toml
+6
View File
@@ -55,6 +55,12 @@ regexes = [
'''123456:ABC-DEF1234ghIkl-zyx57W2v1u123ew11''',
'''ABC-DEF1234ghIkl-zyx57W2v1u123ew11''',
'''Bearer abcdef0123456789ABCDEF0123456789''',
# Fixtures for the JSON-quoted value, scheme-prefixed header and
# unprefixed-assignment tests: a JSON `api_key`, a `client-token`, an
# `authorization: Token`, an Azure `AccountKey=` and an npm `_authToken=`.
# `generic-api-key` extracts the value alone, so the FAKE marker lives
# inside the value rather than in the key.
'''FAKEfake0123456789[A-Za-z0-9]*={0,2}''',
'''xoxb-1234567890-abcdefghij''',
# ---- Canary string used by the e2e test in tests/e2e/ ----
AGENTS.md
+70 -10
View File
@@ -197,6 +197,9 @@ evals/ live A/B harness; workspace member, not shipped.
companions/ai-memory-importer/ standalone OMC + external-conversation importer; NOT a root
workspace member — build/test it with
`--manifest-path companions/ai-memory-importer/Cargo.toml`.
companions/ai-memory-macos/ Swift menu bar wrapper; NOT a root workspace member —
`swift test --package-path companions/ai-memory-macos`
and `./companions/ai-memory-macos/build.sh`.
hooks/ per-agent lifecycle hook bundles (shell/native).
bin/ host wrapper scripts (`ai-memory`, `deploy`, `release`).
docker/ Dockerfile, compose files, TLS proxy templates.
@@ -272,6 +275,10 @@ no tiers.
`cargo test --manifest-path companions/ai-memory-importer/Cargo.toml`
(plus fmt/clippy on the same manifest). Root `--workspace` commands do
not cover it.
- Run the macOS menu bar companion separately:
`swift test --package-path companions/ai-memory-macos`
(plus `./companions/ai-memory-macos/build.sh` to stage `AI Memory.app`).
Root `--workspace` commands do not cover it.
### Platform notes
@@ -315,8 +322,10 @@ no tiers.
Linux — keeping it out of `ci.yml`
is what holds PR feedback near the eight minutes the gating jobs take.
**Add the `windows` label** to a PR touching path handling, file
locking, or git plumbing, so the check runs before the merge rather
than after it.
locking, git plumbing, or the hook bundle, so the corresponding Windows
jobs run before the merge rather than only on the nightly schedule. Both
the Rust test job and the hook-bundle job use this label gate on pull
requests; they also run on manual dispatch.
## Code style guidelines
@@ -440,16 +449,37 @@ Additional boundary rules:
fails), and build `git2` signatures with a fixed `Signature::now(...)`,
never `repo.signature()`. CI cannot catch either — its runners have no
global gitconfig, so the breakage only ever shows up on a developer's box.
- PRs touching scope resolution need table-driven tests for partial
scope, missing explicit scope, active-project precedence, and
cross-workspace isolation.
- PRs touching permissions need tests for root, DB-user, and anonymous
behavior.
- **Security-boundary tests are adversarial and mandatory.**
[`docs/security-boundaries.md`](docs/security-boundaries.md) is the inventory
of every isolation/security guard (per-project and workspace isolation, the
multi-user auth ladder, handoff single-claim + `any_owner` gate, pages-shared/
`author_id`-never-a-read-filter and supersession per invariant #16, the
active-project pointer, the sanitizer boundary, messaging scope, scope-
resolution fail-closed, destructive-op guards, hook backpressure, network
posture), each mapped to the test that would fail if the guard were removed.
Whenever you touch code in a boundary's "Enforcing code" column — or add a
new read/write/admin/hook/cross-scope entry point past one of these guards —
you MUST add or extend an **adversarial** test (attempt the violation, assert
refusal, include a legitimate control) and update that file's row in the same
change. Adding a new isolation dimension means a new row + its tests before
merge. A raw-id or unscoped/cross-project entry point (bare `session_id`/
`run_id`/`page_id`/message id, `global=true`, global `recent`, search) is
guilty until a test proves a foreign id/scope is refused. Prove the test
bites: it must fail with the guard removed and pass with it — a happy-path or
single-tenant test cannot see an isolation defect and does not count. These
guards live at integration level (`multi_session.rs`, `handoff_ownership.rs`,
`agent_messages.rs`, the active-project pointer tests, the MCP permission
suites); put boundary tests there. This subsumes the older "scope resolution"
and "permissions" test rules: PRs touching scope resolution still need
table-driven tests for partial scope, missing explicit scope, active-project
precedence, and cross-workspace isolation; PRs touching permissions still need
root, DB-user, and anonymous cases — now recorded against the inventory.
- New disk+SQL mutations need recovery/rollback tests.
- The recall-eval framework lives at
`crates/ai-memory-consolidate/tests/recall_eval.rs`.
- Tests run with `cargo t` locally and `cargo test --workspace --all-targets`
in CI.
in CI (plus `cargo test --workspace --doc`, which `--all-targets` excludes,
for doc-comment code blocks).
## Security considerations
@@ -487,6 +517,21 @@ Additional boundary rules:
`Added`/`Changed`/`Fixed` heading, past-tense, trailing `(#NNN)`
reference) and update the relevant README/docs references in the same
commit. Internal refactors and test-only churn are exempt.
- **Competitor research keeps the comparison docs in sync — never let them
go stale.** Any new competitor research pass, or a correction to an existing
one, must land its findings in the comparison docs in the *same* change, not
just in a research note: update `docs/comparison.md` (the public camp table +
positioning + "coming from …" migration notes), `docs/research-2026-landscape.md`
(the §3 camp entry + §6 sources, appending per the no-standalone-doc
convention), and `docs/competitive-parity.md` (the migration verdict + the
"did we copy without improving?" audit) wherever the finding applies. When a
competitor is reclassified or a claim is corrected, fix the camp table *and*
every per-tool claim that repeats it — a benchmark number, a "not file-first",
a camp label. This is a recurring failure: the Sept-2026 parity audit found
Supermemory mislabeled as a fact extractor, agentmemory's `0.967` attributed
to the wrong benchmark, and basic-memory's shipped reranking/Teams unrecorded.
Treat a stale claim in `comparison.md` (the doc that promises to be *fair*) as
a defect, not a nicety.
- **CI pacing: fast per merge, full matrix before release.** Every
implementation merge gates on the fast Linux jobs only. The slow
macOS/Windows legs run on a `full-ci` PR label, nightly (windows), or
@@ -504,6 +549,16 @@ Additional boundary rules:
wrong hash makes `brew install` fail for everyone. This is a mandatory,
recurring post-release step (it has been forgotten repeatedly); do not rely
on a contributor PR to the tap to remember it.
- **Reclaim build storage after every release — run `cargo clean`.** The
multi-worktree, multi-target-dir release flow (integration worktrees, per-agent
worktrees, separate `CARGO_TARGET_DIR`s) leaves many stale 100–180 MB test
binaries behind and has exhausted disk (see the `target/` bloat note under
Platform notes). Once a release is tagged and its artifacts are published, run
`cargo clean` (and `cargo clean` in each release worktree / extra target dir
you created, then remove finished `git worktree`s). This is a mandatory
post-release cleanup step, not optional housekeeping — treat it like the tap
bump above. During normal development, `cargo sweep --time 7` weekly is the
lighter-touch equivalent.
- **No version bumps or release tags without explicit user approval.**
Do not bump crate/package versions automatically.
- **PR evaluation:** report pros, cons, and recommended fix, then ask for
@@ -537,12 +592,17 @@ Additional boundary rules:
- [`docs/install.md`](docs/install.md) — installation cookbook for every
supported agent client.
- [`docs/cookbook.md`](docs/cookbook.md) — task-oriented cheat sheet: "I want
to do X" → how (recall, durable rules, importing a knowledge base, two agents
working together).
to do X" → how (recall, durable rules and must-read docs, importing a
knowledge base, retention, two agents working together, several accounts or
an external launcher via `run --env`, the Mac app, CLI, troubleshooting).
- [`docs/comparison.md`](docs/comparison.md) — fair, user-facing comparison
against other memory tools (camps, migration notes, how the field validates
the file-first/pages-over-facts approach). Analysis behind it:
`research-2026-landscape.md`.
- [`docs/competitive-parity.md`](docs/competitive-parity.md) — self-critical
internal audit: per-competitor migration-worthiness (do we do the basics +
add enough to justify switching?), the "did we copy without improving?"
borrowed-ideas verdicts, and documented gap-fill recommendations.
- [`docs/lifecycle-ops.md`](docs/lifecycle-ops.md) — read before touching
purge/rename/backup/restore/reset/reindex/restore-page.
- [`docs/auto-improvement-loop.md`](docs/auto-improvement-loop.md) —
CHANGELOG.md
+1771 -6
View File
File diff suppressed because it is too large Load Diff
CONTRIBUTING.md
+26 -3
View File
@@ -77,9 +77,32 @@ Skipped tests still count as "skipped" in the summary, never hidden, and two
independent things run them anyway: the pre-push hook and CI.
Install the hook once per clone with `scripts/install-git-hooks.sh` (from Git
Bash on Windows). It appends or updates only ai-memory's managed block in
`.git/hooks/pre-push`, preserving any existing hook body. Bypass it on a
work-in-progress branch with `git push --no-verify`.
Bash on Windows). It can run from the main checkout or a linked worktree;
both use the shared repository hook directory. Reinstallation replaces
ai-memory's managed block in place, preserving surrounding user commands and
their order. If `core.hooksPath` is set, the installer stops before writing;
integrate the block through your existing hook manager instead. Incomplete or
duplicate managed markers also stop installation and leave the hook unchanged.
Bypass it on a work-in-progress branch with `git push --no-verify`.
The managed test block clears Git's repository environment and disables global
and system Git configuration for Cargo and its children. Fixture commands can
then use their own repositories without inheriting the checkout being pushed.
The publishing Git process and other hook code retain their configuration, and
the block's shell options stay inside it; a failing test run still fails the
hook even when your own commands follow the block without `set -e`.
Run the installer again to update an existing installation.
The managed test block clears Git's repository environment and disables global
and system Git configuration for Cargo and its children. Fixture commands can
then use their own repositories without inheriting the checkout being pushed.
The publishing Git process and other hook code retain their configuration.
Run the installer again to update an existing installation.
Companions have separate Cargo workspaces. Check each changed companion with
`cargo fmt`, `cargo clippy --all-targets -- -D warnings`, and `cargo test`, passing
its `--manifest-path`. Changes to the lifecycle relay also need the real-server
test documented in [its README](companions/ai-memory-relay/README.md#validation).
Integration tests live in `tests/suite/` per crate and compile into the
crate's own test harness (declare a new file with `mod name;` in
Cargo.lock
Generated
+63 -18
View File
@@ -33,7 +33,7 @@ dependencies = [
[[package]]
name = "ai-memory-cli"
version = "2.3.1"
version = "2.5.1"
dependencies = [
"ai-memory-consolidate",
"ai-memory-core",
@@ -60,11 +60,13 @@ dependencies = [
"reqwest 0.12.28",
"rmcp",
"rstest",
"rusqlite",
"rustls",
"secrecy",
"serde",
"serde_json",
"sha2",
"socket2",
"sysinfo",
"tar",
"tempfile",
@@ -77,11 +79,12 @@ dependencies = [
"tracing-appender",
"tracing-subscriber",
"uuid",
"zip",
]
[[package]]
name = "ai-memory-consolidate"
version = "2.3.1"
version = "2.5.1"
dependencies = [
"ai-memory-core",
"ai-memory-llm",
@@ -91,7 +94,9 @@ dependencies = [
"anyhow",
"async-trait",
"git2",
"icu_normalizer",
"jiff",
"regex",
"rusqlite",
"schemars",
"serde",
@@ -106,8 +111,9 @@ dependencies = [
[[package]]
name = "ai-memory-core"
version = "2.3.1"
version = "2.5.1"
dependencies = [
"icu_normalizer",
"jiff",
"regex",
"schemars",
@@ -121,7 +127,7 @@ dependencies = [
[[package]]
name = "ai-memory-eval"
version = "2.3.1"
version = "2.5.1"
dependencies = [
"ai-memory-consolidate",
"ai-memory-core",
@@ -130,6 +136,7 @@ dependencies = [
"clap",
"jiff",
"reqwest 0.12.28",
"schemars",
"secrecy",
"serde",
"serde_json",
@@ -143,7 +150,7 @@ dependencies = [
[[package]]
name = "ai-memory-hooks"
version = "2.3.1"
version = "2.5.1"
dependencies = [
"ai-memory-consolidate",
"ai-memory-core",
@@ -168,7 +175,7 @@ dependencies = [
[[package]]
name = "ai-memory-llm"
version = "2.3.1"
version = "2.5.1"
dependencies = [
"ai-memory-core",
"anyhow",
@@ -197,7 +204,7 @@ dependencies = [
[[package]]
name = "ai-memory-mcp"
version = "2.3.1"
version = "2.5.1"
dependencies = [
"ai-memory-consolidate",
"ai-memory-core",
@@ -214,7 +221,9 @@ dependencies = [
"http",
"jiff",
"rmcp",
"rusqlite",
"schemars",
"secrecy",
"serde",
"serde_json",
"subtle",
@@ -230,7 +239,7 @@ dependencies = [
[[package]]
name = "ai-memory-store"
version = "2.3.1"
version = "2.5.1"
dependencies = [
"ai-memory-core",
"anyhow",
@@ -256,11 +265,11 @@ dependencies = [
[[package]]
name = "ai-memory-test-support"
version = "2.3.1"
version = "2.5.1"
[[package]]
name = "ai-memory-web"
version = "2.3.1"
version = "2.5.1"
dependencies = [
"ai-memory-core",
"ai-memory-store",
@@ -283,7 +292,7 @@ dependencies = [
[[package]]
name = "ai-memory-wiki"
version = "2.3.1"
version = "2.5.1"
dependencies = [
"ai-memory-core",
"ai-memory-llm",
@@ -315,7 +324,7 @@ dependencies = [
[[package]]
name = "ai-memory-workstream"
version = "2.3.1"
version = "2.5.1"
dependencies = [
"ai-memory-core",
"anyhow",
@@ -1339,6 +1348,7 @@ checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c"
dependencies = [
"crc32fast",
"miniz_oxide",
"zlib-rs",
]
[[package]]
@@ -1956,6 +1966,9 @@ dependencies = [
"icu_properties",
"icu_provider",
"smallvec",
"utf16_iter",
"utf8_iter",
"write16",
"zerovec",
]
@@ -3219,9 +3232,9 @@ dependencies = [
[[package]]
name = "rmcp"
version = "1.7.0"
version = "2.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0810a9f717d9828f475fe1f629f4c305c8464b7f496c3a854b58d29e65f4058e"
checksum = "14db48ee17a9ba61810ab1a9c1beb7d06d8136ae39ac25a1137f10d357af01af"
dependencies = [
"async-trait",
"base64 0.22.1",
@@ -3251,9 +3264,9 @@ dependencies = [
[[package]]
name = "rmcp-macros"
version = "1.7.0"
version = "2.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6aefac48c364756e97f04c0401ba3231e8607882c7c1d92da0437dc16307904d"
checksum = "783d787bf21813b285f13019adc49e11af501c658890c1e519f31f937c68b7e3"
dependencies = [
"darling 0.23.0",
"proc-macro2",
@@ -3764,9 +3777,9 @@ dependencies = [
[[package]]
name = "sse-stream"
version = "0.2.3"
version = "0.2.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f3962b63f038885f15bce2c6e02c0e7925c072f1ac86bb60fd44c5c6b762fb72"
checksum = "c25ac7aff0abd1dbc474536e40416e1102c7dd9bfba0b9861c6d357f835dcfb4"
dependencies = [
"bytes",
"futures-util",
@@ -4405,6 +4418,12 @@ dependencies = [
"serde",
]
[[package]]
name = "utf16_iter"
version = "1.0.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c8232dd3cdaed5356e0f716d285e4b40b932ac434100fe9b7e0e8e935b9e6246"
[[package]]
name = "utf8_iter"
version = "1.0.4"
@@ -5082,6 +5101,12 @@ dependencies = [
"wasmparser",
]
[[package]]
name = "write16"
version = "1.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d1890f4022759daae28ed4fe62859b1236caebfc61ede2f63ed4e695f3f6d936"
[[package]]
name = "writeable"
version = "0.6.3"
@@ -5214,13 +5239,33 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2d04a6b5381502aa6087c94c669499eb1602eb9c5e8198e534de571f7154809b"
dependencies = [
"crc32fast",
"flate2",
"indexmap",
"memchr",
"typed-path",
"zopfli",
]
[[package]]
name = "zlib-rs"
version = "0.6.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3be3d40e40a133f9c916ee3f9f4fa2d9d63435b5fbe1bfc6d9dae0aa0ada1513"
[[package]]
name = "zmij"
version = "1.0.21"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa"
[[package]]
name = "zopfli"
version = "0.8.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f05cd8797d63865425ff89b5c4a48804f35ba0ce8d125800027ad6017d2b5249"
dependencies = [
"bumpalo",
"crc32fast",
"log",
"simd-adler32",
]
Cargo.toml
+20 -12
View File
@@ -35,7 +35,7 @@ default-members = [
]
[workspace.package]
version = "2.3.1"
version = "2.5.1"
edition = "2024"
rust-version = "1.95"
license = "MIT"
@@ -44,16 +44,16 @@ authors = ["Fabio Akita <boss@akitaonrails.com>"]
[workspace.dependencies]
# Inter-crate dependencies.
ai-memory-core = { path = "crates/ai-memory-core", version = "2.3.1" }
ai-memory-store = { path = "crates/ai-memory-store", version = "2.3.1" }
ai-memory-wiki = { path = "crates/ai-memory-wiki", version = "2.3.1" }
ai-memory-mcp = { path = "crates/ai-memory-mcp", version = "2.3.1" }
ai-memory-hooks = { path = "crates/ai-memory-hooks", version = "2.3.1" }
ai-memory-llm = { path = "crates/ai-memory-llm", version = "2.3.1" }
ai-memory-consolidate = { path = "crates/ai-memory-consolidate", version = "2.3.1" }
ai-memory-web = { path = "crates/ai-memory-web", version = "2.3.1" }
ai-memory-workstream = { path = "crates/ai-memory-workstream", version = "2.3.1" }
ai-memory-test-support = { path = "crates/ai-memory-test-support", version = "2.1.1" }
ai-memory-core = { path = "crates/ai-memory-core", version = "2.4.0" }
ai-memory-store = { path = "crates/ai-memory-store", version = "2.4.0" }
ai-memory-wiki = { path = "crates/ai-memory-wiki", version = "2.4.0" }
ai-memory-mcp = { path = "crates/ai-memory-mcp", version = "2.4.0" }
ai-memory-hooks = { path = "crates/ai-memory-hooks", version = "2.4.0" }
ai-memory-llm = { path = "crates/ai-memory-llm", version = "2.4.0" }
ai-memory-consolidate = { path = "crates/ai-memory-consolidate", version = "2.4.0" }
ai-memory-web = { path = "crates/ai-memory-web", version = "2.4.0" }
ai-memory-workstream = { path = "crates/ai-memory-workstream", version = "2.4.0" }
ai-memory-test-support = { path = "crates/ai-memory-test-support", version = "2.3.2" }
# (Workspace shared deps follow below)
@@ -136,7 +136,7 @@ winapi-util = "0.1"
# `aws_lc_rs`, which is exactly the C/JNI toolchain the bridge avoids by using
# rmcp's `reqwest-tls-no-provider`. `ring` is already in the lockfile via reqwest 0.12.
rustls = { version = "0.23", default-features = false, features = ["ring", "logging", "std", "tls12"] }
rmcp = { version = "1.7", features = ["server", "macros", "transport-io", "transport-streamable-http-server", "schemars"] }
rmcp = { version = "2.2", features = ["server", "macros", "transport-io", "transport-streamable-http-server", "schemars"] }
schemars = "1"
axum = "0.8"
# Same `http` 1.x rmcp/axum use, so handlers can read the injected
@@ -149,9 +149,17 @@ tower-http = { version = "0.6", features = ["fs", "cors"] }
# Backup / restore archive format.
tar = "0.4"
flate2 = "1"
# Native Windows release zip extract/write for `ai-memory upgrade` (#801).
# Already in the lockfile via candle-core. Keep default-features off so we
# do not pull bzip2/`libbz2-rs-sys` (license not in deny allowlist); Deflate
# covers PowerShell Compress-Archive release zips and Stored test fixtures.
zip = { version = "8.6", default-features = false, features = ["deflate"] }
# Text munging / hook payload sanitisation.
regex = "1"
# Unicode NFC for the portable page-path key. Already in the tree via
# idna -> idna_adapter, so this adds no new code to the binary.
icu_normalizer = "2"
# Wiki git versioning.
git2 = { version = "0.21", default-features = false, features = ["vendored-libgit2"] }
DATA_HANDLING.md
+20 -6
View File
@@ -52,7 +52,8 @@ and each requires a deliberate config change to turn on.
opt-in and sanitized"). Persisting the coding assistant's final-turn text
requires a double opt-in: `capture_assistant` on the server *and*
`install-hooks --capture-assistant` on the client. Once enabled, captured
text flows into consolidation/reviewer prompts and — only if you have
text rides in the session's automatic handoff to the next session and
flows into consolidation/reviewer prompts and — only if you have
separately configured a cloud LLM provider — is sent to that provider. The
flag is global to the install; there is no per-project exclusion once it's
on.
@@ -101,11 +102,24 @@ them directly — ai-memory is a conduit, not a party to that relationship.
## Deletion / retention
There is no built-in retention-expiry policy; data persists until removed.
Deletion is filesystem-level: removing the data directory removes everything
in it. For scoped deletion, a per-project purge operation exists and is
isolation-safe — it cannot delete files or rows belonging to a different
`(workspace_id, project_id)` (`SECURITY.md`, "Per-project isolation").
Most data persists until removed, but two mechanisms expire some of it without
operator action, both on by default:
- **Decay eviction of cold episodic pages.** The daily forget sweep
(`[maintenance] forget_sweep_interval_secs`, default 86400) tombstones
episodic pages whose relevance score falls below `[decay] cold_threshold`
(default 0.20, with an ~80-day survival floor) and hard-deletes those
tombstones after `[decay] hard_delete_after_days` (default 180). Semantic,
procedural and pinned pages are exempt, and raw observations are kept
(`observation_retention_days = 0`).
- **TTL pages.** A page written with `expires_at` is hidden after its expiry
and hard-deleted by the next forget sweep; a TTL outranks `pinned`.
Everything else persists until removed. Deletion is filesystem-level: removing
the data directory removes everything in it. For scoped deletion, a per-project
purge operation exists and is isolation-safe — it cannot delete files or rows
belonging to a different `(workspace_id, project_id)` (`SECURITY.md`,
"Per-project isolation").
## Network exposure (if you run the server non-loopback)
README.md
+155 -55
View File
@@ -53,17 +53,31 @@ ai-memory is what's on the other side of those walls.
uses **zero LLM calls**: capture, search, and handoffs all work with no
API key at all.
- **It ages gracefully, without an LLM.** Memory decays on a schedule you
can tune per tier, and the memory you actually use decays *slower* — open a
page, search it, or reach it through a link and it earns its keep. When
episodic notes go cold they can be compacted down to their durable facts
(file paths, error codes, decisions) instead of dropped, near-duplicates
collapse into one, and likely contradictions get flagged — all with **zero
API calls**. Nothing is hard-deleted: the original stays in git and the
version chain (`restore-page` brings it back). Access-weighted retention is
always on because it can only ever keep memory *longer*; the parts that
rewrite or drop content (compaction, dedup, per-tier curves) stay off by
default until you turn them on.
- **And it can dream, if you let it.** Point it at an LLM and an opt-in
background pass will, while you're idle, rewrite whole clusters of cold
notes into single coherent pages — cancelling the moment you come back to
work. It never deletes a source (the pre-merge versions stay reachable),
it's off by default, and it's gated on a recall eval before it could ever
become default behavior. The zero-LLM path above is what runs unless you
ask for more.
- **It tells you the truth about itself.** One self-contained binary.
Purge commands that say exactly what "deleted" means. A measured write
ceiling (~700/s) instead of a guessed one. An audit log of every
mutation. Boring, in the way infrastructure should be.
**Coming from Mem0, Zep, mcp-memory-service, Hindsight, OpenViking, or Claude
Code's built-in memory?** [How ai-memory compares](docs/comparison.md) is a
fair, specific rundown — where each approach wins, where ai-memory differs, the
published benchmark, and how the field has independently validated the
file-first, pages-over-facts bet.
## How it works
```
@@ -120,10 +134,53 @@ caveats is in [`docs/support-matrix.md`](docs/support-matrix.md).
| VS Code Copilot | MCP-only |
| Zed | MCP-only |
| Muse Code | MCP-only |
| Hermes Agent | Community |
| Hermes Agent | Supported |
| LLM/auth providers | Supported |
| Embedding providers | Supported |
## Coming from another tool?
Most agent-memory tools optimize one thing — extracting atomic facts per turn,
a temporal knowledge graph, an agent-editable memory OS, or a hosted context
API. ai-memory optimizes something different: a **git-backed markdown wiki as
the source of truth**, with a derived index for retrieval, captured
automatically from lifecycle hooks, shared across agents, machines, and people,
and working with **zero LLM calls by default**. Here's what carries over from
each, and what you gain:
| Coming from… | What's similar | What you gain |
|---|---|---|
| **Mem0 / fact extractors** (LangMem) | Automatic per-turn capture | Memory compiles into readable **pages** you own and edit, not opaque fact rows; retrieval fuses FTS + entity + graph (+ optional vectors), not vector-only |
| **Zep / Graphiti** (temporal KG) | Temporal reasoning, typed relations | Bi-temporal-lite (`as_of`, version-filtered search) and typed edges without standing up a graph database — on one binary |
| **mcp-memory-service** (closest sibling) | SQLite + local embeddings, hook capture, typed edges, honest numbers — and, on 2.4, per-tier decay curves, extractive compression, DBSCAN cold-cluster dedup, access reinforcement, and contradiction flagging | Human-editable markdown **pages** instead of fact-rows, cross-agent claim-once handoffs, and the same aging machinery done **zero-LLM by default, reversibly** (supersede-not-delete + `restore-page`), and **off by default** |
| **basic-memory** (file-first sibling) | Markdown-on-disk as the source of truth | Automatic lifecycle capture and a derived FTS/entity/graph index on top, cross-agent handoffs, and multi-user sharing built in |
| **Claude Code built-in memory** | "Remember my project" convenience, zero setup | Synced across machines and agents, searchable, team-capable, and captures tool lifecycle — not a per-laptop `MEMORY.md` |
| **Hindsight / OpenViking** (hosted, LLM-required) | Living pages / document memory with a background consolidation loop — and, on 2.4, belief-strength confidence plus an opt-in LLM "dream" rewrite of cold clusters | A self-contained binary that runs zero-LLM by default and keeps memory in files you own; per-project team sharing instead of strict per-bank isolation; the dream/belief features are **opt-in, off by default, and never delete a source** (vs a mandatory LLM loop) |
| **Supermemory / LiquidLM** (hosted memory API) | A managed second brain with automatic ingestion | Git-versioned markdown you own, no required API spend, offline operation, and per-project team sharing — ai-memory remembers *this repo*, not a general vault |
The consistent theme: **files you own** (git-backed markdown), a **zero-LLM
default**, **one self-contained binary**, **cross-agent + cross-machine + team**
sharing, **automatic lifecycle capture**, and **typed, claim-once handoffs**.
Opt-in features (LLM consolidation, vector search, the "dream" consolidation
pass, belief-strength in ranking) stay opt-in — and the zero-LLM aging path
(per-tier decay, extractive compaction, dedup, contradiction flagging,
access-weighted retention) works with no API key at all.
**Built on the shoulders of:** the
[Karpathy LLM Wiki](https://gist.github.com/karpathy/442a6bf555914893e9891c11519de94f)
(compile-not-retrieve),
[agentmemory](https://github.com/rohitg00/agentmemory) (this project is its Rust
successor), [basic-memory](https://github.com/basicmachines-co/basic-memory)
(markdown-on-disk truth),
[cognee](https://github.com/topoteretes/cognee) (pipeline composition and
triplet embeddings),
[Hermes Agent](https://github.com/NousResearch/hermes-agent) (the
self-improvement loop), and [A-MEM](https://arxiv.org/abs/2502.12110)
(Zettelkasten-style atomic notes).
The full, fair rundown — where each approach wins, where ai-memory differs, the
published benchmark — is in [How ai-memory compares](docs/comparison.md).
## Quick start
### Arch Linux (AUR)
@@ -137,6 +194,18 @@ yay -S ai-memory-bin # prebuilt Linux x86_64/aarch64 binary
yay -S ai-memory # builds from source
```
### Fedora (RPM)
Download the `x86_64` or `aarch64` RPM from the
[latest release](https://github.com/akitaonrails/ai-memory/releases/latest),
then install it:
```bash
sudo dnf install ./ai-memory-*.rpm
```
Then follow the native Linux service instructions in [`docs/install.md`](docs/install.md).
Single-user workstation:
```bash
@@ -152,6 +221,37 @@ System service installs use `/var/lib/ai-memory` and `/etc/ai-memory/` via the
packaged unit. Full user-service, system-service, auth, and provider setup is in
[`docs/install.md#arch-linux-native-packages-aur`](docs/install.md#arch-linux-native-packages-aur).
### macOS (menu bar app)
A self-contained `.app` that bundles the native `ai-memory` binary and
`hooks/` tree, starts the existing LaunchAgent, and opens `/web`,
`ai-memory status`, and `config.toml` from the menu bar. Wiki, SQLite,
config, and models stay in `~/Library/Application Support/ai-memory`, so
replacing the app is an update — it does not rewrite that tree.
Needs a Rust toolchain and Xcode / Swift 6 (same as a source build):
```bash
git clone https://github.com/akitaonrails/ai-memory
cd ai-memory
./companions/ai-memory-macos/build.sh
open "companions/ai-memory-macos/dist/AI Memory.app"
```
Drag **AI Memory.app** to `/Applications`, then **Install & Start Server**
from the menu extra (no Dock icon). When the status item is green, wire an
agent with the bundled binary:
```bash
BIN="/Applications/AI Memory.app/Contents/Resources/runtime/ai-memory"
"$BIN" install-mcp --client claude-code --apply
"$BIN" install-hooks --agent claude-code --apply
```
Prebuilt tarball and launchd-without-the-app paths:
[`docs/macos.md`](docs/macos.md). Companion details:
[`companions/ai-memory-macos`](companions/ai-memory-macos).
### Docker
You need: Docker or Podman + an agent CLI from the [Support Matrix](#support-matrix),
@@ -225,8 +325,14 @@ wrapper automatically uses Podman when Docker is not installed. Set
On Linux/macOS, that's it. Start a Claude Code session as usual - every
prompt and tool call now lands in ai-memory, and the next session you
open in this project will see a handoff with where you left off.
On macOS, the native release binary is also supported and recommended when you
do not need Docker; see [`docs/macos.md`](docs/macos.md).
On macOS the native binary is the recommended path when you do not need
Docker — either the [menu bar app](#macos-menu-bar-app) above or a
[release tarball / launchd agent](docs/macos.md). Later updates for that
path use `ai-memory upgrade` (checksum-verified GitHub release replace + hook
refresh) — see
[`docs/install.md#keeping-ai-memory-up-to-date`](docs/install.md#keeping-ai-memory-up-to-date).
The same native upgrade path covers Windows x86_64 zip installs under a
writable prefix (see [`docs/windows.md`](docs/windows.md) Scenario C).
Wiring another agent is the same two commands with a different name —
`--client codex`, `--agent codex`, and so on for every row of the support
@@ -243,7 +349,8 @@ way to launch: the first time it runs a harness it auto-installs that harness's
ai-memory hooks + MCP if they are missing (so capture and recall just work —
no separate `install-hooks`/`install-mcp` step to forget), it wires the right
project scope by construction, and it adds cross-harness *session* continuity on
top of shared memory. Everything is idempotent and one-time per harness.
top of shared memory. Everything is idempotent and one-time per harness and
config home.
```bash
ai-memory run claude
@@ -257,7 +364,9 @@ Auto-wiring is on by default; opt out with `ai-memory run --no-autowire` or
`ai-memory run`).
`ai-memory uninstall --apply` removes everything ai-memory installed,
and only what it installed. Install commands are idempotent and write
and only what it installed. It also clears `ai-memory run`'s auto-wire
record, so the next managed launch wires that harness again; to keep it
unwired, launch with `--no-autowire` or set `AI_MEMORY_RUN_AUTOWIRE=false`. Install commands are idempotent and write
timestamped backups next to any file they touch.
## Everyday use
@@ -297,7 +406,8 @@ your machine can reach it. From there, hardening is incremental: a bearer
token for the LAN, per-user accounts, OIDC device auth for hooks, TLS via
a reverse proxy. Capture is sanitized at a typed privacy boundary before
anything is stored, and per-repository `[capture]` rules can exclude
paths or invert to allowlist mode.
paths or invert to allowlist mode. A repository can also route its capture
to a different server than the one the hooks were installed against.
The full model is in [`docs/security.md`](docs/security.md),
[`docs/users.md`](docs/users.md), and
@@ -341,56 +451,46 @@ diagram, crate breakdown, schema notes, and invariants.
## Docs
### For users
| File | What it is |
|---|---|
| [`docs/cookbook.md`](docs/cookbook.md) | **Task-oriented cheat sheet.** "I want to do X" → how: recall prior work, keep a rule a project must follow, import an existing knowledge base (OKF norms/specs) and have a project read a specific document, and get two agents/repos working together. Start here if you're unsure what ai-memory can do for you. |
| [`docs/install.md`](docs/install.md) | **Installation cookbook.** Every agent CLI, every alternative (curl, source build, no-docker, no-auth), and the server-on-a-different-machine (homelab/LAN) walkthrough. Read after the Quick start if your setup doesn't match the happy path. |
| [`docs/usage.md`](docs/usage.md) | Handoffs, proactive memory queries, slim routing snippet + managed Agent Skills, migration from other memory tools, web UI, raw-wiki inspection, and rules-vs-facts workflow. |
| [`docs/managed-workstreams.md`](docs/managed-workstreams.md) | Optional `ai-memory run` continuity across Claude Code, Codex, OpenCode, OpenCode 2 beta, Pi, Crush, Kimi Code, Command Code, Kiro CLI v2/v3, OMP, Grok Build CLI, and Antigravity CLI: automatic harness selection, native resume, argument forwarding, ledger search, privacy, and recovery. The preferred way to launch — it auto-installs a harness's hooks + MCP on first run. |
| [`docs/agent-messaging.md`](docs/agent-messaging.md) | Cross-project agent-to-agent messaging: a directed, claim-once inbox/queue so an agent in one project can hand a self-contained request to an agent in another, plus the on-start "you have mail" notice. Four `memory_message_*` MCP tools + `ai-memory message` CLI. |
| [`docs/managed-harness-contributions.md`](docs/managed-harness-contributions.md) | Protocol and acceptance bar for contributors adding managed resume, read-only transcript import, and startup context delivery to another harness. |
| [`docs/marker-file.md`](docs/marker-file.md) | `.ai-memory.toml` workspace/project routing for multi-client trees, mono-repos, worktrees, and work/personal separation. |
| [`docs/cookbook.md`](docs/cookbook.md) | **Task-oriented cheat sheet.** "I want to do X" → how: recall prior work, keep a project rule, import an existing knowledge base, get two agents/repos working together. Start here. |
| [`docs/install.md`](docs/install.md) | **Installation cookbook.** Every agent CLI, every alternative (curl, source build, no-docker, no-auth), and the server-on-a-different-machine walkthrough. |
| [`docs/usage.md`](docs/usage.md) | Handoffs, proactive memory queries, slim routing snippet + managed Agent Skills, web UI, raw-wiki inspection, and rules-vs-facts workflow. |
| [`docs/managed-workstreams.md`](docs/managed-workstreams.md) | Optional `ai-memory run` continuity across harnesses: auto harness selection, native resume, argument forwarding, ledger search, privacy, and recovery. |
| [`docs/agent-messaging.md`](docs/agent-messaging.md) | Cross-project agent-to-agent messaging: a directed, claim-once inbox/queue plus the on-start "you have mail" notice. |
| [`docs/marker-file.md`](docs/marker-file.md) | `.ai-memory.toml` workspace/project routing for multi-client trees, mono-repos, worktrees, and work/personal separation, plus per-repository server profiles. |
| [`docs/auto-scope.md`](docs/auto-scope.md) | `[auto_scope]` modes for shared servers: default single-slot routing, session-aware isolation, and multi-user `per_actor` behavior. |
| [`docs/macos.md`](docs/macos.md) | macOS install paths: native release binary (recommended), source build, the Docker wrapper, hook-platform notes, and current macOS limitations. |
| [`docs/windows.md`](docs/windows.md) | Windows install modes: full WSL2, native Windows with Docker Desktop, prebuilt native release zip, native source builds, and current hook/MCP harness caveats. |
| [`docs/macos.md`](docs/macos.md) | macOS install paths: menu bar app, native release tarball, source build, Docker wrapper, launchd, and current limitations. |
| [`docs/windows.md`](docs/windows.md) | Windows install modes: full WSL2, native Windows with Docker Desktop, prebuilt native release zip, native source builds, and caveats. |
| [`docs/mcp-install.md`](docs/mcp-install.md) | Per-client MCP and lifecycle notes, handoff-injection limits, and community bridge guidance. |
| [`docs/deploy.md`](docs/deploy.md) | Homelab deploy: bin/deploy, bearer-token auth, pointers to the TLS guide. |
| [`docs/users.md`](docs/users.md) | **Multi-user attribution and human login.** Four-rung bearer ladder, password sessions, `ai-memory user` / `api-key` walkthrough, brownfield `aim_` migration. |
| [`docs/https-via-proxy.md`](docs/https-via-proxy.md) | **HTTPS via a reverse proxy.** When you need TLS (multi-user, non-loopback) and when you don't (loopback / stdio). Copy-paste docker compose templates for Caddy + Let's Encrypt, Caddy + internal CA (LAN-only), Cloudflare Tunnel (no open ports), and external cert files; plus native-Caddy + nginx recipes. The "thinking you're secure when you're not" failure modes explicitly called out. |
| [`docs/lifecycle-ops.md`](docs/lifecycle-ops.md) | **Read before running purge / rename / backup / restore / reset / reindex / restore-page.** Safety matrix for state-touching commands, per-project disk layout (how isolation actually works), checkpoint-based page recovery, and operator workflows for "fresh start", "snapshot before risky op", "drop one project", and rebuilding SQLite from wiki files. |
| [`docs/auto-improvement-loop.md`](docs/auto-improvement-loop.md) | Auto-improvement design notes: Hermes-inspired scheduled review, auto-approval default, manual review opt-in, pending proposal storage, and curator work. |
| [`docs/companion-crates.md`](docs/companion-crates.md) | Boundary and implementation plan for optional companion projects, including the standalone importer at [`companions/ai-memory-importer`](companions/ai-memory-importer), without widening core ai-memory. |
| [`docs/llm-provider-comparison.md`](docs/llm-provider-comparison.md) | Empirical notes behind the recommended LLM defaults. |
| [`DATA_HANDLING.md`](DATA_HANDLING.md) | **Data-flow reference for security/legal review.** What's stored, what's local-only, the two opt-in paths that send data externally, and how deletion/retention work. |
| [`docs/sso.md`](docs/sso.md) | Enterprise identity: the existing OIDC device-auth flow, what it does and doesn't cover, and how to front the server with an OIDC-aware gateway. |
| [`docs/airgapped-install.md`](docs/airgapped-install.md) | Offline/air-gapped install: self-contained build, checksum-verified release binaries, and the offline path for local embedding models. |
| [`docs/llm-provider-fallback.md`](docs/llm-provider-fallback.md) | Proposed opt-in fallback-chain design for transient LLM-provider failures; not yet a supported configuration surface. |
| [`docs/users.md`](docs/users.md) | **Multi-user attribution and human login.** Four-rung bearer ladder, password sessions, `ai-memory user` / `api-key` walkthrough, brownfield migration. |
| [`docs/https-via-proxy.md`](docs/https-via-proxy.md) | **HTTPS via a reverse proxy.** When you need TLS and when you don't, with copy-paste Caddy / nginx / Cloudflare Tunnel templates and the "secure when you're not" failure modes. |
| [`docs/lifecycle-ops.md`](docs/lifecycle-ops.md) | **Read before purge / rename / backup / restore / reset / reindex / restore-page.** Safety matrix, per-project disk layout, checkpoint page recovery, and operator workflows. |
| [`docs/backup.md`](docs/backup.md) | Backing up the wiki + data dir to a remote git repository: what to include, what to exclude, scheduled push pattern, restore, and security posture. Companion to `docs/lifecycle-ops.md` (which covers the on-box `ai-memory backup` snapshot). |
| [`docs/llm-providers.md`](docs/llm-providers.md) | Provider configuration for consolidation and embeddings. |
| [`docs/security.md`](docs/security.md) | The full security model. |
| [`docs/support-matrix.md`](docs/support-matrix.md) | The full agent/platform matrix with notes. |
| [`docs/use-cases.md`](docs/use-cases.md) | Scenario walkthroughs. |
| [`DATA_HANDLING.md`](DATA_HANDLING.md) | **Data-flow reference for security/legal review.** What's stored, what's local-only, the two opt-in external paths, and how deletion/retention work. |
| [`docs/sso.md`](docs/sso.md) | Enterprise identity: the OIDC device-auth flow, its scope, and how to front the server with an OIDC-aware gateway. |
| [`docs/airgapped-install.md`](docs/airgapped-install.md) | Offline/air-gapped install: self-contained build, checksum-verified release binaries, and offline local embedding models. |
| [`docs/MIGRATION-2.0.md`](docs/MIGRATION-2.0.md) | Upgrading an existing store to 2.0: the backup-gated automatic migration and how to restore. |
| [`docs/benchmarks/`](docs/benchmarks/README.md) | Published retrieval-quality numbers with provenance, reproducible from the in-repo harness. |
| [`docs/okf.md`](docs/okf.md) | The wiki is natively an Open Knowledge Format (OKF v0.2) bundle; design and field mapping. |
### For contributors
| File | What it is |
|---|---|
| [`docs/ARCHITECTURE.md`](docs/ARCHITECTURE.md) | Operational summary: data flow, crate layout, cross-cutting invariants, schema. |
| [`docs/design-decisions.md`](docs/design-decisions.md) | The full v1 spec. |
| Research docs under `docs/` | Karpathy LLM Wiki notes, Hermes Agent, agentmemory / basic-memory / cognee / hindsight deep-dives, the 2026 landscape survey (Zep/Graphiti, Letta, Mem0, mcp-memory-service, OpenViking, …), and lessons-learned from upstream issues. |
- [`docs/support-matrix.md`](docs/support-matrix.md) - the full agent/platform matrix with notes.
- [`docs/use-cases.md`](docs/use-cases.md) - scenario walkthroughs.
- [`docs/llm-providers.md`](docs/llm-providers.md) - provider configuration.
- [`docs/security.md`](docs/security.md) - the full security model.
- [`docs/comparison.md`](docs/comparison.md) - how ai-memory compares to other memory tools, fairly, and how the field validates the approach.
- [`docs/research-2026-landscape.md`](docs/research-2026-landscape.md) - how the field looks and where we sit in it.
- [`docs/ROADMAP-2.0.md`](docs/ROADMAP-2.0.md) - the plan for the 2.0 release, one item at a time.
- [`docs/okf.md`](docs/okf.md) - the wiki is natively an Open Knowledge Format (OKF v0.2) bundle; design and field mapping.
- [`docs/typed-edges.md`](docs/typed-edges.md) - typed relation edges (`causes` / `fixes` / `contradicts`) and how lint uses them.
- [`docs/temporal.md`](docs/temporal.md) - ingestion-time validity on the entity index and page versions, and `as_of` time-travel queries (entity timeline + version-filtered FTS).
- [`docs/local-embeddings.md`](docs/local-embeddings.md) - in-process embeddings with no API key (`embedding_provider = "local"`).
- [`docs/experience.md`](docs/experience.md) - the opt-in cross-session abstraction pass: knowledge visible only across trajectories.
- [`docs/MIGRATION-2.0.md`](docs/MIGRATION-2.0.md) - upgrading an existing store to 2.0: the backup-gated automatic migration and how to restore.
- [`docs/benchmarks/`](docs/benchmarks/README.md) - published retrieval-quality numbers with provenance, reproducible from the in-repo harness.
## Influences and prior art
- **[Karpathy LLM Wiki](https://gist.github.com/karpathy/442a6bf555914893e9891c11519de94f)** - the compile-not-retrieve pattern.
- **[agentmemory](https://github.com/rohitg00/agentmemory)** - most of the right ideas; this project is the Rust successor.
- **[basic-memory](https://github.com/basicmachines-co/basic-memory)** - the markdown-on-disk source-of-truth model.
- **[cognee](https://github.com/topoteretes/cognee)** - pipeline composition and triplet embeddings.
- **[Hermes Agent](https://github.com/NousResearch/hermes-agent)** - the self-improvement loop: post-turn review, approval gates, and curator boundaries.
- **[A-MEM](https://arxiv.org/abs/2502.12110)** - Zettelkasten-style atomic notes with link evolution.
| [`docs/managed-harness-contributions.md`](docs/managed-harness-contributions.md) | Protocol and acceptance bar for adding managed resume, transcript import, and startup context delivery to another harness. |
| [`docs/companion-crates.md`](docs/companion-crates.md) | Optional companion projects: the [importer](companions/ai-memory-importer) and [external lifecycle relay](companions/ai-memory-relay). |
| [`docs/external-lifecycle.md`](docs/external-lifecycle.md) | External lifecycle producers: per-execution native capture suppression, preserved handoffs, batch ingestion and stable retry identity. |
| [`docs/auto-improvement-loop.md`](docs/auto-improvement-loop.md) | Auto-improvement design notes: scheduled review, auto-approval default, manual review opt-in, pending proposal storage, and curator work. |
## License
SECURITY.md
+3
View File
@@ -101,6 +101,9 @@ what the project is and is not designed to defend against.
`[REDACTED]`.
- Captured assistant text flows into the consolidation and reviewer prompts,
and — if you configure a cloud LLM provider — is sent to that provider.
The latest excerpt of a session also rides in its automatic handoff, so
the next session that claims the baton receives it as startup context. It
is not rendered into the git-tracked session page.
- The opt-in is **global** to the install: there is no per-project marker to
exclude a sensitive repository once the flag is on (assistant text is not
path-attributable). Turn the server flag off to disable it everywhere.
bin/ai-memory
+28 -4
View File
@@ -9,7 +9,7 @@
#
# Special wrapper-only subcommands (not forwarded to the binary):
# ai-memory upgrade Pull the latest image + remind to re-stage hooks.
# ai-memory run ... Use a cached native client so it can exec host agents.
# ai-memory run ... Use a local native client so it can exec host agents.
# ai-memory show ... Use that client for host project/harness discovery.
# ai-memory continue ... Use that client to resume the newest linked checkout.
# ai-memory workstreams Use that client to inspect the host checkout identity.
@@ -55,6 +55,10 @@ DATA_VOLUME="${AI_MEMORY_DATA_VOLUME:-ai-memory-data}"
CACHE_DIR="${XDG_CACHE_HOME:-${HOME}/.cache}/ai-memory"
VERSION_CHECK_FILE="${CACHE_DIR}/last-version-check"
HOOKS_STAGE_DIR="${HOME}/.local/share/ai-memory/hooks"
# `run` auto-wires hooks whose command is this client's own path, so a cache
# flush would break every hook until the next managed launch. Keep it with the
# host's ai-memory data instead of under CACHE_DIR.
NATIVE_RUNNER_DIR="${XDG_DATA_HOME:-${HOME}/.local/share}/ai-memory/native-runner"
WRAPPER_URL="${AI_MEMORY_WRAPPER_URL:-https://github.com/akitaonrails/ai-memory/releases/latest/download/ai-memory-wrapper}"
WRAPPER_SHA256_URL="${AI_MEMORY_WRAPPER_SHA256_URL:-${WRAPPER_URL}.sha256}"
@@ -354,7 +358,7 @@ cmd_upgrade() {
fi
echo "→ pulling ${IMAGE}"
"${DOCKER}" pull "${IMAGE}"
rm -f "${CACHE_DIR}/native-runner/last-check"
rm -f "${NATIVE_RUNNER_DIR}/last-check"
local found_agents=()
if [ -d "${HOOKS_STAGE_DIR}" ]; then
@@ -455,7 +459,7 @@ cmd_upgrade() {
# Managed launch/discovery commands must execute on the host: checkouts, harnesses,
# and native transcript stores are host resources, not contents of the helper
# container. Keep a checksum-verified release client beside the wrapper cache.
# container. Keep a checksum-verified release client under NATIVE_RUNNER_DIR.
native_host_binary() {
if [ -n "${AI_MEMORY_NATIVE_BIN:-}" ]; then
[ -x "${AI_MEMORY_NATIVE_BIN}" ] || {
@@ -486,7 +490,7 @@ native_host_binary() {
esac
artifact="ai-memory-${os}-${arch}"
base="https://github.com/akitaonrails/ai-memory/releases/latest/download/${artifact}.tar.gz"
native_dir="${CACHE_DIR}/native-runner"
native_dir="${NATIVE_RUNNER_DIR}"
binary="${native_dir}/ai-memory"
archive="${native_dir}/${artifact}.tar.gz"
check_file="${native_dir}/last-check"
@@ -533,6 +537,13 @@ native_host_binary() {
[ -x "${tmp}/ai-memory" ] || chmod +x "${tmp}/ai-memory"
mv "${tmp}/${artifact}.tar.gz" "${archive}"
mv "${tmp}/ai-memory" "${binary}"
# `run` auto-wires hooks through this client, and install-hooks looks for
# its script bundle beside the binary. Without it, auto-wiring fails for
# every script-based harness on a host where nothing was staged yet.
if [ -d "${tmp}/hooks" ]; then
rm -rf "${native_dir}/hooks"
mv "${tmp}/hooks" "${native_dir}/hooks"
fi
rm -rf "${tmp}"
touch "${check_file}"
fi
@@ -604,6 +615,7 @@ for var in \
VOYAGE_API_KEY \
LLM_API_KEY \
EMBEDDING_API_KEY \
OPENCODE_API_KEY \
RUST_LOG
do
if [ -n "${!var:-}" ]; then
@@ -611,6 +623,18 @@ do
fi
done
# These two are presence-based, not non-empty-based like the loop above: an
# operator sets one to the empty string to deliberately clear a
# config.toml-configured prefix without editing the file (see
# Config::load's figment overlay), and a present-but-empty value must reach
# the container for that to work — `[ -n ]` above would drop it, making the
# wrapper indistinguishable from the var never having been set at all.
for var in AI_MEMORY_EMBEDDING_QUERY_PREFIX AI_MEMORY_EMBEDDING_DOCUMENT_PREFIX; do
if [ -n "${!var+x}" ]; then
ENV_ARGS+=(-e "${var}")
fi
done
# The wrapper itself runs the CLI inside a short-lived helper container, while
# the README server runs in the long-lived ai-memory container and publishes
# 127.0.0.1:49374 on the host. Inside a normal bridge-network helper,
bin/ai-memory.ps1
+33
View File
@@ -145,19 +145,28 @@ foreach ($Name in @(
"AI_MEMORY_LLM_PROVIDER",
"AI_MEMORY_LLM_MODEL",
"AI_MEMORY_LLM_BASE_URL",
"AI_MEMORY_COPILOT_CLIENT_ID",
"AI_MEMORY_EMBEDDING_PROVIDER",
"AI_MEMORY_EMBEDDING_MODEL",
"AI_MEMORY_EMBEDDING_BASE_URL",
"AI_MEMORY_EMBEDDING_DIM",
"AI_MEMORY_ALLOWED_HOSTS",
"AI_MEMORY_WORKSTREAM_ID",
"CLAUDE_CONFIG_DIR",
"CLAUDE_CODE_SESSION_ID",
"ANTHROPIC_API_KEY",
"ANTHROPIC_OAUTH_TOKEN",
"CLAUDE_CODE_OAUTH_TOKEN",
"OPENAI_API_KEY",
"GEMINI_API_KEY",
"GOOGLE_API_KEY",
"COPILOT_GITHUB_TOKEN",
"GITHUB_COPILOT_API_TOKEN",
"COPILOT_API_URL",
"VOYAGE_API_KEY",
"LLM_API_KEY",
"EMBEDDING_API_KEY",
"OPENCODE_API_KEY",
"RUST_LOG"
)) {
if (-not [string]::IsNullOrEmpty([Environment]::GetEnvironmentVariable($Name))) {
@@ -165,6 +174,30 @@ foreach ($Name in @(
}
}
# Presence-based, not non-empty-based like the loop above: an operator sets
# one of these to the empty string to deliberately clear a
# config.toml-configured prefix without editing the file (see
# Config::load's figment overlay), and a present-but-empty value must reach
# the container for that to work — `IsNullOrEmpty` above would drop it,
# making the wrapper indistinguishable from the variable never having been
# set at all. `GetEnvironmentVariable` returns `$null` only when the
# variable is truly unset, and `""` when it is set-but-empty, so a `-ne
# $null` check is exactly the presence test needed here.
#
# An operator's own `$env:NAME = ''` additionally needs PowerShell 7.5+
# (first built on .NET 9) to leave a set-but-empty variable rather than
# deleting it; not exercised on a real pwsh runtime.
# https://learn.microsoft.com/en-us/dotnet/api/system.environment.setenvironmentvariable
# https://learn.microsoft.com/en-us/powershell/scripting/whats-new/what-s-new-in-powershell-75
foreach ($Name in @(
"AI_MEMORY_EMBEDDING_QUERY_PREFIX",
"AI_MEMORY_EMBEDDING_DOCUMENT_PREFIX"
)) {
if ($null -ne [Environment]::GetEnvironmentVariable($Name)) {
$DockerArgs += @("-e", $Name)
}
}
# Docker Desktop gives Windows no host networking for Linux containers, so a
# thin-client command (status, search, bootstrap, ...) reaches the loopback-
# published server from this helper container through Docker Desktop's host
companions/ai-memory-importer/Cargo.lock
Generated
+4 -4
View File
@@ -929,9 +929,9 @@ dependencies = [
[[package]]
name = "rustls"
version = "0.23.41"
version = "0.23.45"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6b92b125634d9b795e7beca796cc790df15a7fb38323bf3196fda83292d06b1f"
checksum = "0d41d731c7d2f962d1ccc364cec258de3c0e93b38c2fb3ba97ac74513048d634"
dependencies = [
"once_cell",
"ring",
@@ -965,9 +965,9 @@ dependencies = [
[[package]]
name = "rustls-webpki"
version = "0.103.13"
version = "0.103.15"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "61c429a8649f110dddef65e2a5ad240f747e85f7758a6bccc7e5777bd33f756e"
checksum = "f3c3cf1d8b1e7d4927e2d154c3fcb02979afb9939629c62cd9048d4f07b60ac2"
dependencies = [
"ring",
"rustls-pki-types",
companions/ai-memory-importer/src/main.rs
+40 -4
View File
@@ -1089,11 +1089,23 @@ struct ParsedMarkdown {
pinned: bool,
}
/// Split an OMC page into frontmatter fields and body. The fence is found
/// the way `ai_memory_wiki::markdown::parse` finds it: past a leading UTF-8
/// BOM, and with either LF or CRLF fence lines, since a wiki checked out on
/// Windows with `core.autocrlf=true` has CRLF throughout. Missing it dropped
/// kind, tier, tags and pin, and imported the YAML block as body text.
fn parse_markdown(input: &str) -> Result<ParsedMarkdown> {
let mut out = ParsedMarkdown::default();
let body = if let Some(rest) = input.strip_prefix("---\n") {
if let Some(end) = rest.find("\n---\n") {
let yaml = &rest[..end];
let input = input.strip_prefix('\u{FEFF}').unwrap_or(input);
let (rest, newline) = if let Some(rest) = input.strip_prefix("---\r\n") {
(Some(rest), "\r\n")
} else {
(input.strip_prefix("---\n"), "\n")
};
let close = format!("\n---{newline}");
let body = if let Some(rest) = rest {
if let Some(end) = rest.find(&close) {
let yaml = rest[..end].trim_end_matches('\r');
let value: serde_yaml::Value =
serde_yaml::from_str(yaml).context("parse YAML frontmatter")?;
if let Some(map) = value.as_mapping() {
@@ -1103,7 +1115,7 @@ fn parse_markdown(input: &str) -> Result<ParsedMarkdown> {
out.pinned = yaml_bool(map, "pinned").unwrap_or(false);
out.tags = yaml_tags(map);
}
rest[end + "\n---\n".len()..].to_owned()
rest[end + close.len()..].to_owned()
} else {
input.to_owned()
}
@@ -1454,6 +1466,30 @@ mod tests {
assert_eq!(parsed.body, "# Body\ntext");
}
/// A wiki checked out on Windows with `core.autocrlf=true`, or saved
/// there by an editor, has CRLF line endings; one saved as "UTF-8 with
/// BOM" opens with U+FEFF. Either used to miss the fence, so kind, tier,
/// tags and pin were dropped and the YAML block was imported as body.
#[test]
fn parses_omc_frontmatter_with_crlf_or_a_bom() {
let crlf = "---\r\ntitle: T\r\nkind: rule\r\ntier: procedural\r\ntags: [a, b]\r\npinned: true\r\n---\r\n# Body\r\ntext";
let parsed = parse_markdown(crlf).unwrap();
assert_eq!(parsed.title.as_deref(), Some("T"));
assert_eq!(parsed.kind.as_deref(), Some("rule"));
assert_eq!(parsed.tier.as_deref(), Some("procedural"));
assert_eq!(parsed.tags, vec!["a", "b"]);
assert!(parsed.pinned);
assert_eq!(
parsed.body, "# Body\r\ntext",
"the body keeps its line endings"
);
let bom = parse_markdown("\u{feff}---\ntier: procedural\n---\n# Body\ntext").unwrap();
assert_eq!(bom.tier.as_deref(), Some("procedural"));
assert_eq!(bom.title.as_deref(), Some("Body"));
assert_eq!(bom.body, "# Body\ntext");
}
#[test]
fn planning_rejects_unknown_tier_before_live_write() {
let td = tempdir().unwrap();
companions/ai-memory-macos/.gitignore
+5
View File
@@ -0,0 +1,5 @@
.build/
.swiftpm/
dist/
*.xcodeproj/xcuserdata/
*.xcworkspace/xcuserdata/
companions/ai-memory-macos/Info.plist
+37
View File
@@ -0,0 +1,37 @@
<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
<plist version="1.0">
<dict>
<key>CFBundleDevelopmentRegion</key>
<string>en</string>
<key>CFBundleExecutable</key>
<string>AIMemoryMenu</string>
<key>CFBundleIdentifier</key>
<string>com.github.akitaonrails.ai-memory-menu</string>
<key>CFBundleInfoDictionaryVersion</key>
<string>6.0</string>
<key>CFBundleName</key>
<string>AI Memory</string>
<key>CFBundleDisplayName</key>
<string>AI Memory</string>
<key>CFBundlePackageType</key>
<string>APPL</string>
<key>CFBundleShortVersionString</key>
<string>1.0.0</string>
<key>CFBundleVersion</key>
<string>1</string>
<key>LSMinimumSystemVersion</key>
<string>14.0</string>
<key>LSUIElement</key>
<true/>
<key>NSHighResolutionCapable</key>
<true/>
<key>NSPrincipalClass</key>
<string>NSApplication</string>
<key>NSAppTransportSecurity</key>
<dict>
<key>NSAllowsLocalNetworking</key>
<true/>
</dict>
</dict>
</plist>
companions/ai-memory-macos/Package.swift
+28
View File
@@ -0,0 +1,28 @@
// swift-tools-version: 6.0
import PackageDescription
let package = Package(
name: "AIMemoryMenu",
platforms: [.macOS(.v14)],
products: [
.executable(name: "AIMemoryMenu", targets: ["AIMemoryMenu"]),
],
targets: [
.target(
name: "AIMemoryMenuCore",
path: "Sources/AIMemoryMenuCore"
),
.executableTarget(
name: "AIMemoryMenu",
dependencies: ["AIMemoryMenuCore"],
path: "Sources/AIMemoryMenu"
),
.testTarget(
name: "AIMemoryMenuTests",
dependencies: ["AIMemoryMenuCore"],
path: "Tests/AIMemoryMenuTests",
resources: [.copy("fixtures")]
),
]
)
companions/ai-memory-macos/README.md
+59
View File
@@ -0,0 +1,59 @@
# AI Memory (macOS menu bar)
Self-contained macOS accessory app that **ships the `ai-memory` runtime**, **governs the LaunchAgent**, and **opens the surfaces the tool already has**. It is a wrapper, not a second operator console.
This companion is not a root Cargo workspace member. Durable memory stays in the user data directory so replacing the `.app` does not rewrite wiki, SQLite, config, models, or logs.
| In the `.app` (replaceable) | In the user data dir (survives updates) |
|---|---|
| Swift menu bar UI | wiki, SQLite, `config.toml` |
| `ai-memory` binary | hook spool, `auth.json`, capture-mode |
| bundled `hooks/` | downloaded embedding models |
| LaunchAgent template | rendered plist in `~/Library/LaunchAgents/` |
| | logs in `~/Library/Logs/ai-memory/` |
Data directory: `~/Library/Application Support/ai-memory` (the binary’s existing macOS default). Optional override in Settings writes `AI_MEMORY_DATA_DIR` into the LaunchAgent plist only.
This app does **not** replace `ai-memory status`, `/web`, or hand-editing `config.toml`. Those stay the real tools; the menu opens them.
## Build
From the repository root (needs a Rust toolchain and Xcode / Swift 6):
```bash
chmod +x companions/ai-memory-macos/build.sh
./companions/ai-memory-macos/build.sh
open "companions/ai-memory-macos/dist/AI Memory.app"
```
`build.sh` compiles `ai-memory` with Cargo, compiles the Swift menu extra, and stages:
```text
AI Memory.app/Contents/Resources/runtime/
ai-memory
hooks/
packaging/launchd/com.github.akitaonrails.ai-memory.plist
```
Drag the `.app` to `/Applications` for a stable LaunchAgent path. Notarization, Developer ID, and a Homebrew cask are out of this companion’s first version.
## Use
1. Open the app (menu bar extra; no Dock icon).
2. **Install & Start Server** — runs bundled `ai-memory init` if `config.toml` is missing, renders the existing launchd template, and `launchctl bootstrap`s `com.github.akitaonrails.ai-memory`.
3. The status item turns green when `GET /admin/status` succeeds.
4. **Open Web UI**, **Show Status…** (bundled `ai-memory status`), **Open Config**, **Open Data Directory**, **Open Logs**.
Updates: replace `/Applications/AI Memory.app`. The data dir is untouched. If the helper path inside the bundle changed, **Restart Server** re-renders the plist.
## Tests
```bash
swift test --package-path companions/ai-memory-macos
```
Root `cargo t` / `cargo tf` do not cover this package.
## Open in Xcode
Open `companions/ai-memory-macos/Package.swift`. `swift run` from the package directory will not include the staged runtime; use `build.sh` (or set `AI_MEMORY_MENU_RUNTIME` at the tarball-equivalent `runtime/` directory) to govern the service.
companions/ai-memory-macos/Sources/AIMemoryMenu/App.swift
+32
View File
@@ -0,0 +1,32 @@
import AppKit
import SwiftUI
import AIMemoryMenuCore
@main
struct AIMemoryMenuApp: App {
@State private var model = AppModel()
init() {
NSApplication.shared.setActivationPolicy(.accessory)
}
var body: some Scene {
MenuBarExtra {
MenuBarView()
.environment(model)
} label: {
MenuBarLabel(icon: model.icon)
}
.menuBarExtraStyle(.menu)
Window("Status", id: "status") {
StatusOutputView()
.environment(model)
}
Settings {
SettingsView()
.environment(model)
}
}
}
companions/ai-memory-macos/Sources/AIMemoryMenu/AppModel.swift
+329
View File
@@ -0,0 +1,329 @@
import AppKit
import Foundation
import AIMemoryMenuCore
import Observation
@MainActor
@Observable
final class AppModel {
var settings: AppSettings
var report: StatusReport?
var launchd: LaunchdState = .notInstalled
var icon: IconState = .unknown
var lastError: String?
var statusOutput: String = ""
var isBusy = false
var runtimeMissing = false
var tokenConfigured = false
var lastHTTPStatus: Int?
@ObservationIgnored private var tokenStore: any TokenStore
@ObservationIgnored private var health = HealthClient()
@ObservationIgnored private var runner: any CommandRunning
@ObservationIgnored private var pollTask: Task<Void, Never>?
@ObservationIgnored private var fetchFailed = false
@ObservationIgnored private var startingDeadline: Date?
init(
settings: AppSettings = .load(),
tokenStore: any TokenStore = KeychainTokenStore(),
runner: any CommandRunning = ProcessRunner()
) {
self.settings = settings
self.tokenStore = tokenStore
self.runner = runner
self.tokenConfigured = tokenStore.read()?.isEmpty == false
startPolling()
}
var headlineVersion: String {
switch icon {
case .ok:
if let report {
return "Server running · v\(report.version)"
}
return "Server running"
case .degraded:
return "Server running · warnings"
case .authRequired:
return "Server running · auth required"
case .starting:
return "Server is starting…"
case .unreachable:
return launchd == .running ? "LaunchAgent up · server unreachable" : "Server down"
case .notInstalled:
return "LaunchAgent not installed"
case .unknown:
return "Checking server…"
}
}
var statisticLines: [String] {
guard let report else {
return []
}
return report.statisticLines
}
var showsServerStatus: Bool {
report != nil
}
var dataDir: URL {
settings.resolvedDataDir
}
var configURL: URL {
dataDir.appending(path: "config.toml")
}
var logURL: URL {
FileManager.default.homeDirectoryForCurrentUser
.appending(path: "Library/Logs/ai-memory/stderr.log")
}
func startPolling() {
guard pollTask == nil else { return }
pollTask = Task { [weak self] in
while let self, !Task.isCancelled {
await self.poll()
let interval: Duration = self.icon == .starting ? .seconds(1) : .seconds(15)
try? await Task.sleep(for: interval)
}
}
}
func poll() async {
let controller = launchdController()
launchd = controller.state()
lastHTTPStatus = nil
do {
report = try await health.fetch(baseURL: settings.serverURL, token: tokenStore.read())
fetchFailed = false
lastError = nil
} catch {
report = nil
fetchFailed = true
if let health = error as? HealthError, case let .http(code) = health {
lastHTTPStatus = code
fetchFailed = code != 401 && code != 403
}
lastError = (error as? HealthError).map(Self.describe) ?? error.localizedDescription
}
let derived = IconState.derived(
launchd: launchd,
report: report,
fetchFailed: fetchFailed,
httpStatus: lastHTTPStatus
)
let starting = startingDeadline.map { Date() < $0 } ?? false
icon = IconState.applyingStartGrace(derived, isStarting: starting)
if icon != .starting {
startingDeadline = nil
}
runtimeMissing = RuntimeLayout.resolve() == nil
tokenConfigured = tokenStore.read()?.isEmpty == false
}
func installAndStart() async {
markStarting()
await runBusy {
let runtime = try self.requireRuntime()
let controller = self.launchdController()
if FirstRun.needsInit(dataDir: self.dataDir) {
_ = try BundledCLI(runtime: runtime, runner: self.runner).initDataDir(self.dataDir)
}
let template = try String(contentsOf: runtime.plistTemplate, encoding: .utf8)
if controller.state() == .running {
try? controller.bootout()
}
try controller.writePlist(
template: template,
binary: runtime.binary,
dataDir: self.settings.dataDirOverride
)
try controller.bootstrap()
}
clearStartingIfFailed()
await poll()
}
func start() async {
markStarting()
await runBusy {
let runtime = try self.requireRuntime()
let controller = self.launchdController()
let template = try String(contentsOf: runtime.plistTemplate, encoding: .utf8)
try controller.writePlist(
template: template,
binary: runtime.binary,
dataDir: self.settings.dataDirOverride
)
try controller.bootstrap()
}
clearStartingIfFailed()
await poll()
}
func stop() async {
startingDeadline = nil
await runBusy {
try self.launchdController().bootout()
}
await poll()
}
func restart() async {
markStarting()
await runBusy {
let controller = self.launchdController()
let runtime = try self.requireRuntime()
let installed = (try? String(contentsOf: controller.plistDestination, encoding: .utf8)) ?? ""
if LaunchdController.programArgumentsBinary(inPlist: installed) != runtime.binary.path {
try? controller.bootout()
let template = try String(contentsOf: runtime.plistTemplate, encoding: .utf8)
try controller.writePlist(
template: template,
binary: runtime.binary,
dataDir: self.settings.dataDirOverride
)
try controller.bootstrap()
} else {
try controller.kickstart()
}
}
clearStartingIfFailed()
await poll()
}
func refreshStatusOutput() async {
await runBusy {
let runtime = try self.requireRuntime()
let result = try BundledCLI(runtime: runtime, runner: self.runner).status(
serverURL: self.settings.serverURL,
token: self.tokenStore.read(),
dataDir: self.settings.dataDirOverride
)
self.statusOutput = result.combinedOutput
}
}
func openWebUI() {
NSWorkspace.shared.open(HealthURL.webUI(from: settings.serverURL))
}
func openConfig() {
revealOrOpen(configURL)
}
func openDataDirectory() {
NSWorkspace.shared.open(dataDir)
}
func openLogs() {
revealOrOpen(logURL)
}
func saveSettings(
serverURLString: String,
dataDirOverride: String,
token: String?,
clearToken: Bool
) {
if let url = URL(string: serverURLString), url.scheme != nil {
settings.serverURL = url
}
let trimmed = dataDirOverride.trimmingCharacters(in: .whitespacesAndNewlines)
settings.dataDirOverride = trimmed.isEmpty ? nil : URL(fileURLWithPath: trimmed)
settings.save()
if clearToken {
try? tokenStore.clear()
} else if let token {
let value = token.trimmingCharacters(in: .whitespacesAndNewlines)
if !value.isEmpty {
try? tokenStore.save(value)
}
}
tokenConfigured = tokenStore.read()?.isEmpty == false
Task { await poll() }
}
private func markStarting() {
startingDeadline = Date().addingTimeInterval(45)
icon = .starting
report = nil
lastError = nil
pollTask?.cancel()
pollTask = nil
startPolling()
}
private func clearStartingIfFailed() {
if lastError != nil {
startingDeadline = nil
}
}
private func launchdController() -> LaunchdController {
LaunchdController(home: FileManager.default.homeDirectoryForCurrentUser, runner: runner)
}
private func requireRuntime() throws -> RuntimeLayout {
guard let runtime = RuntimeLayout.resolve() else {
runtimeMissing = true
throw CliError.missingRuntime
}
return runtime
}
private func runBusy(_ work: () throws -> Void) async {
isBusy = true
defer { isBusy = false }
do {
try work()
lastError = nil
} catch {
lastError = error.localizedDescription
if let cli = error as? CliError, case .failed(_, let output) = cli {
statusOutput = output
lastError = output
}
}
}
private func revealOrOpen(_ url: URL) {
if FileManager.default.fileExists(atPath: url.path) {
NSWorkspace.shared.activateFileViewerSelecting([url])
} else {
NSWorkspace.shared.open(url.deletingLastPathComponent())
}
}
private static func describe(_ error: HealthError) -> String {
switch error {
case .badURL:
"Invalid server URL"
case .http(let code) where code == 401 || code == 403:
"Server is up but this app is not authorized (HTTP \(code)). Add a bearer in Settings."
case .http(let code):
"Server returned HTTP \(code)"
case .decode:
"Could not read /admin/status"
case .transport(let message):
message
}
}
}
extension CliError: LocalizedError {
public var errorDescription: String? {
switch self {
case .timeout:
"Timed out running ai-memory"
case .missingRuntime:
"Bundled ai-memory runtime is missing. Build with companions/ai-memory-macos/build.sh"
case .failed(_, let output):
output.isEmpty ? "ai-memory command failed" : output
}
}
}
companions/ai-memory-macos/Sources/AIMemoryMenu/MenuBarView.swift
+127
View File
@@ -0,0 +1,127 @@
import AppKit
import SwiftUI
import AIMemoryMenuCore
struct MenuBarView: View {
@Environment(AppModel.self) private var model
@Environment(\.openWindow) private var openWindow
var body: some View {
Text(model.headlineVersion)
.disabled(true)
.onAppear {
Task { await model.poll() }
}
ForEach(Array(model.statisticLines.enumerated()), id: \.offset) { _, line in
Text(line)
.font(.system(.body, design: .monospaced))
.disabled(true)
}
if model.runtimeMissing {
Text("Runtime not bundled — run build.sh")
.disabled(true)
}
Divider()
serviceButtons
Divider()
Button("Open Web UI") {
model.openWebUI()
}
if model.showsServerStatus {
Button("Show Status…") {
Task {
await model.refreshStatusOutput()
NSApp.activate(ignoringOtherApps: true)
openWindow(id: "status")
}
}
}
Button("Open Config") {
model.openConfig()
}
Button("Open Data Directory") {
model.openDataDirectory()
}
Button("Open Logs") {
model.openLogs()
}
Divider()
SettingsLink {
Text("Settings…")
}
Button("Quit") {
NSApp.terminate(nil)
}
}
@ViewBuilder
private var serviceButtons: some View {
switch model.launchd {
case .notInstalled:
Button("Install & Start Server") {
Task { await model.installAndStart() }
}
.disabled(model.isBusy || model.icon == .starting)
case .stopped:
Button("Start Server") {
Task { await model.start() }
}
.disabled(model.isBusy || model.icon == .starting)
case .running:
Button("Stop Server") {
Task { await model.stop() }
}
.disabled(model.isBusy)
Button("Restart Server") {
Task { await model.restart() }
}
.disabled(model.isBusy || model.icon == .starting)
}
}
}
struct MenuBarLabel: View {
var icon: IconState
var body: some View {
// Menu extras flatten SwiftUI tint to a template image, so a
// non-template NSImage is what actually changes with server state.
Image(nsImage: StatusDot.image(for: icon))
.accessibilityLabel(icon.accessibilityLabel)
}
}
enum StatusDot {
static func image(for icon: IconState) -> NSImage {
let size = NSSize(width: 18, height: 18)
let image = NSImage(size: size, flipped: false) { rect in
let inset = rect.insetBy(dx: 3, dy: 3)
color(for: icon).setFill()
NSBezierPath(ovalIn: inset).fill()
if icon == .unknown || icon == .notInstalled {
NSColor.windowBackgroundColor.setStroke()
let stroke = NSBezierPath(ovalIn: inset.insetBy(dx: 0.5, dy: 0.5))
stroke.lineWidth = 1
stroke.stroke()
}
return true
}
image.isTemplate = false
return image
}
private static func color(for icon: IconState) -> NSColor {
switch icon {
case .ok:
NSColor.systemGreen
case .starting:
NSColor.systemOrange
case .degraded, .authRequired:
NSColor.systemYellow
case .unreachable:
NSColor.systemRed
case .notInstalled, .unknown:
NSColor.systemGray
}
}
}
companions/ai-memory-macos/Sources/AIMemoryMenu/SettingsView.swift
+104
View File
@@ -0,0 +1,104 @@
import ServiceManagement
import SwiftUI
import AIMemoryMenuCore
struct SettingsView: View {
@Environment(AppModel.self) private var model
@State private var serverURL: String = ""
@State private var dataDir: String = ""
@State private var token: String = ""
@State private var launchAtLogin = SMAppService.mainApp.status == .enabled
@State private var loginError: String?
@State private var saveTask: Task<Void, Never>?
var body: some View {
Form {
Section("Server") {
TextField("URL", text: $serverURL)
.onChange(of: serverURL) { _, _ in scheduleSave() }
.onSubmit { persist() }
SecureField(
model.tokenConfigured ? "Bearer token (saved in Keychain)" : "Bearer token (optional)",
text: $token
)
.onChange(of: token) { _, _ in scheduleSave() }
.onSubmit { persist() }
if model.tokenConfigured {
Button("Clear saved token") {
token = ""
model.saveSettings(
serverURLString: serverURL,
dataDirOverride: dataDir,
token: nil,
clearToken: true
)
}
}
}
Section("Data") {
TextField("Data directory override", text: $dataDir, prompt: Text(FirstRun.defaultDataDir().path))
.onChange(of: dataDir) { _, _ in scheduleSave() }
.onSubmit { persist() }
Text("Empty keeps ~/Library/Application Support/ai-memory. Changing this does not move existing files.")
.font(.caption)
.foregroundStyle(.secondary)
}
Section("Login") {
Toggle("Launch menu bar app at login", isOn: $launchAtLogin)
.onChange(of: launchAtLogin) { _, enabled in
setLoginItem(enabled)
}
if let loginError {
Text(loginError)
.font(.caption)
.foregroundStyle(.red)
}
}
}
.formStyle(.grouped)
.frame(minWidth: 480, minHeight: 320)
.onAppear {
serverURL = model.settings.serverURL.absoluteString
dataDir = model.settings.dataDirOverride?.path ?? ""
launchAtLogin = SMAppService.mainApp.status == .enabled
}
.onDisappear {
persist()
}
}
private func scheduleSave() {
saveTask?.cancel()
saveTask = Task { @MainActor in
try? await Task.sleep(for: .milliseconds(400))
guard !Task.isCancelled else { return }
persist()
}
}
private func persist() {
saveTask?.cancel()
model.saveSettings(
serverURLString: serverURL,
dataDirOverride: dataDir,
token: token.isEmpty ? nil : token,
clearToken: false
)
}
private func setLoginItem(_ enabled: Bool) {
let currentlyEnabled = SMAppService.mainApp.status == .enabled
guard enabled != currentlyEnabled else { return }
do {
if enabled {
try SMAppService.mainApp.register()
} else {
try SMAppService.mainApp.unregister()
}
loginError = nil
} catch {
loginError = error.localizedDescription
launchAtLogin = SMAppService.mainApp.status == .enabled
}
}
}
companions/ai-memory-macos/Sources/AIMemoryMenu/StatusOutputView.swift
+35
View File
@@ -0,0 +1,35 @@
import AppKit
import SwiftUI
struct StatusOutputView: View {
@Environment(AppModel.self) private var model
var body: some View {
VStack(alignment: .leading, spacing: 8) {
HStack {
Text("ai-memory status")
.font(.headline)
Spacer()
Button("Refresh") {
Task { await model.refreshStatusOutput() }
}
.disabled(model.isBusy)
Button("Copy") {
NSPasteboard.general.clearContents()
NSPasteboard.general.setString(model.statusOutput, forType: .string)
}
.disabled(model.statusOutput.isEmpty)
}
ScrollView {
Text(model.statusOutput.isEmpty ? "Running ai-memory status…" : model.statusOutput)
.font(.system(.body, design: .monospaced))
.textSelection(.enabled)
.frame(maxWidth: .infinity, alignment: .leading)
.padding(8)
}
.background(Color(nsColor: .textBackgroundColor))
}
.padding()
.frame(minWidth: 560, minHeight: 360)
}
}
companions/ai-memory-macos/Sources/AIMemoryMenuCore/AppSettings.swift
+49
View File
@@ -0,0 +1,49 @@
import Foundation
public struct AppSettings: Equatable, Sendable {
public var serverURL: URL
public var dataDirOverride: URL?
public static let defaultServerURL = URL(string: "http://127.0.0.1:49374")!
public init(serverURL: URL = defaultServerURL, dataDirOverride: URL? = nil) {
self.serverURL = serverURL
self.dataDirOverride = dataDirOverride
}
public var resolvedDataDir: URL {
dataDirOverride ?? FirstRun.defaultDataDir()
}
public static func load(defaults: UserDefaults = .standard) -> AppSettings {
let url: URL
if let stored = defaults.string(forKey: Keys.serverURL),
let parsed = URL(string: stored)
{
url = parsed
} else {
url = defaultServerURL
}
let override: URL?
if let path = defaults.string(forKey: Keys.dataDir), !path.isEmpty {
override = URL(fileURLWithPath: path)
} else {
override = nil
}
return AppSettings(serverURL: url, dataDirOverride: override)
}
public func save(defaults: UserDefaults = .standard) {
defaults.set(serverURL.absoluteString, forKey: Keys.serverURL)
if let dataDirOverride {
defaults.set(dataDirOverride.path, forKey: Keys.dataDir)
} else {
defaults.removeObject(forKey: Keys.dataDir)
}
}
private enum Keys {
static let serverURL = "serverURL"
static let dataDir = "dataDirOverride"
}
}
companions/ai-memory-macos/Sources/AIMemoryMenuCore/BundledRuntime.swift
+101
View File
@@ -0,0 +1,101 @@
import Foundation
/// Layout of `Contents/Resources/runtime/`: the release-tarball sibling pair
/// (`ai-memory` + `hooks/`) plus the launchd template.
public struct RuntimeLayout: Equatable, Sendable {
public var root: URL
public init(root: URL) {
self.root = root
}
public var binary: URL {
root.appending(path: "ai-memory")
}
public var hooks: URL {
root.appending(path: "hooks")
}
public var plistTemplate: URL {
root.appending(path: "packaging/launchd/\(LaunchdController.label).plist")
}
public func validate(fileManager: FileManager = .default) -> Bool {
fileManager.isExecutableFile(atPath: binary.path)
&& fileManager.fileExists(atPath: hooks.path)
&& fileManager.fileExists(atPath: plistTemplate.path)
}
/// Prefers the staged bundle resource; `AI_MEMORY_MENU_RUNTIME` is a
/// developer override for `swift run` without wrapping an `.app`.
public static func resolve(
bundle: Bundle = .main,
environment: [String: String] = ProcessInfo.processInfo.environment,
fileManager: FileManager = .default
) -> RuntimeLayout? {
if let env = environment["AI_MEMORY_MENU_RUNTIME"], !env.isEmpty {
let layout = RuntimeLayout(root: URL(fileURLWithPath: env))
if layout.validate(fileManager: fileManager) {
return layout
}
}
if let resourceRoot = bundle.resourceURL {
let layout = RuntimeLayout(root: resourceRoot.appending(path: "runtime"))
if layout.validate(fileManager: fileManager) {
return layout
}
}
return nil
}
}
public enum FirstRun {
public static func needsInit(dataDir: URL, fileManager: FileManager = .default) -> Bool {
!fileManager.fileExists(atPath: dataDir.appending(path: "config.toml").path)
}
public static func defaultDataDir(fileManager: FileManager = .default) -> URL {
let base = fileManager.urls(for: .applicationSupportDirectory, in: .userDomainMask).first
?? URL(fileURLWithPath: NSHomeDirectory())
.appending(path: "Library/Application Support")
return base.appending(path: "ai-memory")
}
}
public struct BundledCLI: Sendable {
public var runtime: RuntimeLayout
public var runner: any CommandRunning
public init(runtime: RuntimeLayout, runner: any CommandRunning = ProcessRunner()) {
self.runtime = runtime
self.runner = runner
}
public func initDataDir(_ dataDir: URL) throws -> ProcessResult {
try run(arguments: ["--data-dir", dataDir.path, "init"], extraEnv: [:])
}
public func status(serverURL: URL, token: String?, dataDir: URL?) throws -> ProcessResult {
var args: [String] = []
var env: [String: String] = [
"AI_MEMORY_SERVER_URL": serverURL.absoluteString,
]
if let dataDir {
args.append(contentsOf: ["--data-dir", dataDir.path])
}
if let token, !token.isEmpty {
env["AI_MEMORY_AUTH_TOKEN"] = token
}
args.append("status")
return try run(arguments: args, extraEnv: env)
}
private func run(arguments: [String], extraEnv: [String: String]) throws -> ProcessResult {
let result = try runner.run(binary: runtime.binary, arguments: arguments, extraEnv: extraEnv)
if result.exitCode != 0 {
throw CliError.failed(result.exitCode, result.combinedOutput)
}
return result
}
}
companions/ai-memory-macos/Sources/AIMemoryMenuCore/HealthClient.swift
+60
View File
@@ -0,0 +1,60 @@
import Foundation
public enum HealthError: Error, Equatable, Sendable {
case badURL
case http(Int)
case decode
case transport(String)
}
public struct HealthClient: Sendable {
public var timeout: TimeInterval
private let session: URLSession
public init(timeout: TimeInterval = 3) {
self.timeout = timeout
let config = URLSessionConfiguration.ephemeral
config.timeoutIntervalForRequest = timeout
config.timeoutIntervalForResource = timeout
config.requestCachePolicy = .reloadIgnoringLocalCacheData
config.waitsForConnectivity = false
session = URLSession(configuration: config)
}
public func fetch(baseURL: URL, token: String?) async throws -> StatusReport {
let url = baseURL.appending(path: "admin/status")
var request = URLRequest(url: url)
request.timeoutInterval = timeout
request.cachePolicy = .reloadIgnoringLocalCacheData
request.setValue("application/json", forHTTPHeaderField: "Accept")
if let token, !token.isEmpty {
request.setValue("Bearer \(token)", forHTTPHeaderField: "Authorization")
}
let data: Data
let response: URLResponse
do {
(data, response) = try await session.data(for: request)
} catch {
throw HealthError.transport(error.localizedDescription)
}
guard let http = response as? HTTPURLResponse else {
throw HealthError.transport("non-HTTP response")
}
guard (200 ..< 300).contains(http.statusCode) else {
throw HealthError.http(http.statusCode)
}
do {
return try JSONDecoder().decode(StatusReport.self, from: data)
} catch {
throw HealthError.decode
}
}
}
public enum HealthURL {
public static func webUI(from baseURL: URL) -> URL {
baseURL.appending(path: "web")
}
}
companions/ai-memory-macos/Sources/AIMemoryMenuCore/HealthModels.swift
+281
View File
@@ -0,0 +1,281 @@
import Foundation
/// Wire shape of `GET /admin/status`. Unknown fields are ignored so an
/// older or newer server still decodes the headlines this wrapper shows.
public struct StatusReport: Decodable, Equatable, Sendable {
public var version: String
public var dataDir: String?
public var bind: String?
public var counts: StatusCounts
public var writeQueue: WriteQueue?
public var providers: ProviderHealthSnapshot?
public var ingest: IngestSnapshot?
public init(
version: String,
dataDir: String? = nil,
bind: String? = nil,
counts: StatusCounts,
writeQueue: WriteQueue? = nil,
providers: ProviderHealthSnapshot? = nil,
ingest: IngestSnapshot? = nil
) {
self.version = version
self.dataDir = dataDir
self.bind = bind
self.counts = counts
self.writeQueue = writeQueue
self.providers = providers
self.ingest = ingest
}
enum CodingKeys: String, CodingKey {
case version
case bind
case counts
case providers
case ingest
case dataDir = "data_dir"
case writeQueue = "write_queue"
}
public var isDegraded: Bool {
if let queue = writeQueue, queue.queued > 0 {
return true
}
if providers?.llm.status == "error" {
return true
}
if providers?.embedding.status == "error" {
return true
}
return false
}
public var llmHeadline: String {
roleHeadline(label: "LLM", role: providers?.llm)
}
public var embeddingHeadline: String {
roleHeadline(label: "Embed", role: providers?.embedding)
}
/// Lines shown in the menu extra. Counts come from `GET /admin/status`.
public var statisticLines: [String] {
var lines: [String] = []
if let bind, !bind.isEmpty {
lines.append("Bind \(bind)")
}
lines.append(contentsOf: [
"Pages \(counts.pagesLatest) (all versions \(counts.pagesAll))",
"Sessions \(counts.sessions)",
"Observations \(counts.observations)",
llmHeadline,
embeddingHeadline,
])
if let queue = writeQueue, queue.queued > 0 {
lines.append("Write queue \(queue.queued)/\(queue.capacity)")
}
if let ingest {
lines.append("Ingest accepted \(ingest.accepted)")
if ingest.droppedByPolicy > 0 {
lines.append("Dropped by policy \(ingest.droppedByPolicy)")
}
lines.append("Last write \(Self.lastWriteLabel(ingest.lastPersistedMs))")
}
return lines
}
private func roleHeadline(label: String, role: ProviderRoleHealth?) -> String {
guard let role else {
return "\(label) unknown"
}
var text = "\(label) \(role.status)"
if let provider = role.provider, !provider.isEmpty {
if let model = role.model, !model.isEmpty {
text += " \(provider)/\(model)"
} else {
text += " \(provider)"
}
}
return text
}
public static func lastWriteLabel(_ unixMs: UInt64?) -> String {
guard let unixMs else {
return "—"
}
let nowMs = UInt64(max(0, Date().timeIntervalSince1970 * 1000))
let ageMs = nowMs > unixMs ? nowMs - unixMs : 0
let secs = ageMs / 1000
if secs < 60 {
return "\(secs)s ago"
}
if secs < 3600 {
return "\(secs / 60)m ago"
}
if secs < 86_400 {
return "\(secs / 3600)h ago"
}
return "\(secs / 86_400)d ago"
}
}
public struct IngestSnapshot: Decodable, Equatable, Sendable {
public var accepted: UInt64
public var droppedByPolicy: UInt64
public var shedSaturated: UInt64
public var shedRateLimited: UInt64
public var lastPersistedMs: UInt64?
public init(
accepted: UInt64 = 0,
droppedByPolicy: UInt64 = 0,
shedSaturated: UInt64 = 0,
shedRateLimited: UInt64 = 0,
lastPersistedMs: UInt64? = nil
) {
self.accepted = accepted
self.droppedByPolicy = droppedByPolicy
self.shedSaturated = shedSaturated
self.shedRateLimited = shedRateLimited
self.lastPersistedMs = lastPersistedMs
}
enum CodingKeys: String, CodingKey {
case accepted
case droppedByPolicy = "dropped_by_policy"
case shedSaturated = "shed_saturated"
case shedRateLimited = "shed_rate_limited"
case lastPersistedMs = "last_persisted_ms"
}
public init(from decoder: Decoder) throws {
let container = try decoder.container(keyedBy: CodingKeys.self)
accepted = try container.decodeIfPresent(UInt64.self, forKey: .accepted) ?? 0
droppedByPolicy = try container.decodeIfPresent(UInt64.self, forKey: .droppedByPolicy) ?? 0
shedSaturated = try container.decodeIfPresent(UInt64.self, forKey: .shedSaturated) ?? 0
shedRateLimited = try container.decodeIfPresent(UInt64.self, forKey: .shedRateLimited) ?? 0
lastPersistedMs = try container.decodeIfPresent(UInt64.self, forKey: .lastPersistedMs)
}
}
public struct StatusCounts: Decodable, Equatable, Sendable {
public var pagesLatest: UInt64
public var pagesAll: UInt64
public var sessions: UInt64
public var observations: UInt64
public init(pagesLatest: UInt64, pagesAll: UInt64, sessions: UInt64, observations: UInt64) {
self.pagesLatest = pagesLatest
self.pagesAll = pagesAll
self.sessions = sessions
self.observations = observations
}
enum CodingKeys: String, CodingKey {
case pagesLatest = "pages_latest"
case pagesAll = "pages_all"
case sessions
case observations
}
}
/// JSON encoding of the server's `(queued, capacity)` tuple.
public struct WriteQueue: Decodable, Equatable, Sendable {
public var queued: Int
public var capacity: Int
public init(queued: Int, capacity: Int) {
self.queued = queued
self.capacity = capacity
}
public init(from decoder: Decoder) throws {
var container = try decoder.unkeyedContainer()
queued = try container.decode(Int.self)
capacity = try container.decode(Int.self)
}
}
public struct ProviderHealthSnapshot: Decodable, Equatable, Sendable {
public var llm: ProviderRoleHealth
public var embedding: ProviderRoleHealth
public init(llm: ProviderRoleHealth, embedding: ProviderRoleHealth) {
self.llm = llm
self.embedding = embedding
}
}
public struct ProviderRoleHealth: Decodable, Equatable, Sendable {
public var status: String
public var provider: String?
public var model: String?
public init(status: String, provider: String? = nil, model: String? = nil) {
self.status = status
self.provider = provider
self.model = model
}
}
public enum IconState: Equatable, Sendable {
case unknown
case notInstalled
case unreachable
case starting
case authRequired
case ok
case degraded
public static func derived(
launchd: LaunchdState,
report: StatusReport?,
fetchFailed: Bool,
httpStatus: Int? = nil
) -> IconState {
if let report {
return report.isDegraded ? .degraded : .ok
}
if let httpStatus, httpStatus == 401 || httpStatus == 403 {
return .authRequired
}
if fetchFailed {
return launchd == .notInstalled ? .notInstalled : .unreachable
}
return .unknown
}
/// Keep the "starting" overlay until `/admin/status` answers or the grace expires.
public static func applyingStartGrace(_ derived: IconState, isStarting: Bool) -> IconState {
guard isStarting else {
return derived
}
switch derived {
case .ok, .degraded, .authRequired:
return derived
default:
return .starting
}
}
public var accessibilityLabel: String {
switch self {
case .unknown:
"ai-memory, status unknown"
case .notInstalled:
"ai-memory, not installed"
case .unreachable:
"ai-memory, server down"
case .starting:
"ai-memory, server is starting"
case .authRequired:
"ai-memory, running, authentication required"
case .ok:
"ai-memory, server running"
case .degraded:
"ai-memory, running with warnings"
}
}
}
companions/ai-memory-macos/Sources/AIMemoryMenuCore/LaunchdController.swift
+136
View File
@@ -0,0 +1,136 @@
import Darwin
import Foundation
public enum LaunchdState: Equatable, Sendable {
case notInstalled
case stopped
case running
}
public struct LaunchdController: Sendable {
public static let label = "com.github.akitaonrails.ai-memory"
public var home: URL
public var runner: any CommandRunning
public init(home: URL, runner: any CommandRunning = ProcessRunner()) {
self.home = home
self.runner = runner
}
public var plistDestination: URL {
home.appending(path: "Library/LaunchAgents")
.appending(path: "\(Self.label).plist")
}
public var logDirectory: URL {
home.appending(path: "Library/Logs/ai-memory")
}
public static func renderTemplate(
_ template: String,
binary: URL,
home: URL,
dataDir: URL?
) -> String {
var rendered = template
.replacingOccurrences(of: "__AI_MEMORY_BIN__", with: binary.path)
.replacingOccurrences(of: "__HOME__", with: home.path)
if let dataDir {
let env = """
<key>EnvironmentVariables</key>
<dict>
<key>AI_MEMORY_DATA_DIR</key>
<string>\(xmlEscape(dataDir.path))</string>
</dict>
"""
if let range = rendered.range(of: "</dict>", options: .backwards) {
rendered.replaceSubrange(range, with: env + "</dict>")
}
}
return rendered
}
public static func programArgumentsBinary(inPlist plist: String) -> String? {
// First <string> after ProgramArguments is the executable.
guard let argsRange = plist.range(of: "<key>ProgramArguments</key>") else {
return nil
}
let rest = plist[argsRange.upperBound...]
guard let start = rest.range(of: "<string>") else {
return nil
}
let after = rest[start.upperBound...]
guard let end = after.range(of: "</string>") else {
return nil
}
return String(after[..<end.lowerBound])
}
public func writePlist(template: String, binary: URL, dataDir: URL?) throws {
let fm = FileManager.default
try fm.createDirectory(
at: plistDestination.deletingLastPathComponent(),
withIntermediateDirectories: true
)
try fm.createDirectory(at: logDirectory, withIntermediateDirectories: true)
let body = Self.renderTemplate(template, binary: binary, home: home, dataDir: dataDir)
try body.write(to: plistDestination, atomically: true, encoding: .utf8)
}
public func state() -> LaunchdState {
let plistExists = FileManager.default.fileExists(atPath: plistDestination.path)
let result = try? runner.run(
binary: URL(fileURLWithPath: "/bin/launchctl"),
arguments: ["print", domainService],
extraEnv: [:]
)
if let result, result.exitCode == 0 {
if result.stdout.contains("state = running") {
return .running
}
return .stopped
}
return plistExists ? .stopped : .notInstalled
}
public func bootstrap() throws {
try runLaunchctl(["bootstrap", domain, plistDestination.path])
}
public func bootout() throws {
try runLaunchctl(["bootout", domainService])
}
public func kickstart() throws {
try runLaunchctl(["kickstart", "-k", domainService])
}
private func runLaunchctl(_ arguments: [String]) throws {
let result = try runner.run(
binary: URL(fileURLWithPath: "/bin/launchctl"),
arguments: arguments,
extraEnv: [:]
)
if result.exitCode != 0 {
throw CliError.failed(result.exitCode, result.combinedOutput)
}
}
private var domain: String {
"gui/\(getuid())"
}
private var domainService: String {
"\(domain)/\(Self.label)"
}
private static func xmlEscape(_ value: String) -> String {
value
.replacingOccurrences(of: "&", with: "&amp;")
.replacingOccurrences(of: "<", with: "&lt;")
.replacingOccurrences(of: ">", with: "&gt;")
.replacingOccurrences(of: "\"", with: "&quot;")
}
}
companions/ai-memory-macos/Sources/AIMemoryMenuCore/ProcessRunner.swift
+78
View File
@@ -0,0 +1,78 @@
import Foundation
public struct ProcessResult: Equatable, Sendable {
public var exitCode: Int32
public var stdout: String
public var stderr: String
public init(exitCode: Int32, stdout: String, stderr: String) {
self.exitCode = exitCode
self.stdout = stdout
self.stderr = stderr
}
public var combinedOutput: String {
let out = stdout.trimmingCharacters(in: .whitespacesAndNewlines)
let err = stderr.trimmingCharacters(in: .whitespacesAndNewlines)
if err.isEmpty {
return out
}
if out.isEmpty {
return err
}
return out + "\n" + err
}
}
public protocol CommandRunning: Sendable {
func run(binary: URL, arguments: [String], extraEnv: [String: String]) throws -> ProcessResult
}
public struct ProcessRunner: CommandRunning {
public var timeout: TimeInterval
public init(timeout: TimeInterval = 30) {
self.timeout = timeout
}
public func run(binary: URL, arguments: [String], extraEnv: [String: String]) throws -> ProcessResult {
let process = Process()
process.executableURL = binary
process.arguments = arguments
var env = ProcessInfo.processInfo.environment
for (key, value) in extraEnv {
env[key] = value
}
process.environment = env
let stdout = Pipe()
let stderr = Pipe()
process.standardOutput = stdout
process.standardError = stderr
try process.run()
let deadline = Date().addingTimeInterval(timeout)
while process.isRunning, Date() < deadline {
Thread.sleep(forTimeInterval: 0.05)
}
if process.isRunning {
process.terminate()
throw CliError.timeout
}
let outData = stdout.fileHandleForReading.readDataToEndOfFile()
let errData = stderr.fileHandleForReading.readDataToEndOfFile()
return ProcessResult(
exitCode: process.terminationStatus,
stdout: String(data: outData, encoding: .utf8) ?? "",
stderr: String(data: errData, encoding: .utf8) ?? ""
)
}
}
public enum CliError: Error, Equatable, Sendable {
case timeout
case missingRuntime
case failed(Int32, String)
}
companions/ai-memory-macos/Sources/AIMemoryMenuCore/TokenStore.swift
+106
View File
@@ -0,0 +1,106 @@
import Foundation
import Security
public protocol TokenStore: Sendable {
func read() -> String?
func save(_ token: String) throws
func clear() throws
}
public struct MemoryTokenStore: TokenStore, Sendable {
private let box: LockingBox<String?>
public init(_ initial: String? = nil) {
box = LockingBox(initial)
}
public func read() -> String? {
box.value
}
public func save(_ token: String) throws {
box.value = token
}
public func clear() throws {
box.value = nil
}
}
/// Tiny mutex so `MemoryTokenStore` can be `Sendable` in tests.
final class LockingBox<Value>: @unchecked Sendable {
private let lock = NSLock()
private var storage: Value
init(_ value: Value) {
storage = value
}
var value: Value {
get {
lock.lock()
defer { lock.unlock() }
return storage
}
set {
lock.lock()
defer { lock.unlock() }
storage = newValue
}
}
}
public struct KeychainTokenStore: TokenStore, Sendable {
public static let service = "com.github.akitaonrails.ai-memory-menu"
public static let account = "bearer"
public init() {}
public func read() -> String? {
let query: [String: Any] = [
kSecClass as String: kSecClassGenericPassword,
kSecAttrService as String: Self.service,
kSecAttrAccount as String: Self.account,
kSecReturnData as String: true,
kSecMatchLimit as String: kSecMatchLimitOne,
]
var item: CFTypeRef?
let status = SecItemCopyMatching(query as CFDictionary, &item)
guard status == errSecSuccess, let data = item as? Data else {
return nil
}
return String(data: data, encoding: .utf8)
}
public func save(_ token: String) throws {
try clear()
let data = Data(token.utf8)
let query: [String: Any] = [
kSecClass as String: kSecClassGenericPassword,
kSecAttrService as String: Self.service,
kSecAttrAccount as String: Self.account,
kSecValueData as String: data,
kSecAttrAccessible as String: kSecAttrAccessibleAfterFirstUnlock,
]
let status = SecItemAdd(query as CFDictionary, nil)
guard status == errSecSuccess else {
throw KeychainError.unhandled(status)
}
}
public func clear() throws {
let query: [String: Any] = [
kSecClass as String: kSecClassGenericPassword,
kSecAttrService as String: Self.service,
kSecAttrAccount as String: Self.account,
]
let status = SecItemDelete(query as CFDictionary)
guard status == errSecSuccess || status == errSecItemNotFound else {
throw KeychainError.unhandled(status)
}
}
}
public enum KeychainError: Error, Equatable, Sendable {
case unhandled(OSStatus)
}
companions/ai-memory-macos/Tests/AIMemoryMenuTests/AppSettingsTests.swift
+23
View File
@@ -0,0 +1,23 @@
import Foundation
import Testing
@testable import AIMemoryMenuCore
struct AppSettingsTests {
@Test func loadAndSaveRoundTrip() throws {
let name = "ai-memory-menu-settings-test-\(UUID().uuidString)"
let suite = try #require(UserDefaults(suiteName: name))
defer { suite.removePersistentDomain(forName: name) }
var settings = AppSettings.load(defaults: suite)
#expect(settings.serverURL == AppSettings.defaultServerURL)
#expect(settings.dataDirOverride == nil)
settings.serverURL = URL(string: "http://127.0.0.1:8080")!
settings.dataDirOverride = URL(fileURLWithPath: "/tmp/custom-memory")
settings.save(defaults: suite)
let loaded = AppSettings.load(defaults: suite)
#expect(loaded.serverURL.absoluteString == "http://127.0.0.1:8080")
#expect(loaded.dataDirOverride?.path == "/tmp/custom-memory")
}
}
companions/ai-memory-macos/Tests/AIMemoryMenuTests/HealthModelsTests.swift
+108
View File
@@ -0,0 +1,108 @@
import Foundation
import Testing
@testable import AIMemoryMenuCore
struct HealthModelsTests {
@Test func decodesStatusFixture() throws {
let url = try #require(Bundle.module.url(forResource: "status", withExtension: "json", subdirectory: "fixtures"))
let data = try Data(contentsOf: url)
let report = try JSONDecoder().decode(StatusReport.self, from: data)
#expect(report.version == "2.3.2")
#expect(report.counts.pagesLatest == 138)
#expect(report.counts.sessions == 27)
#expect(report.writeQueue == WriteQueue(queued: 0, capacity: 128))
#expect(report.providers?.llm.status == "ok")
#expect(report.providers?.embedding.status == "disabled")
#expect(report.ingest?.accepted == 4198)
#expect(!report.isDegraded)
#expect(report.llmHeadline.contains("ok"))
let stats = report.statisticLines
#expect(stats.contains { $0.hasPrefix("Pages 138") })
#expect(stats.contains { $0.hasPrefix("Sessions 27") })
#expect(stats.contains { $0.hasPrefix("Observations 4198") })
#expect(stats.contains { $0.hasPrefix("Bind ") })
#expect(stats.contains { $0.hasPrefix("Ingest accepted 4198") })
}
@Test func decodesOlderPayloadWithoutWriteQueueOrProviders() throws {
let json = """
{
"version": "1.0.0",
"counts": {
"pages_latest": 1,
"pages_all": 1,
"sessions": 0,
"observations": 0
}
}
""".data(using: .utf8)!
let report = try JSONDecoder().decode(StatusReport.self, from: json)
#expect(report.writeQueue == nil)
#expect(report.providers == nil)
#expect(!report.isDegraded)
}
@Test func degradedWhenProviderErrorsOrQueueIsBusy() {
let ok = StatusReport(
version: "2.3.2",
counts: StatusCounts(pagesLatest: 1, pagesAll: 1, sessions: 0, observations: 0),
writeQueue: WriteQueue(queued: 0, capacity: 8),
providers: ProviderHealthSnapshot(
llm: ProviderRoleHealth(status: "ok"),
embedding: ProviderRoleHealth(status: "disabled")
)
)
#expect(!ok.isDegraded)
var queued = ok
queued.writeQueue = WriteQueue(queued: 1, capacity: 8)
#expect(queued.isDegraded)
var llmError = ok
llmError.writeQueue = WriteQueue(queued: 0, capacity: 8)
llmError.providers = ProviderHealthSnapshot(
llm: ProviderRoleHealth(status: "error"),
embedding: ProviderRoleHealth(status: "ok")
)
#expect(llmError.isDegraded)
var embedError = ok
embedError.providers = ProviderHealthSnapshot(
llm: ProviderRoleHealth(status: "ok"),
embedding: ProviderRoleHealth(status: "error")
)
#expect(embedError.isDegraded)
}
}
struct IconStateTests {
@Test func derivedStates() {
let report = StatusReport(
version: "2.3.2",
counts: StatusCounts(pagesLatest: 1, pagesAll: 1, sessions: 0, observations: 0)
)
#expect(IconState.derived(launchd: .running, report: report, fetchFailed: false) == .ok)
var degraded = report
degraded.writeQueue = WriteQueue(queued: 3, capacity: 8)
#expect(IconState.derived(launchd: .running, report: degraded, fetchFailed: false) == .degraded)
#expect(IconState.derived(launchd: .notInstalled, report: nil, fetchFailed: true) == .notInstalled)
#expect(IconState.derived(launchd: .stopped, report: nil, fetchFailed: true) == .unreachable)
#expect(IconState.derived(launchd: .running, report: nil, fetchFailed: true) == .unreachable)
#expect(IconState.derived(launchd: .notInstalled, report: nil, fetchFailed: false) == .unknown)
#expect(
IconState.derived(
launchd: .running,
report: nil,
fetchFailed: true,
httpStatus: 401
) == .authRequired
)
#expect(
IconState.applyingStartGrace(.unreachable, isStarting: true) == .starting
)
#expect(IconState.applyingStartGrace(.ok, isStarting: true) == .ok)
#expect(IconState.applyingStartGrace(.unreachable, isStarting: false) == .unreachable)
}
}
companions/ai-memory-macos/Tests/AIMemoryMenuTests/LaunchdControllerTests.swift
+74
View File
@@ -0,0 +1,74 @@
import Foundation
import Testing
@testable import AIMemoryMenuCore
struct LaunchdControllerTests {
@Test func substitutesPlaceholders() {
let template = """
<string>__AI_MEMORY_BIN__</string>
<string>__HOME__/Library/Logs/ai-memory/stderr.log</string>
"""
let rendered = LaunchdController.renderTemplate(
template,
binary: URL(fileURLWithPath: "/Applications/AI Memory.app/Contents/Resources/runtime/ai-memory"),
home: URL(fileURLWithPath: "/Users/ada"),
dataDir: nil
)
#expect(rendered.contains("/Applications/AI Memory.app/Contents/Resources/runtime/ai-memory"))
#expect(rendered.contains("/Users/ada/Library/Logs/ai-memory/stderr.log"))
#expect(!rendered.contains("__AI_MEMORY_BIN__"))
#expect(!rendered.contains("EnvironmentVariables"))
}
@Test func injectsDataDirEnvironment() {
let template = """
<dict>
<key>Label</key>
<string>com.github.akitaonrails.ai-memory</string>
</dict>
"""
let rendered = LaunchdController.renderTemplate(
template,
binary: URL(fileURLWithPath: "/bin/ai-memory"),
home: URL(fileURLWithPath: "/Users/ada"),
dataDir: URL(fileURLWithPath: "/Users/ada/.ai-memory")
)
#expect(rendered.contains("<key>AI_MEMORY_DATA_DIR</key>"))
#expect(rendered.contains("<string>/Users/ada/.ai-memory</string>"))
#expect(rendered.contains("<key>EnvironmentVariables</key>"))
}
@Test func rendersCheckedInLaunchdTemplate() throws {
let templateURL = repoRoot()
.appending(path: "packaging/launchd/com.github.akitaonrails.ai-memory.plist")
let template = try String(contentsOf: templateURL, encoding: .utf8)
let binary = URL(fileURLWithPath: "/Applications/AI Memory.app/Contents/Resources/runtime/ai-memory")
let home = URL(fileURLWithPath: "/Users/ada")
let rendered = LaunchdController.renderTemplate(template, binary: binary, home: home, dataDir: nil)
#expect(LaunchdController.programArgumentsBinary(inPlist: rendered) == binary.path)
#expect(rendered.contains("/Users/ada/Library/Logs/ai-memory/stderr.log"))
#expect(rendered.contains("serve"))
#expect(rendered.contains("--enable-web"))
#expect(!rendered.contains("__HOME__"))
}
@Test func xmlEscapesDataDir() {
let template = "<dict></dict>"
let rendered = LaunchdController.renderTemplate(
template,
binary: URL(fileURLWithPath: "/bin/ai-memory"),
home: URL(fileURLWithPath: "/Users/ada"),
dataDir: URL(fileURLWithPath: "/tmp/a&b<c>")
)
#expect(rendered.contains("/tmp/a&amp;b&lt;c&gt;"))
}
}
private func repoRoot(file: String = #filePath) -> URL {
URL(fileURLWithPath: file)
.deletingLastPathComponent() // Tests/AIMemoryMenuTests
.deletingLastPathComponent() // Tests
.deletingLastPathComponent() // companions/ai-memory-macos
.deletingLastPathComponent() // companions
.deletingLastPathComponent() // repo
}
companions/ai-memory-macos/Tests/AIMemoryMenuTests/RuntimeAndCliTests.swift
+132
View File
@@ -0,0 +1,132 @@
import Foundation
import Testing
@testable import AIMemoryMenuCore
struct RuntimeAndCliTests {
@Test func layoutPointsAtTarballSiblings() {
let root = URL(fileURLWithPath: "/tmp/runtime")
let layout = RuntimeLayout(root: root)
#expect(layout.binary.path.hasSuffix("/runtime/ai-memory"))
#expect(layout.hooks.path.hasSuffix("/runtime/hooks"))
#expect(layout.plistTemplate.path.hasSuffix("/packaging/launchd/com.github.akitaonrails.ai-memory.plist"))
}
@Test func validateRequiresBinaryHooksAndPlist() throws {
let dir = FileManager.default.temporaryDirectory
.appending(path: "ai-memory-menu-runtime-\(UUID().uuidString)")
defer { try? FileManager.default.removeItem(at: dir) }
let layout = RuntimeLayout(root: dir)
#expect(!layout.validate())
try FileManager.default.createDirectory(at: layout.hooks, withIntermediateDirectories: true)
try FileManager.default.createDirectory(
at: layout.plistTemplate.deletingLastPathComponent(),
withIntermediateDirectories: true
)
try Data().write(to: layout.binary)
try Data().write(to: layout.plistTemplate)
#expect(!layout.validate())
try FileManager.default.setAttributes([.posixPermissions: 0o755], ofItemAtPath: layout.binary.path)
#expect(layout.validate())
}
@Test func resolvePrefersEnvironmentOverride() throws {
let dir = FileManager.default.temporaryDirectory
.appending(path: "ai-memory-menu-env-\(UUID().uuidString)")
defer { try? FileManager.default.removeItem(at: dir) }
let layout = RuntimeLayout(root: dir)
try FileManager.default.createDirectory(at: layout.hooks, withIntermediateDirectories: true)
try FileManager.default.createDirectory(
at: layout.plistTemplate.deletingLastPathComponent(),
withIntermediateDirectories: true
)
try Data().write(to: layout.binary)
try FileManager.default.setAttributes([.posixPermissions: 0o755], ofItemAtPath: layout.binary.path)
try Data().write(to: layout.plistTemplate)
let resolved = RuntimeLayout.resolve(
bundle: Bundle.main,
environment: ["AI_MEMORY_MENU_RUNTIME": dir.path]
)
#expect(resolved?.root.path == dir.path)
}
@Test func needsInitOnlyWhenConfigMissing() throws {
let dir = FileManager.default.temporaryDirectory
.appending(path: "ai-memory-menu-init-\(UUID().uuidString)")
defer { try? FileManager.default.removeItem(at: dir) }
try FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true)
#expect(FirstRun.needsInit(dataDir: dir))
try "bind = \"127.0.0.1:49374\"\n".write(
to: dir.appending(path: "config.toml"),
atomically: true,
encoding: .utf8
)
#expect(!FirstRun.needsInit(dataDir: dir))
}
@Test func bundledCLIInvokesInitAndStatus() throws {
let runner = MockRunner()
runner.results.append(ProcessResult(exitCode: 0, stdout: "initialized\n", stderr: ""))
runner.results.append(ProcessResult(exitCode: 0, stdout: "ai-memory 2.3.2 (server)\n", stderr: ""))
let cli = BundledCLI(
runtime: RuntimeLayout(root: URL(fileURLWithPath: "/tmp/runtime")),
runner: runner
)
let dataDir = URL(fileURLWithPath: "/tmp/data")
_ = try cli.initDataDir(dataDir)
_ = try cli.status(
serverURL: URL(string: "http://127.0.0.1:49374")!,
token: "secret",
dataDir: dataDir
)
#expect(runner.calls.count == 2)
#expect(runner.calls[0].arguments == ["--data-dir", "/tmp/data", "init"])
#expect(runner.calls[1].arguments == ["--data-dir", "/tmp/data", "status"])
#expect(runner.calls[1].extraEnv["AI_MEMORY_SERVER_URL"] == "http://127.0.0.1:49374")
#expect(runner.calls[1].extraEnv["AI_MEMORY_AUTH_TOKEN"] == "secret")
}
@Test func bundledCLISurfacesNonZeroExit() {
let runner = MockRunner()
runner.results.append(ProcessResult(exitCode: 2, stdout: "", stderr: "could not reach server\n"))
let cli = BundledCLI(
runtime: RuntimeLayout(root: URL(fileURLWithPath: "/tmp/runtime")),
runner: runner
)
do {
_ = try cli.status(
serverURL: URL(string: "http://127.0.0.1:49374")!,
token: nil,
dataDir: nil
)
Issue.record("expected failure")
} catch let CliError.failed(code, output) {
#expect(code == 2)
#expect(output.contains("could not reach server"))
} catch {
Issue.record("wrong error \(error)")
}
}
}
private final class MockRunner: CommandRunning, @unchecked Sendable {
struct Call {
var binary: URL
var arguments: [String]
var extraEnv: [String: String]
}
var calls: [Call] = []
var results: [ProcessResult] = []
func run(binary: URL, arguments: [String], extraEnv: [String: String]) throws -> ProcessResult {
calls.append(Call(binary: binary, arguments: arguments, extraEnv: extraEnv))
if results.isEmpty {
return ProcessResult(exitCode: 0, stdout: "", stderr: "")
}
return results.removeFirst()
}
}
companions/ai-memory-macos/Tests/AIMemoryMenuTests/fixtures/status.json
+31
View File
@@ -0,0 +1,31 @@
{
"version": "2.3.2",
"data_dir": "/Users/test/Library/Application Support/ai-memory",
"bind": "127.0.0.1:49374",
"db_path": "/Users/test/Library/Application Support/ai-memory/db/memory.sqlite",
"counts": {
"pages_latest": 138,
"pages_all": 162,
"sessions": 27,
"observations": 4198
},
"write_queue": [0, 128],
"providers": {
"llm": {
"status": "ok",
"provider": "openai",
"model": "gpt-4.1"
},
"embedding": {
"status": "disabled"
},
"llm_candidates": []
},
"ingest": {
"accepted": 4198,
"dropped_by_policy": 12,
"shed_saturated": 0,
"shed_rate_limited": 0,
"last_persisted_ms": 1700000000000
}
}
companions/ai-memory-macos/build.sh
Executable
+51
View File
@@ -0,0 +1,51 @@
#!/usr/bin/env bash
# Build AI Memory.app: Swift menu bar + staged ai-memory runtime (binary + hooks).
set -euo pipefail
COMPANION="$(cd "$(dirname "$0")" && pwd)"
ROOT="$(cd "$COMPANION/../.." && pwd)"
DIST="${1:-$COMPANION/dist/AI Memory.app}"
CONFIG="${CONFIGURATION:-release}"
echo "building ai-memory ($CONFIG) from $ROOT"
if [[ "$CONFIG" == "release" ]]; then
cargo build --release --bin ai-memory --manifest-path "$ROOT/Cargo.toml"
BINARY="$ROOT/target/release/ai-memory"
SWIFT_FLAGS=(-c release)
else
cargo build --bin ai-memory --manifest-path "$ROOT/Cargo.toml"
BINARY="$ROOT/target/debug/ai-memory"
SWIFT_FLAGS=(-c debug)
fi
echo "building AIMemoryMenu"
# SwiftUI @State needs libSwiftUIMacros; Command Line Tools alone do not ship it.
XCODE_DEVELOPER="${DEVELOPER_DIR:-/Applications/Xcode.app/Contents/Developer}"
if [[ ! -d "$XCODE_DEVELOPER/Platforms/MacOSX.platform" ]]; then
echo "error: install Xcode or set DEVELOPER_DIR to Xcode.app/Contents/Developer" >&2
exit 1
fi
export DEVELOPER_DIR="$XCODE_DEVELOPER"
swift build "${SWIFT_FLAGS[@]}" --package-path "$COMPANION"
BIN_PATH="$(swift build "${SWIFT_FLAGS[@]}" --package-path "$COMPANION" --show-bin-path)"
APP="$DIST"
echo "staging $APP"
rm -rf "$APP"
mkdir -p "$APP/Contents/MacOS"
mkdir -p "$APP/Contents/Resources/runtime/packaging/launchd"
cp "$BIN_PATH/AIMemoryMenu" "$APP/Contents/MacOS/AIMemoryMenu"
cp "$COMPANION/Info.plist" "$APP/Contents/Info.plist"
printf 'APPL????' > "$APP/Contents/PkgInfo"
cp "$BINARY" "$APP/Contents/Resources/runtime/ai-memory"
chmod 755 "$APP/Contents/Resources/runtime/ai-memory"
rsync -a --delete "$ROOT/hooks/" "$APP/Contents/Resources/runtime/hooks/"
cp "$ROOT/packaging/launchd/com.github.akitaonrails.ai-memory.plist" \
"$APP/Contents/Resources/runtime/packaging/launchd/"
echo "built $APP"
echo "runtime: $APP/Contents/Resources/runtime/ai-memory"
echo "open with: open \"$APP\""
companions/ai-memory-relay/Cargo.lock
Generated
+1717
View File
File diff suppressed because it is too large Load Diff
companions/ai-memory-relay/Cargo.toml
+38
View File
@@ -0,0 +1,38 @@
# Separate workspace, as described in docs/companion-crates.md.
[workspace]
[package]
name = "ai-memory-relay"
version = "0.1.0"
edition = "2024"
publish = false
license = "MIT"
description = "Persistent queue for external ai-memory lifecycle events"
[[bin]]
name = "ai-memory-relay"
path = "src/main.rs"
[lints.rust]
unsafe_code = "forbid"
[dependencies]
anyhow = "1"
# Read the bearer token separately to keep it out of clap's Debug output.
clap = { version = "4", features = ["derive"] }
fs2 = "0.4"
# Native roots match the workspace (#492).
reqwest = { version = "0.12", default-features = false, features = [
"blocking",
"json",
"rustls-tls-native-roots",
] }
rusqlite = { version = "0.32", features = ["bundled"] }
serde = { version = "1", features = ["derive"] }
# Default map ordering makes content comparisons independent of input key order.
serde_json = "1"
sha2 = "0.10"
url = "2"
[dev-dependencies]
tempfile = "3"
companions/ai-memory-relay/README.md
+215
View File
@@ -0,0 +1,215 @@
# ai-memory-relay
A companion CLI for sending externally captured lifecycle events to ai-memory.
It records events in a local SQLite queue and sends them through `POST /hook/batch`.
Use it when an orchestrator needs delivery to survive an unavailable server or a
process restart.
The relay has its own Cargo workspace. It does not open ai-memory's database or
wiki, launch agents, or claim handoffs. See the
[external lifecycle contract](../../docs/external-lifecycle.md) for capture
ownership and the [companion policy](../../docs/companion-crates.md) for the core
boundary.
## Build and use
From the repository root:
```sh
cargo build --locked --manifest-path companions/ai-memory-relay/Cargo.toml
```
The executable is `companions/ai-memory-relay/target/debug/ai-memory-relay`, unless
`CARGO_TARGET_DIR` selects another target directory. The commands below assume
the executable is on `PATH` and an ai-memory server is running.
Initialize a queue. Its parent directory must exist; the relay creates the queue
directory itself.
```sh
ai-memory-relay init \
--queue-dir "$HOME/.ai-memory-relay" \
--server-url http://127.0.0.1:49374 \
--producer my-orchestrator \
--actor developer \
--workspace work \
--project example
```
This binds the queue to one destination and scope. Repeating an identical `init`
is safe. A different binding is rejected. `actor` is a stable namespace used to
derive retry keys; authentication still comes from the server's bearer token.
It does not select or impersonate a server user.
When the orchestrator launches a harness whose lifecycle it captures, pass
`AI_MEMORY_CAPTURE_OWNER=my-orchestrator` to that process. Native capture is then
suppressed as described in the lifecycle contract. Supported handoff delivery
remains active, and the agent can keep using MCP for retrieval and deliberate
memory writes. A standalone harness without that context keeps normal hooks.
The orchestrator supplies a JSON array, for example `events.json`:
```json
[
{
"event_id": "event-001",
"agent": "claude-code",
"event": "session-start",
"body": {
"session_id": "native-session-123",
"cwd": "/workspace/example"
}
},
{
"event_id": "event-002",
"agent": "claude-code",
"event": "user-prompt-submit",
"body": {
"session_id": "native-session-123",
"cwd": "/workspace/example",
"prompt": "Check the failing test."
}
}
]
```
```sh
ai-memory-relay enqueue --queue-dir "$HOME/.ai-memory-relay" --file events.json
ai-memory-relay flush --queue-dir "$HOME/.ai-memory-relay"
ai-memory-relay status --queue-dir "$HOME/.ai-memory-relay"
```
`enqueue` validates the whole array and commits it in one transaction. Each body
must be an object with explicit `session_id` and `cwd` strings. Use the harness's
native identity and canonical hook event names. Other body fields keep their
values, including an `_ai_memory_capture` block. The relay does not translate
provider transcripts or infer parent agents and workflows.
For authenticated servers, supply `AI_MEMORY_AUTH_TOKEN` through the flush
process's environment. The relay does not store or print it. The destination
must use HTTP or HTTPS, with no credentials, query string or fragment in the
base URL. Redirects are refused.
## Delivery and retries
Each event gets a deterministic `ingest_key` from the producer, actor, native
agent/session identity, event name and producer-assigned `event_id`, following
the lifecycle contract. `extension` carries the producer and `source_event`
carries the event name.
Re-enqueueing the same identity and body is recognized while the pending event
or receipt remains in the queue. JSON object key order does not matter. Reusing
that identity with different content rejects the entire input array. Two equal
bodies with different event IDs remain separate events.
Only the oldest pending event from each session enters a batch. The next event
for that session is eligible after acknowledgement. This keeps `session-end`
behind earlier events while allowing other sessions to proceed. Order is the
queue's committed enqueue order; producers must serialize events within a
session before enqueueing them. Separate queue directories do not coordinate
session order.
One process can flush a queue at a time. Other processes can enqueue during a
flush. HTTP runs outside the SQLite write transaction. If a response is lost
after the server commits, the event stays pending and a retry uses the same key.
The relay validates the complete batch acknowledgement before changing its
queue. It honors both contiguous and noncontiguous accepted indexes, including
partial success in HTTP 200 or 429 responses. Malformed acknowledgements,
authentication failures and transport errors retain unacknowledged events.
A failed head blocks its own session. HTTP 429 ends the flush; the caller should
back off before scheduling another one.
An acknowledgement can also mean the server deliberately dropped an event, for
example because of capture policy or a session identity collision. A drained
queue does not prove that every event became an observation. The relay catches
a native session changing agents within one queue, but server ownership checks
still apply across queues and producers.
The server's ingest keys expire after 30 days. The relay records the first
attempt before sending and retains expired pending events without replaying
them automatically. Unsent events can remain offline longer. Keep the host
clock accurate; the retry deadline uses wall-clock timestamps. Recreating a
queue discards its attempt history and receipts, so it is not a safe way to
recover an ambiguous delivery after the retry window.
`flush` attempts at most 64 batches by default; `--max-batches` changes that
count. Requests, transport retries and the total flush duration also have finite
bounds. Unfinished work stays queued for the next invocation.
Exit codes are:
| Code | Meaning |
| --- | --- |
| `0` | Command succeeded; a flush has no pending events. |
| `2` | Command failed. Inspect stderr; unacknowledged events remain queued. |
| `3` | Flush ended with events still pending. |
`status` emits a JSON object with counters, including `pending_items`,
`pending_bytes`, `pending_sessions`, `expired_items` and `receipts`. It does not
print event bodies. Receipts count acknowledgements, including policy drops.
## Local data and recovery
The queue contains producer-supplied event bodies before server sanitization.
The producer must apply its capture exclusions before enqueueing. The relay
does not load the project's `[capture] ignore_paths` settings or inspect tool
payloads for sensitive paths. Server sanitization cannot protect a local queue
that already contains those payloads.
On Unix, new queue directories use mode `0700` and files use `0600`. Existing
queue paths must already be private. Queue symlinks and Windows reparse points
are rejected. On Windows, the relay does not set ACLs; use a directory restricted
to the account that runs the orchestrator.
SQLite uses WAL mode and `synchronous=FULL`. Keep the database and its sidecar
files together. To move or back up a queue, stop producers and flush processes
first. Reopen the same queue after a restart; do not edit its SQLite tables to
mark events as delivered. An unknown queue schema is rejected.
The queue uses these limits:
| Resource | Limit |
| --- | --- |
| Pending events | 50,000 |
| Pending body bytes | 64 MiB |
| Pending events plus retained receipts | 200,000 |
| Retained session identities | 50,000 |
| One event body | 256 KiB |
| One input file | 32 MiB and 10,000 events |
| One HTTP batch | 8 MiB and 256 events |
| Acknowledgement body | 64 KiB |
Reaching a limit rejects new input instead of dropping older events. Receipts
and unused session records can expire after the 30-day retry window. Pending
events are retained. Review expired entries and producer failures before
choosing how to archive a queue; the CLI has no command that discards them
automatically. These are logical record limits, not a fixed SQLite file size.
## Validation
The package needs separate checks because it is outside the root workspace:
```sh
cargo fmt --check --manifest-path companions/ai-memory-relay/Cargo.toml
cargo clippy --locked --manifest-path companions/ai-memory-relay/Cargo.toml --all-targets -- -D warnings
cargo test --locked --manifest-path companions/ai-memory-relay/Cargo.toml
```
Run the integration test against locally built binaries:
```sh
cargo build --workspace
cargo build --locked --manifest-path companions/ai-memory-relay/Cargo.toml
uv run --no-project python tests/e2e/external_relay_smoke.py \
--ai-memory-bin "$PWD/target/debug/ai-memory" \
--relay-bin "$PWD/companions/ai-memory-relay/target/debug/ai-memory-relay"
```
Adjust the binary paths if `CARGO_TARGET_DIR` is set; Windows executables end in
`.exe`. The test starts an isolated real ai-memory server and invokes its native
hooks. It checks persistent observations for offline recovery, retries after a
lost response, distinct event identities and 15 concurrent sessions. It also
checks handoff delivery with native capture suppressed, rejected authentication
and acknowledged session collisions. Temporary data and logs are retained, and
the test prints their directory.
companions/ai-memory-relay/src/ack.rs
+98
View File
@@ -0,0 +1,98 @@
//! Acknowledgement parsing for `POST /hook/batch`.
//!
//! Validate the entire response before releasing any queued event. An
//! inconsistent acknowledgement leaves the whole batch pending for retry.
use serde::Deserialize;
/// The server's ack. Unknown fields are accepted (the server may add some);
/// a missing `accepted` is a malformed ack and preserves the batch.
#[derive(Debug, Clone, Deserialize)]
pub struct BatchAck {
/// Contiguous leading prefix committed, oldest-first.
pub accepted: usize,
/// Non-contiguous committed indexes, when per-source rate limiting skipped items.
#[serde(default)]
pub accepted_indices: Option<Vec<usize>>,
/// Item that failed processing after earlier skips.
#[serde(default)]
pub failed_index: Option<usize>,
}
/// Why an ack was refused. The batch stays pending in every case.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AckRejected(pub String);
impl std::fmt::Display for AckRejected {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.0)
}
}
/// Validate an ack against the batch it answers, returning the indexes that may
/// be released.
///
/// `accepted` is the contiguous leading prefix *even when* `accepted_indices` is
/// present, so the two must agree. Duplicated or out-of-range indexes, a
/// disagreeing `accepted`, a `failed_index` outside the batch, and a
/// `failed_index` that also claims to be accepted are all refusals.
pub fn validate(batch_len: usize, ack: &BatchAck) -> Result<Vec<usize>, AckRejected> {
let reject = |detail: String| Err(AckRejected(detail));
let accepted: Vec<usize> = match &ack.accepted_indices {
Some(indices) => {
if let Some(bad) = indices.iter().find(|idx| **idx >= batch_len) {
return reject(format!(
"accepted_indices contains {bad}, outside a {batch_len}-item batch"
));
}
if indices.windows(2).any(|pair| pair[0] >= pair[1]) {
return reject(
"accepted_indices must be strictly ascending and duplicate-free".into(),
);
}
let prefix = indices
.iter()
.enumerate()
.take_while(|(pos, idx)| pos == *idx)
.count();
if prefix != ack.accepted {
return reject(format!(
"accepted={} disagrees with the {prefix}-item contiguous prefix of accepted_indices",
ack.accepted
));
}
indices.clone()
}
None => {
if ack.accepted > batch_len {
return reject(format!(
"accepted={} exceeds the {batch_len} items sent",
ack.accepted
));
}
(0..ack.accepted).collect()
}
};
if let Some(failed) = ack.failed_index {
if failed >= batch_len {
return reject(format!(
"failed_index={failed} is outside a {batch_len}-item batch"
));
}
if accepted.contains(&failed) {
return reject(format!(
"failed_index={failed} is also reported as accepted"
));
}
// The server fails fast: it stops at `failed_index` and processes
// nothing after it. An ack claiming a later item committed describes a
// run that cannot have happened, so the batch is preserved whole.
if let Some(after) = accepted.iter().find(|idx| **idx > failed) {
return reject(format!(
"accepted index {after} comes after failed_index={failed}, which the server \
never processes past"
));
}
}
Ok(accepted)
}
companions/ai-memory-relay/src/fsguard.rs
+220
View File
@@ -0,0 +1,220 @@
//! Filesystem guards for the queue directory.
//!
//! Event bodies reach this directory before server sanitization. Unix paths
//! must be private; existing permissions are checked without changing them.
//!
//! Windows: only symlink/reparse-point rejection is implemented. No ACL
//! hardening is performed, so a queue directory there is exactly as private as
//! the operator made it.
use std::path::{Component, Path, PathBuf};
use anyhow::{Context, Result, bail};
/// SQLite sidecars that must be checked alongside the database itself.
pub const SIDECARS: &[&str] = &[
"relay.sqlite-wal",
"relay.sqlite-shm",
"relay.sqlite-journal",
];
/// The queue database file name.
pub const DB_FILE: &str = "relay.sqlite";
/// Reject a path that is a symlink or a Windows reparse point.
///
/// A missing path is fine: it is about to be created inside a directory that
/// was itself checked.
pub fn reject_symlink(path: &Path) -> Result<()> {
let meta = match std::fs::symlink_metadata(path) {
Ok(meta) => meta,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(()),
Err(e) => return Err(e).with_context(|| format!("inspect {}", path.display())),
};
if meta.file_type().is_symlink() {
bail!(
"{} is a symlink; the relay refuses to follow one into a queue directory",
path.display()
);
}
#[cfg(windows)]
{
use std::os::windows::fs::MetadataExt;
const FILE_ATTRIBUTE_REPARSE_POINT: u32 = 0x400;
if meta.file_attributes() & FILE_ATTRIBUTE_REPARSE_POINT != 0 {
bail!(
"{} is a reparse point; the relay refuses to follow one into a queue directory",
path.display()
);
}
}
Ok(())
}
/// Check one queue file: not a symlink/reparse point, a regular file, not
/// hardlinked elsewhere, and owner-only on Unix. A missing file passes.
pub fn check_queue_file(path: &Path) -> Result<()> {
reject_symlink(path)?;
let meta = match std::fs::symlink_metadata(path) {
Ok(meta) => meta,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(()),
Err(e) => return Err(e).with_context(|| format!("inspect {}", path.display())),
};
if !meta.is_file() {
bail!(
"{} exists but is not a regular file; refusing to use it as queue state",
path.display()
);
}
#[cfg(unix)]
{
use std::os::unix::fs::MetadataExt;
use std::os::unix::fs::PermissionsExt;
if meta.nlink() > 1 {
bail!(
"{} has {} hard links; refusing to write queue state through a shared inode",
path.display(),
meta.nlink()
);
}
let mode = meta.permissions().mode();
if mode & 0o077 != 0 {
bail!(
"{} is mode {:o}; queue state must be owner-only (chmod 600) \
(existing permissions are unchanged)",
path.display(),
mode & 0o7777
);
}
}
Ok(())
}
/// Prepare `--queue-dir`.
///
/// A directory the relay creates is created owner-only. An existing directory
/// must already be private and must be either empty or an existing relay queue;
/// anything else is refused untouched, so pointing the relay at a shared
/// directory can never re-permission it.
pub fn prepare_queue_dir(dir: &Path) -> Result<PathBuf> {
if dir.components().any(|c| c == Component::ParentDir) {
bail!(
"--queue-dir must not contain a `..` component: {}",
dir.display()
);
}
reject_symlink(dir)?;
match std::fs::symlink_metadata(dir) {
Ok(meta) if !meta.is_dir() => bail!("--queue-dir {} is not a directory", dir.display()),
Ok(_) => check_existing_dir(dir)?,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => create_private_dir(dir)?,
Err(e) => return Err(e).with_context(|| format!("inspect {}", dir.display())),
}
for name in std::iter::once(DB_FILE).chain(SIDECARS.iter().copied()) {
check_queue_file(&dir.join(name))?;
}
std::fs::canonicalize(dir).with_context(|| format!("resolve {}", dir.display()))
}
/// An existing directory must be private and dedicated to this queue.
fn check_existing_dir(dir: &Path) -> Result<()> {
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let mode = std::fs::metadata(dir)?.permissions().mode();
if mode & 0o077 != 0 {
bail!(
"--queue-dir {} is mode {:o}; it must be owner-only (chmod 700) before the relay \
will store event bodies in it. The relay does not re-permission a directory it \
did not create",
dir.display(),
mode & 0o7777
);
}
}
let mut foreign = Vec::new();
for entry in std::fs::read_dir(dir).with_context(|| format!("read {}", dir.display()))? {
let name = entry?.file_name();
let name = name.to_string_lossy().to_string();
let known = name == DB_FILE || name == "flush.lock" || SIDECARS.contains(&name.as_str());
if !known {
foreign.push(name);
}
}
if !foreign.is_empty() {
foreign.sort();
foreign.truncate(5);
bail!(
"--queue-dir {} already holds unrelated files ({}); point the relay at an empty or \
existing queue directory instead of sharing one",
dir.display(),
foreign.join(", ")
);
}
Ok(())
}
#[cfg(unix)]
fn create_private_dir(dir: &Path) -> Result<()> {
use std::os::unix::fs::DirBuilderExt;
if let Some(parent) = dir.parent()
&& !parent.as_os_str().is_empty()
&& !parent.exists()
{
bail!(
"parent of --queue-dir {} does not exist; create it deliberately first",
dir.display()
);
}
std::fs::DirBuilder::new()
.mode(0o700)
.create(dir)
.with_context(|| format!("create {}", dir.display()))
}
#[cfg(not(unix))]
fn create_private_dir(dir: &Path) -> Result<()> {
if let Some(parent) = dir.parent()
&& !parent.as_os_str().is_empty()
&& !parent.exists()
{
bail!(
"parent of --queue-dir {} does not exist; create it deliberately first",
dir.display()
);
}
std::fs::create_dir(dir).with_context(|| format!("create {}", dir.display()))
}
/// Create a queue file owner-only, or leave an existing one alone.
///
/// SQLite copies the database mode onto its sidecars. Set 0600 at creation so
/// WAL and SHM files are private from their first write.
#[cfg(unix)]
pub fn create_private_file(path: &Path) -> Result<()> {
use std::os::unix::fs::OpenOptionsExt;
match std::fs::OpenOptions::new()
.create_new(true)
.write(true)
.mode(0o600)
.open(path)
{
Ok(_) => Ok(()),
// Another process won the race and made it; its own guard applies.
Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => Ok(()),
Err(e) => Err(e).with_context(|| format!("create {}", path.display())),
}
}
/// No ACL hardening on non-Unix targets; documented in the README.
#[cfg(not(unix))]
pub fn create_private_file(path: &Path) -> Result<()> {
match std::fs::OpenOptions::new()
.create_new(true)
.write(true)
.open(path)
{
Ok(_) => Ok(()),
Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => Ok(()),
Err(e) => Err(e).with_context(|| format!("create {}", path.display())),
}
}
companions/ai-memory-relay/src/http.rs
+200
View File
@@ -0,0 +1,200 @@
//! The only network surface: `POST <server>/hook/batch`.
//!
//! Per-item URLs are always constructed here from the bound destination and the
//! queue's own fields. An input file can never steer a request anywhere.
use std::io::Read;
use std::time::Duration;
use anyhow::{Context, Result, bail};
use url::Url;
use crate::queue::{Binding, PendingItem};
/// Bearer material. Never stored on disk, never printed, never in a Debug line.
pub struct Token(String);
impl std::fmt::Debug for Token {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("Token(<redacted>)")
}
}
impl Token {
/// Read `AI_MEMORY_AUTH_TOKEN` from the environment of *this* flush.
///
/// The value is never echoed, not even when it is rejected.
pub fn from_env() -> Result<Option<Self>> {
let Ok(raw) = std::env::var("AI_MEMORY_AUTH_TOKEN") else {
return Ok(None);
};
let trimmed = raw.trim();
if trimmed.is_empty() {
return Ok(None);
}
if trimmed.chars().any(|c| c.is_control()) {
bail!("AI_MEMORY_AUTH_TOKEN contains control characters; refusing to build a request");
}
Ok(Some(Self(trimmed.to_owned())))
}
fn expose(&self) -> &str {
&self.0
}
}
/// Normalize and harden `--server-url`.
///
/// http/https only, a host, no userinfo, no query, no fragment. No host
/// allowlist: a remote HTTPS server is a legitimate deployment.
pub fn normalize_server_url(raw: &str) -> Result<String> {
let url = Url::parse(raw).context("invalid --server-url")?;
if !matches!(url.scheme(), "http" | "https") {
bail!("--server-url must be http or https, got {:?}", url.scheme());
}
if !url.username().is_empty() || url.password().is_some() {
bail!("--server-url must not carry userinfo; pass credentials in AI_MEMORY_AUTH_TOKEN");
}
if url.query().is_some() {
bail!("--server-url must not carry a query string");
}
if url.fragment().is_some() {
bail!("--server-url must not carry a fragment");
}
let host = url.host_str().context("--server-url needs a host")?;
let authority = match url.port() {
Some(port) => format!("{host}:{port}"),
None => host.to_owned(),
};
let path = url.path().trim_end_matches('/');
Ok(format!("{}://{}{}", url.scheme(), authority, path))
}
/// The `{url, body}` pair the batch endpoint expects.
#[derive(Debug, serde::Serialize)]
pub struct BatchItem {
pub url: String,
pub body: serde_json::Value,
}
/// Build one item's hook URL from the binding plus the queued event.
///
/// `extension` carries the producer namespace, `source_event` repeats the
/// canonical event explicitly (both are required to preserve provenance), and
/// `ingest_key` is the stable key persisted with the event at enqueue time.
pub fn item_url(base: &str, binding: &Binding, item: &PendingItem) -> String {
let query = url::form_urlencoded::Serializer::new(String::new())
.append_pair("event", &item.event)
.append_pair("agent", &item.agent)
.append_pair("workspace", &binding.workspace)
.append_pair("project", &binding.project)
.append_pair("extension", &binding.producer)
.append_pair("source_event", &item.event)
.append_pair("ingest_key", &item.ingest_key)
.finish();
format!("{base}/hook?{query}")
}
/// Largest ack body the relay will read. A 256-index ack is under 2 KiB; 64 KiB
/// leaves room for future fields while refusing to stream an unbounded response
/// from a server that is not behaving like ai-memory.
pub const MAX_ACK_BYTES: usize = 64 * 1024;
/// One completed HTTP exchange. The body is bounded and never logged.
pub struct Delivery {
pub status: u16,
pub body: Vec<u8>,
}
impl std::fmt::Debug for Delivery {
// Length only: an ack body is server output, and output never becomes a log
// line here by accident.
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Delivery")
.field("status", &self.status)
.field("body_bytes", &self.body.len())
.finish()
}
}
/// Blocking batch sender.
pub struct Sender {
client: reqwest::blocking::Client,
base: String,
token: Option<Token>,
}
impl std::fmt::Debug for Sender {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Sender")
.field("base", &self.base)
.field("token", &self.token)
.finish()
}
}
impl Sender {
/// Build the client. Redirects are disabled: a 3xx from the ingestion
/// endpoint must never resend the bearer header, or the batch, elsewhere.
/// The read timeout is per request, because a batch's budget scales with the
/// number of items it carries.
pub fn new(base: String, token: Option<Token>, connect_timeout: Duration) -> Result<Self> {
let client = reqwest::blocking::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.connect_timeout(connect_timeout)
.build()
.context("build HTTP client")?;
Ok(Self {
client,
base,
token,
})
}
/// POST one batch under an explicit timeout.
///
/// Transport failures carry a short class label only: never the request, the
/// response, or the bearer header.
pub fn post_batch(&self, items: &[BatchItem], timeout: Duration) -> Result<Delivery> {
let mut request = self
.client
.post(format!("{}/hook/batch", self.base))
.timeout(timeout)
.json(items);
if let Some(token) = &self.token {
request = request.bearer_auth(token.expose());
}
let response = request.send().map_err(|e| {
anyhow::anyhow!("hook batch transport failure: {}", transport_class(&e))
})?;
let status = response.status().as_u16();
// Bounded read: `bytes()` would accept whatever the peer sends.
let mut body = Vec::new();
let mut limited = response.take(MAX_ACK_BYTES as u64 + 1);
if limited.read_to_end(&mut body).is_err() {
body.clear();
}
if body.len() > MAX_ACK_BYTES {
bail!(
"hook batch ack exceeded {MAX_ACK_BYTES} bytes; refusing to parse it \
(the batch stays pending)"
);
}
Ok(Delivery { status, body })
}
}
/// A short, content-free label for a transport error.
fn transport_class(error: &reqwest::Error) -> &'static str {
if error.is_timeout() {
"timeout"
} else if error.is_connect() {
"connect"
} else if error.is_redirect() {
"redirect refused"
} else if error.is_request() {
"request"
} else {
"io"
}
}
companions/ai-memory-relay/src/identity.rs
+224
View File
@@ -0,0 +1,224 @@
//! Envelope validation and the stable retry identity.
//!
//! Validate input before enqueueing so malformed events cannot block a batch.
//! Body values, including `_ai_memory_capture`, are preserved for the server.
//! These checks do not sanitize locally queued content.
use serde::Deserialize;
use sha2::{Digest, Sha256};
/// `extension` accepts up to 64 ASCII token characters (docs/external-lifecycle.md).
pub const MAX_PRODUCER_LEN: usize = 64;
/// A stable adapter-side operator namespace. Never a bearer token.
pub const MAX_ACTOR_LEN: usize = 64;
/// `source_event` accepts 128 (docs/external-lifecycle.md).
pub const MAX_EVENT_LEN: usize = 128;
/// Producer-assigned event id. It only feeds the ingest-key hash, so it is
/// bounded free-form text rather than a
/// token: an orchestrator numbering events `run/123/event/2` is normal.
pub const MAX_EVENT_ID_LEN: usize = 128;
/// Wire `agent` such as `claude-code` or `codex`.
pub const MAX_AGENT_LEN: usize = 64;
/// Workspace / project names, URL-encoded into the item query.
pub const MAX_SCOPE_LEN: usize = 128;
/// Native session id, preserved exactly as the harness minted it.
pub const MAX_SESSION_ID_LEN: usize = 256;
/// Native cwd, preserved exactly as the harness reported it.
pub const MAX_CWD_LEN: usize = 4096;
/// Canonicalized body bytes accepted for one event.
pub const MAX_BODY_BYTES: usize = 256 * 1024;
/// Items in one `POST /hook/batch` (server: `MAX_HOOK_BATCH_ITEMS`).
pub const MAX_BATCH_ITEMS: usize = 256;
/// Server ingest keys expire after 30 days; nothing older may be auto-resent.
pub const RETRY_WINDOW_MS: i64 = 30 * 24 * 60 * 60 * 1000;
/// Lifecycle events that end an execution. A terminal event may only ship once
/// it is its session's oldest pending item, so it can never pass an earlier one.
pub const TERMINAL_EVENTS: &[&str] = &["session-end"];
/// One item of the `enqueue --file` input array.
///
/// `deny_unknown_fields` applies to the envelope only. `body` is an opaque
/// object: unknown keys inside it (a harness payload field, the
/// `_ai_memory_capture` protocol block) are preserved untouched.
#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct InputEvent {
/// Producer-assigned id, unique within the producer/actor/agent/session tuple.
pub event_id: String,
/// Wire agent of the live harness (`claude-code`, `codex`, ...).
pub agent: String,
/// Canonical lifecycle event name, used for both `event` and `source_event`.
pub event: String,
/// Harness payload, delivered verbatim.
pub body: serde_json::Value,
}
/// A validated event, ready to be persisted and later replayed unchanged.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ValidEvent {
pub event_id: String,
pub agent: String,
pub event: String,
pub session_id: String,
pub cwd: String,
/// Canonical JSON of the body exactly as supplied (sorted keys, no reformatting
/// of values). The bytes delivered later are these bytes.
pub body_json: String,
pub body_sha256: String,
pub ingest_key: String,
}
impl ValidEvent {
/// True when this event ends the execution for its session.
pub fn is_terminal(&self) -> bool {
TERMINAL_EVENTS.contains(&self.event.as_str())
}
}
/// `extension`, `source_event` and `ingest_key` share the server's token
/// alphabet: letters, digits, `.`, `_`, `-` and `:`.
fn is_token(value: &str, max: usize) -> bool {
!value.is_empty()
&& value.len() <= max
&& value
.bytes()
.all(|b| b.is_ascii_alphanumeric() || matches!(b, b'.' | b'_' | b'-' | b':'))
}
/// Free-form identity text (session id, cwd, scope names): bounded, non-blank,
/// and free of control characters that would corrupt a URL or a log line.
fn is_plain(value: &str, max: usize) -> bool {
!value.trim().is_empty()
&& value.chars().count() <= max
&& !value.chars().any(|c| c.is_control())
}
/// Validate the fields that identify a producer namespace.
pub fn check_producer(producer: &str) -> Result<(), String> {
is_token(producer, MAX_PRODUCER_LEN)
.then_some(())
.ok_or_else(|| format!("--producer must be 1..={MAX_PRODUCER_LEN} token characters"))
}
/// Validate the adapter-side operator namespace.
pub fn check_actor(actor: &str) -> Result<(), String> {
is_token(actor, MAX_ACTOR_LEN)
.then_some(())
.ok_or_else(|| format!("--actor must be 1..={MAX_ACTOR_LEN} token characters"))
}
/// Validate a workspace or project name.
pub fn check_scope(label: &str, value: &str) -> Result<(), String> {
is_plain(value, MAX_SCOPE_LEN)
.then_some(())
.ok_or_else(|| format!("--{label} must be 1..={MAX_SCOPE_LEN} printable characters"))
}
/// The retry identity from `docs/external-lifecycle.md`, byte-for-byte.
///
/// The tuple defines the namespace of `event_id`. Payload content does not
/// contribute to the key, so a restarted producer can derive the same identity.
pub fn ingest_key(
producer: &str,
actor: &str,
agent: &str,
session_id: &str,
source_event: &str,
event_id: &str,
) -> String {
let identity = serde_json::to_string(&[
"external-capture-v1",
producer,
actor,
agent,
session_id,
source_event,
event_id,
])
.expect("a fixed-size array of strings always serializes");
let digest = Sha256::digest(identity.as_bytes());
let mut key = String::with_capacity(64);
for byte in digest {
use std::fmt::Write as _;
let _ = write!(key, "{byte:02x}");
}
key
}
/// Validate one input item against every bound above.
///
/// Errors name the array index and the producer's own `event_id`. They never
/// quote body content: an invalid payload must not become a log line.
pub fn validate(
index: usize,
input: InputEvent,
producer: &str,
actor: &str,
) -> Result<ValidEvent, String> {
let at = |detail: &str| format!("item {index}: {detail}");
if !is_plain(&input.event_id, MAX_EVENT_ID_LEN) {
return Err(at(&format!(
"event_id must be 1..={MAX_EVENT_ID_LEN} printable characters"
)));
}
let tag = format!("item {index} (event_id {}): ", input.event_id);
let at = |detail: &str| format!("{tag}{detail}");
if !is_token(&input.agent, MAX_AGENT_LEN) {
return Err(at(&format!(
"agent must be 1..={MAX_AGENT_LEN} token characters"
)));
}
if !is_token(&input.event, MAX_EVENT_LEN) {
return Err(at(&format!(
"event must be 1..={MAX_EVENT_LEN} token characters"
)));
}
let serde_json::Value::Object(body) = &input.body else {
return Err(at("body must be a JSON object"));
};
let Some(serde_json::Value::String(session_id)) = body.get("session_id") else {
return Err(at("body.session_id must be an explicit string"));
};
if !is_plain(session_id, MAX_SESSION_ID_LEN) {
return Err(at(&format!(
"body.session_id must be 1..={MAX_SESSION_ID_LEN} printable characters"
)));
}
let Some(serde_json::Value::String(cwd)) = body.get("cwd") else {
return Err(at("body.cwd must be an explicit string"));
};
if !is_plain(cwd, MAX_CWD_LEN) {
return Err(at(&format!(
"body.cwd must be 1..={MAX_CWD_LEN} printable characters"
)));
}
let (session_id, cwd) = (session_id.clone(), cwd.clone());
let body_json = serde_json::to_string(&input.body).map_err(|_| at("body is not encodable"))?;
if body_json.len() > MAX_BODY_BYTES {
return Err(at(&format!(
"body is {} bytes, over the {MAX_BODY_BYTES}-byte limit",
body_json.len()
)));
}
let digest = Sha256::digest(body_json.as_bytes());
let body_sha256 = digest.iter().map(|b| format!("{b:02x}")).collect();
let ingest_key = ingest_key(
producer,
actor,
&input.agent,
&session_id,
&input.event,
&input.event_id,
);
Ok(ValidEvent {
event_id: input.event_id,
agent: input.agent,
event: input.event,
session_id,
cwd,
body_json,
body_sha256,
ingest_key,
})
}
companions/ai-memory-relay/src/lib.rs
+28
View File
@@ -0,0 +1,28 @@
//! Persistent queue for external lifecycle events sent through `POST /hook/batch`.
//!
//! The queue is separate from ai-memory's database and wiki. Events leave it
//! after a validated acknowledgement, which may include a server policy drop.
//! Pending events remain on disk when delivery fails or their retry window
//! expires; reaching a capacity limit rejects new input.
//!
//! Bodies retain their input values before server sanitization. Producers must
//! apply capture exclusions before enqueueing and protect the queue directory.
pub mod ack;
pub mod fsguard;
pub mod http;
pub mod identity;
pub mod queue;
pub mod relay;
/// Wall-clock milliseconds since the Unix epoch.
///
/// Used for the retry window, which is why it is clamped at 0 rather than
/// panicking on a pre-epoch clock: a nonsensical clock must not take the process
/// down mid-flush.
pub fn now_ms() -> i64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| i64::try_from(d.as_millis()).unwrap_or(i64::MAX))
.unwrap_or(0)
}
companions/ai-memory-relay/src/main.rs
+113
View File
@@ -0,0 +1,113 @@
//! Thin CLI over [`ai_memory_relay::relay`]. Parse, call, print, exit.
use std::path::PathBuf;
use ai_memory_relay::relay::{self, FlushOptions, Report};
use clap::{Parser, Subcommand};
/// Queue external lifecycle events for ai-memory.
///
/// Exit codes: 0 success, 2 failure, 3 flush ended with events pending.
#[derive(Debug, Parser)]
#[command(name = "ai-memory-relay", version, about, long_about = None)]
struct Cli {
#[command(subcommand)]
command: Command,
}
#[derive(Debug, Subcommand)]
enum Command {
/// Bind a queue directory to one destination, producer, actor and scope.
Init {
#[arg(long)]
queue_dir: PathBuf,
/// http(s) base URL of the ai-memory server. No userinfo, query or fragment.
#[arg(long)]
server_url: String,
/// Producer namespace, sent as `extension` (also your `AI_MEMORY_CAPTURE_OWNER`).
#[arg(long)]
producer: String,
/// Stable adapter-side operator namespace. Never a bearer token.
#[arg(long)]
actor: String,
#[arg(long)]
workspace: String,
#[arg(long)]
project: String,
},
/// Add events from a JSON array file. All-or-nothing.
Enqueue {
#[arg(long)]
queue_dir: PathBuf,
/// `[{"event_id","agent","event","body"}]`; body needs explicit `session_id` and `cwd`.
#[arg(long)]
file: PathBuf,
},
/// Deliver pending events through POST /hook/batch.
Flush {
#[arg(long)]
queue_dir: PathBuf,
/// Batches attempted in one flush. Finite by default: a flush never loops
/// forever, and what it does not deliver stays queued for the next run.
#[arg(long, default_value_t = 64)]
max_batches: usize,
},
/// Payload-free counters for the queue, as JSON.
Status {
#[arg(long)]
queue_dir: PathBuf,
},
}
fn main() {
let cli = Cli::parse();
let outcome = match cli.command {
Command::Init {
queue_dir,
server_url,
producer,
actor,
workspace,
project,
} => relay::init(
&queue_dir,
&server_url,
&producer,
&actor,
&workspace,
&project,
),
Command::Enqueue { queue_dir, file } => relay::enqueue(&queue_dir, &file),
Command::Flush {
queue_dir,
max_batches,
} => relay::flush(
&queue_dir,
&FlushOptions {
max_batches: max_batches.max(1),
},
),
Command::Status { queue_dir } => relay::status(&queue_dir),
};
std::process::exit(finish(outcome));
}
fn finish(outcome: anyhow::Result<Report>) -> i32 {
match outcome {
Ok(report) => {
for line in &report.summary {
println!("{line}");
}
if let Some(failure) = &report.failure {
eprintln!("error: {failure}");
}
report.exit_code()
}
Err(error) => {
// `{error:#}` prints the context chain, which is built from paths,
// counts and classes only.
eprintln!("error: {error:#}");
2
}
}
}
companions/ai-memory-relay/src/queue.rs
+727
View File
@@ -0,0 +1,727 @@
//! The companion's own durable queue: one private SQLite database per queue
//! directory. It never opens ai-memory's database or wiki.
//!
//! The binding fixes the destination, producer identity and scope. Pending
//! events keep their body values until acknowledgement. Receipts retain keys
//! and content hashes for duplicate recognition during the retry window.
//! Session records reject a native session changing agents within one queue.
use std::collections::{HashMap, HashSet};
use std::path::{Path, PathBuf};
use std::time::Duration;
use anyhow::{Context, Result, bail};
use rusqlite::{Connection, OptionalExtension, TransactionBehavior, params};
use crate::fsguard;
use crate::identity::{MAX_BODY_BYTES, RETRY_WINDOW_MS, TERMINAL_EVENTS, ValidEvent};
/// Undelivered events allowed in one queue.
pub const MAX_PENDING_ITEMS: i64 = 50_000;
/// Undelivered body bytes allowed in one queue.
pub const MAX_PENDING_BYTES: i64 = 64 * 1024 * 1024;
/// Pending events plus retained receipts. Admission reserves one receipt for
/// each new event so acknowledgement cannot exceed the limit.
pub const MAX_RECEIPTS: i64 = 200_000;
/// Session-to-agent pins retained at once, enforced the same way.
pub const MAX_SESSION_PINS: i64 = 50_000;
const SCHEMA_VERSION: &str = "1";
const IDENTITY: &str = "ai-memory-relay-queue";
const SCHEMA: &str = "
CREATE TABLE IF NOT EXISTS meta(
key TEXT PRIMARY KEY,
value TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS binding(
id INTEGER PRIMARY KEY CHECK(id = 1),
server_url TEXT NOT NULL,
producer TEXT NOT NULL,
actor TEXT NOT NULL,
workspace TEXT NOT NULL,
project TEXT NOT NULL,
created_at_ms INTEGER NOT NULL
);
CREATE TABLE IF NOT EXISTS pending(
seq INTEGER PRIMARY KEY AUTOINCREMENT,
ingest_key TEXT NOT NULL UNIQUE,
event_id TEXT NOT NULL,
agent TEXT NOT NULL,
event TEXT NOT NULL,
session_id TEXT NOT NULL,
cwd TEXT NOT NULL,
body_json TEXT NOT NULL,
body_sha256 TEXT NOT NULL,
body_bytes INTEGER NOT NULL,
first_seen_ms INTEGER NOT NULL,
first_attempt_ms INTEGER,
attempts INTEGER NOT NULL DEFAULT 0,
last_error TEXT
);
CREATE INDEX IF NOT EXISTS pending_session ON pending(agent, session_id, seq);
CREATE TABLE IF NOT EXISTS receipt(
ingest_key TEXT PRIMARY KEY,
body_sha256 TEXT NOT NULL,
first_attempt_ms INTEGER NOT NULL,
delivered_at_ms INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS receipt_first_attempt ON receipt(first_attempt_ms);
CREATE TABLE IF NOT EXISTS session_agent(
session_id TEXT PRIMARY KEY,
agent TEXT NOT NULL,
last_seen_ms INTEGER NOT NULL
);
";
/// What `init` bound this queue directory to.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Binding {
pub server_url: String,
pub producer: String,
pub actor: String,
pub workspace: String,
pub project: String,
}
/// Whether `init` created the binding or recognized an identical one.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BindOutcome {
Created,
Recognized,
}
/// Per-file result of `enqueue`.
#[derive(Debug, Default, Clone, PartialEq, Eq)]
pub struct EnqueueReport {
pub accepted: usize,
/// Same id, byte-identical body: already pending or already delivered.
pub recognized: usize,
}
/// One item selected for delivery.
#[derive(Debug, Clone)]
pub struct PendingItem {
pub seq: i64,
pub ingest_key: String,
pub event_id: String,
pub agent: String,
pub event: String,
pub session_id: String,
pub body_json: String,
pub body_bytes: i64,
pub first_seen_ms: i64,
pub first_attempt_ms: Option<i64>,
}
impl PendingItem {
/// Session coordinate: one head per `(agent, session_id)` per batch.
pub fn session(&self) -> (String, String) {
(self.agent.clone(), self.session_id.clone())
}
}
/// A batch plus what was held back while building it.
#[derive(Debug, Default, Clone)]
pub struct Batch {
pub items: Vec<PendingItem>,
/// Sessions whose head is past the retry window. The whole session is held:
/// skipping its head and sending the next event would reorder it.
pub blocked_sessions: usize,
/// Sessions skipped because the byte budget filled first.
pub budget_deferred: usize,
}
/// Counters for `status`. Deliberately payload-free.
#[derive(Debug, Default, Clone, PartialEq, Eq)]
pub struct Stats {
pub pending_items: i64,
pub pending_bytes: i64,
pub pending_sessions: i64,
pub attempted_items: i64,
pub expired_items: i64,
pub oldest_pending_age_ms: i64,
pub max_attempts: i64,
pub receipts: i64,
pub known_sessions: i64,
}
/// The companion's queue database.
pub struct Queue {
conn: Connection,
path: PathBuf,
}
impl std::fmt::Debug for Queue {
// Path only: never connection internals, never row data.
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Queue").field("path", &self.path).finish()
}
}
impl Queue {
/// Open (or create) the queue database inside an already-prepared directory.
///
/// Check an existing file's identity before schema writes. Schema and
/// metadata commit together, so an interrupted initialization can reopen
/// the empty database left by rollback.
pub fn open(dir: &Path) -> Result<Self> {
let path = dir.join(fsguard::DB_FILE);
fsguard::check_queue_file(&path)?;
for name in fsguard::SIDECARS {
fsguard::check_queue_file(&dir.join(name))?;
}
if !path.exists() {
// Own the mode before SQLite ever opens the file: the `-wal` and
// `-shm` sidecars inherit the database's permissions, so creating it
// under the ambient umask would make them world-readable.
fsguard::create_private_file(&path)?;
fsguard::check_queue_file(&path)?;
}
let mut conn = Connection::open(&path)
.with_context(|| format!("open relay queue at {}", path.display()))?;
conn.busy_timeout(Duration::from_secs(10))?;
let state = classify(&conn, &path)?;
conn.execute_batch(
"PRAGMA journal_mode=WAL;\nPRAGMA synchronous=FULL;\nPRAGMA foreign_keys=ON;",
)
.with_context(|| format!("configure relay queue at {}", path.display()))?;
if state != DbState::Ready {
let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
tx.execute_batch(SCHEMA)
.with_context(|| format!("initialize relay queue schema at {}", path.display()))?;
tx.execute(
"INSERT OR REPLACE INTO meta(key, value) VALUES('identity', ?1), ('schema_version', ?2)",
params![IDENTITY, SCHEMA_VERSION],
)?;
tx.commit()?;
}
Ok(Self { conn, path })
}
/// Path of the database, for operator-facing messages.
pub fn path(&self) -> &Path {
&self.path
}
/// Bind the queue to a destination, producer, actor and scope.
///
/// Repeated initialization must preserve the destination and identity of
/// queued events, including events whose acknowledgement was lost.
pub fn bind(&mut self, binding: &Binding, now_ms: i64) -> Result<BindOutcome> {
let tx = self
.conn
.transaction_with_behavior(TransactionBehavior::Immediate)?;
let outcome = match read_binding(&tx)? {
Some(existing) if existing == *binding => BindOutcome::Recognized,
Some(existing) => {
let mut changed = Vec::new();
let mut note = |label: &str, old: &str, new: &str| {
if old != new {
changed.push(format!("{label}: {old:?} -> {new:?}"));
}
};
note("server-url", &existing.server_url, &binding.server_url);
note("producer", &existing.producer, &binding.producer);
note("actor", &existing.actor, &binding.actor);
note("workspace", &existing.workspace, &binding.workspace);
note("project", &existing.project, &binding.project);
bail!(
"binding mismatch: this queue is already bound ({}). \
Use a separate --queue-dir for a different destination or identity",
changed.join(", ")
);
}
None => {
tx.execute(
"INSERT INTO binding(id, server_url, producer, actor, workspace, project, created_at_ms)
VALUES(1, ?1, ?2, ?3, ?4, ?5, ?6)",
params![
binding.server_url,
binding.producer,
binding.actor,
binding.workspace,
binding.project,
now_ms,
],
)?;
BindOutcome::Created
}
};
tx.commit()?;
Ok(outcome)
}
/// The binding, or a clear error telling the operator to run `init` first.
pub fn binding(&self) -> Result<Binding> {
read_binding(&self.conn)?.ok_or_else(|| {
anyhow::anyhow!("queue is not bound yet; run `ai-memory-relay init` first")
})
}
/// Persist a whole validated file, all-or-nothing.
///
/// `first_seen_ms` is stamped here for reporting. It is *not* the retry
/// clock: that starts at the first durable attempt, recorded before the
/// first HTTP request (see [`Queue::stamp_attempt`]).
pub fn enqueue(&mut self, events: &[ValidEvent], now_ms: i64) -> Result<EnqueueReport> {
let tx = self
.conn
.transaction_with_behavior(TransactionBehavior::Immediate)?;
let (mut items, mut bytes): (i64, i64) = tx.query_row(
"SELECT COUNT(*), COALESCE(SUM(body_bytes), 0) FROM pending",
[],
|r| Ok((r.get(0)?, r.get(1)?)),
)?;
// Replay protection is bounded, and the bound covers pending *and*
// receipts: every pending item becomes a receipt when it is
// acknowledged, so counting receipts alone would let the queue commit to
// more protection than the cap allows. Pruning stays reserved for
// entries past the 30-day window, so a full queue refuses new ids rather
// than forgetting a recently delivered one. A duplicate is still
// recognized while full, because recognizing it adds nothing.
let receipts: i64 = tx.query_row("SELECT COUNT(*) FROM receipt", [], |r| r.get(0))?;
let mut pins: i64 = tx.query_row("SELECT COUNT(*) FROM session_agent", [], |r| r.get(0))?;
let mut report = EnqueueReport::default();
let mut seen_keys: HashMap<&str, &str> = HashMap::new();
let mut seen_sessions: HashMap<&str, &str> = HashMap::new();
for event in events {
// A session id belongs to exactly one wire agent. Two agents on one
// native session is the producer bug the server can only answer with
// an acknowledged SessionCollision drop.
let bound: Option<String> = tx
.query_row(
"SELECT agent FROM session_agent WHERE session_id = ?1",
params![event.session_id],
|r| r.get(0),
)
.optional()?;
let bound = bound.or_else(|| {
seen_sessions
.get(event.session_id.as_str())
.map(|a| (*a).to_owned())
});
if bound.is_none() {
pins += 1;
if pins > MAX_SESSION_PINS {
bail!(
"session pin table is full: {MAX_SESSION_PINS} session(s) already pinned \
to an agent. Nothing was enqueued and no pin was forgotten; pins are \
released as they pass the 30-day retry window"
);
}
}
if let Some(bound) = bound
&& bound != event.agent
{
bail!(
"session identity conflict: session {:?} is already bound to agent {:?} in \
this queue, but item with event_id {} claims agent {:?}. Nothing was \
enqueued; a native session belongs to one harness",
event.session_id,
bound,
event.event_id,
event.agent
);
}
seen_sessions.insert(&event.session_id, &event.agent);
if let Some(previous) = seen_keys.insert(&event.ingest_key, &event.body_sha256)
&& previous != event.body_sha256
{
bail!(collision_message(
&event.ingest_key,
"another item in this file"
));
}
let known: Option<String> = tx
.query_row(
"SELECT body_sha256 FROM pending WHERE ingest_key = ?1
UNION ALL
SELECT body_sha256 FROM receipt WHERE ingest_key = ?1",
params![event.ingest_key],
|r| r.get(0),
)
.optional()?;
if let Some(known) = known {
if known != event.body_sha256 {
bail!(collision_message(
&event.ingest_key,
"an item already recorded in this queue"
));
}
report.recognized += 1;
continue;
}
let body_bytes = event.body_json.len() as i64;
items += 1;
bytes += body_bytes;
if items > MAX_PENDING_ITEMS || bytes > MAX_PENDING_BYTES {
bail!(
"queue is full: {items} items / {bytes} bytes would exceed the \
{MAX_PENDING_ITEMS}-item / {MAX_PENDING_BYTES}-byte limit. Nothing was \
enqueued and nothing was dropped; flush the queue first"
);
}
// Reserve the receipt this item will need once it is acknowledged.
let protected = receipts + items;
if protected > MAX_RECEIPTS {
bail!(
"replay protection is full: {receipts} acknowledged key(s) plus {items} \
pending would need {protected} slots, over the {MAX_RECEIPTS} limit. \
Nothing was enqueued and nothing was forgotten; slots are released as \
entries pass the 30-day retry window, or start a fresh --queue-dir"
);
}
tx.execute(
"INSERT INTO pending(
ingest_key, event_id, agent, event, session_id, cwd,
body_json, body_sha256, body_bytes, first_seen_ms)
VALUES(?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)",
params![
event.ingest_key,
event.event_id,
event.agent,
event.event,
event.session_id,
event.cwd,
event.body_json,
event.body_sha256,
body_bytes,
now_ms,
],
)?;
report.accepted += 1;
}
for (session_id, agent) in seen_sessions {
tx.execute(
"INSERT INTO session_agent(session_id, agent, last_seen_ms) VALUES(?1, ?2, ?3)
ON CONFLICT(session_id) DO UPDATE SET last_seen_ms = excluded.last_seen_ms",
params![session_id, agent, now_ms],
)?;
}
tx.commit()?;
Ok(report)
}
/// Pick at most one pending head per `(agent, session_id)`, oldest first,
/// within both an item count and a wire-byte budget.
///
/// One head per session per batch is what makes a partial acknowledgement
/// safe: two events of one session are never in flight together, so no ack
/// subset can commit event 2 while event 1 is still pending, and a terminal
/// event can never overtake an earlier one. `deferred` holds sessions this
/// flush already stopped on; they are skipped without touching their order.
pub fn select_batch(
&self,
limit: usize,
byte_budget: usize,
now_ms: i64,
deferred: &HashSet<(String, String)>,
cost_of: impl Fn(&PendingItem) -> usize,
) -> Result<Batch> {
let limit = limit.clamp(1, crate::identity::MAX_BATCH_ITEMS);
let mut stmt = self.conn.prepare(
"SELECT seq, ingest_key, event_id, agent, event, session_id, body_json,
body_bytes, first_seen_ms, first_attempt_ms
FROM pending ORDER BY seq ASC",
)?;
let mut rows = stmt.query([])?;
let mut seen: HashSet<(String, String)> = HashSet::new();
let mut batch = Batch::default();
let mut spent = 0usize;
while let Some(row) = rows.next()? {
if batch.items.len() >= limit {
break;
}
let item = PendingItem {
seq: row.get(0)?,
ingest_key: row.get(1)?,
event_id: row.get(2)?,
agent: row.get(3)?,
event: row.get(4)?,
session_id: row.get(5)?,
body_json: row.get(6)?,
body_bytes: row.get(7)?,
first_seen_ms: row.get(8)?,
first_attempt_ms: row.get(9)?,
};
let session = item.session();
if !seen.insert(session.clone()) {
continue;
}
if deferred.contains(&session) {
continue;
}
if is_expired(item.first_attempt_ms, now_ms) {
batch.blocked_sessions += 1;
continue;
}
let cost = cost_of(&item);
if !batch.items.is_empty() && spent + cost > byte_budget {
batch.budget_deferred += 1;
continue;
}
debug_assert!(
!TERMINAL_EVENTS.contains(&item.event.as_str()) || self.session_head(&item)?,
"a terminal event was selected while its session had an older pending item"
);
spent += cost;
batch.items.push(item);
}
Ok(batch)
}
fn session_head(&self, item: &PendingItem) -> Result<bool> {
let older: i64 = self.conn.query_row(
"SELECT COUNT(*) FROM pending WHERE agent = ?1 AND session_id = ?2 AND seq < ?3",
params![item.agent, item.session_id, item.seq],
|r| r.get(0),
)?;
Ok(older == 0)
}
/// Record the durable start of the retry window, **before** the request.
///
/// `COALESCE` is the whole point: a restart, a lost ack, or a later retry
/// never moves the clock forward, so the 30-day guard measures from the
/// first time this event was actually put on the wire.
pub fn stamp_attempt(&mut self, keys: &[String], now_ms: i64) -> Result<()> {
let tx = self
.conn
.transaction_with_behavior(TransactionBehavior::Immediate)?;
for key in keys {
tx.execute(
"UPDATE pending
SET first_attempt_ms = COALESCE(first_attempt_ms, ?2),
attempts = attempts + 1
WHERE ingest_key = ?1",
params![key, now_ms],
)?;
}
tx.commit()?;
Ok(())
}
/// Move acknowledged items to compact receipts, in one transaction.
///
/// An acknowledgement means delivered *or* deliberately dropped by server
/// policy (a capture-protocol drop, a subagent drop, a session-collision
/// drop). It never promises an observation was written.
pub fn confirm(&mut self, keys: &[String], now_ms: i64) -> Result<usize> {
if keys.is_empty() {
return Ok(0);
}
let tx = self
.conn
.transaction_with_behavior(TransactionBehavior::Immediate)?;
let mut moved = 0;
for key in keys {
let row: Option<(String, Option<i64>, i64)> = tx
.query_row(
"SELECT body_sha256, first_attempt_ms, first_seen_ms
FROM pending WHERE ingest_key = ?1",
params![key],
|r| Ok((r.get(0)?, r.get(1)?, r.get(2)?)),
)
.optional()?;
let Some((body_sha256, first_attempt_ms, first_seen_ms)) = row else {
continue;
};
tx.execute(
"INSERT INTO receipt(ingest_key, body_sha256, first_attempt_ms, delivered_at_ms)
VALUES(?1, ?2, ?3, ?4)
ON CONFLICT(ingest_key) DO UPDATE SET delivered_at_ms = excluded.delivered_at_ms",
params![
key,
body_sha256,
first_attempt_ms.unwrap_or(first_seen_ms),
now_ms
],
)?;
tx.execute("DELETE FROM pending WHERE ingest_key = ?1", params![key])?;
moved += 1;
}
tx.commit()?;
Ok(moved)
}
/// Charge one item with a delivery failure. The class is a short label,
/// never a server body and never payload.
pub fn record_failure(&mut self, key: &str, class: &str) -> Result<()> {
let class: String = class.chars().filter(|c| !c.is_control()).take(80).collect();
self.conn.execute(
"UPDATE pending SET last_error = ?2 WHERE ingest_key = ?1",
params![key, class],
)?;
Ok(())
}
/// Release expired receipts and session records without pending events.
pub fn prune(&mut self, now_ms: i64) -> Result<usize> {
let cutoff = now_ms.saturating_sub(RETRY_WINDOW_MS);
let tx = self
.conn
.transaction_with_behavior(TransactionBehavior::Immediate)?;
let receipts = tx.execute(
"DELETE FROM receipt WHERE first_attempt_ms < ?1",
params![cutoff],
)?;
let sessions = tx.execute(
"DELETE FROM session_agent WHERE last_seen_ms < ?1
AND session_id NOT IN (SELECT session_id FROM pending)",
params![cutoff],
)?;
tx.commit()?;
Ok(receipts + sessions)
}
/// Payload-free counters.
pub fn stats(&self, now_ms: i64) -> Result<Stats> {
let cutoff = now_ms.saturating_sub(RETRY_WINDOW_MS);
let (pending_items, pending_bytes, oldest, max_attempts, attempted_items) =
self.conn.query_row(
"SELECT COUNT(*), COALESCE(SUM(body_bytes), 0), COALESCE(MIN(first_seen_ms), 0),
COALESCE(MAX(attempts), 0), COUNT(first_attempt_ms)
FROM pending",
[],
|r| {
Ok((
r.get::<_, i64>(0)?,
r.get::<_, i64>(1)?,
r.get::<_, i64>(2)?,
r.get::<_, i64>(3)?,
r.get::<_, i64>(4)?,
))
},
)?;
let pending_sessions: i64 = self.conn.query_row(
"SELECT COUNT(*) FROM (SELECT DISTINCT agent, session_id FROM pending)",
[],
|r| r.get(0),
)?;
let expired_items: i64 = self.conn.query_row(
"SELECT COUNT(*) FROM pending WHERE first_attempt_ms IS NOT NULL
AND first_attempt_ms <= ?1",
params![cutoff],
|r| r.get(0),
)?;
let receipts: i64 = self
.conn
.query_row("SELECT COUNT(*) FROM receipt", [], |r| r.get(0))?;
let known_sessions: i64 =
self.conn
.query_row("SELECT COUNT(*) FROM session_agent", [], |r| r.get(0))?;
Ok(Stats {
pending_items,
pending_bytes,
pending_sessions,
attempted_items,
expired_items,
oldest_pending_age_ms: if pending_items == 0 {
0
} else {
now_ms.saturating_sub(oldest)
},
max_attempts,
receipts,
known_sessions,
})
}
}
/// Past the retry window an event may not be auto-resent.
///
/// Refuse the boundary itself. This comparison depends on an accurate host
/// clock; a backwards clock change can extend retries past server key expiry.
pub fn is_expired(first_attempt_ms: Option<i64>, now_ms: i64) -> bool {
match first_attempt_ms {
Some(started) => now_ms.saturating_sub(started) >= RETRY_WINDOW_MS,
None => false,
}
}
fn collision_message(key: &str, other: &str) -> String {
format!(
"identity collision: ingest key {key} was already used by {other} with different content. \
The same producer event id must always carry the same body; nothing was enqueued. \
Body content is deliberately not shown. Per-body limit: {MAX_BODY_BYTES} bytes"
)
}
/// What an existing file turned out to be.
///
/// Schema and identity commit together. Recovery therefore accepts an empty
/// database or an identified queue; generic table names do not establish ownership.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum DbState {
/// A SQLite file with no user objects: empty, or rolled back to empty.
Fresh,
/// A relay queue of this exact schema version.
Ready,
}
/// Classify an existing file before touching it. Never mutates.
fn classify(conn: &Connection, path: &Path) -> Result<DbState> {
let refuse = |detail: String| -> anyhow::Error {
anyhow::anyhow!(
"{} is not a usable relay queue ({detail}); nothing was modified. \
Point --queue-dir at an empty directory or restore the original queue",
path.display()
)
};
// Tables, views, triggers and standalone indexes all count: any user object
// means the file belongs to something, and that something is not this relay
// unless it also carries the identity row.
let objects: Result<i64, rusqlite::Error> = conn.query_row(
"SELECT COUNT(*) FROM sqlite_master WHERE name NOT LIKE 'sqlite_%'",
[],
|r| r.get(0),
);
// A file that is not a SQLite database fails right here, before any write.
let objects = objects.map_err(|e| refuse(format!("cannot read its object list: {e}")))?;
if objects == 0 {
return Ok(DbState::Fresh);
}
let rows: Result<Vec<(String, String)>, rusqlite::Error> = (|| {
let mut stmt =
conn.prepare("SELECT key, value FROM meta WHERE key IN ('identity','schema_version')")?;
let mapped = stmt.query_map([], |r| Ok((r.get(0)?, r.get(1)?)))?;
mapped.collect()
})();
let meta: HashMap<String, String> = rows
.map_err(|e| refuse(format!("it carries no readable relay metadata: {e}")))?
.into_iter()
.collect();
match (
meta.get("identity").map(String::as_str),
meta.get("schema_version").map(String::as_str),
) {
(Some(IDENTITY), Some(SCHEMA_VERSION)) => Ok(DbState::Ready),
(Some(IDENTITY), Some(other)) => Err(refuse(format!(
"schema version {other} but this build speaks {SCHEMA_VERSION}"
))),
(Some(other), _) => Err(refuse(format!("identity is {other:?}"))),
_ => Err(refuse(
"it holds user objects but no relay identity row".into(),
)),
}
}
fn read_binding(conn: &Connection) -> Result<Option<Binding>> {
Ok(conn
.query_row(
"SELECT server_url, producer, actor, workspace, project FROM binding WHERE id = 1",
[],
|r| {
Ok(Binding {
server_url: r.get(0)?,
producer: r.get(1)?,
actor: r.get(2)?,
workspace: r.get(3)?,
project: r.get(4)?,
})
},
)
.optional()?)
}
companions/ai-memory-relay/src/relay.rs
+556
View File
@@ -0,0 +1,556 @@
//! CLI command operations and reports.
//!
//! Reports include queue metadata and counts. Event bodies and bearer tokens
//! are kept out of diagnostics.
use std::collections::HashSet;
use std::io::Read;
use std::path::Path;
use std::time::{Duration, Instant};
use anyhow::{Context, Result, bail};
use crate::ack::{self, BatchAck};
use crate::fsguard;
use crate::http::{self, BatchItem, Sender, Token};
use crate::identity::{self, InputEvent, MAX_BATCH_ITEMS, ValidEvent};
use crate::queue::{self, Batch, BindOutcome, Binding, PendingItem, Queue};
/// Largest `enqueue --file` input accepted, before parsing.
pub const MAX_INPUT_BYTES: u64 = 32 * 1024 * 1024;
/// Most events accepted from one input file.
pub const MAX_INPUT_ITEMS: usize = 10_000;
/// Wire-byte budget for one batch. The server's `/hook` body limit is 10 MiB
/// (`serve.rs`); core's own spool drain uses 8 MiB, and so does this.
pub const MAX_BATCH_BYTES: usize = 8 * 1024 * 1024;
/// JSON framing charged per item on top of its URL and body: `{"url":"","body":},`.
const ITEM_FRAMING_BYTES: usize = 32;
/// Transport retries per batch, on top of the first attempt.
const TRANSPORT_RETRIES: usize = 2;
/// Request timeout budget per item; a batch gets this times its item count.
const PER_EVENT_TIMEOUT: Duration = Duration::from_secs(2);
/// Absolute ceiling for one batch request, however many items it carries.
const MAX_BATCH_TIMEOUT: Duration = Duration::from_secs(120);
/// Wall-clock ceiling for one whole flush.
const TOTAL_BUDGET: Duration = Duration::from_secs(900);
/// Ceiling for one backoff sleep between transport retries.
const MAX_BACKOFF: Duration = Duration::from_secs(2);
/// TCP connect timeout.
const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
/// The one knob a `flush` exposes. Everything else is a constant above, so the
/// public surface cannot be tuned into an unbounded run.
#[derive(Debug, Clone)]
pub struct FlushOptions {
/// Batches attempted in one flush. Finite by construction.
pub max_batches: usize,
}
impl Default for FlushOptions {
fn default() -> Self {
Self { max_batches: 64 }
}
}
/// What a command wants the process to exit with.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Report {
pub summary: Vec<String>,
/// Work remains: unacknowledged items are still queued.
pub pending: bool,
/// Something went wrong. The queue is intact either way.
pub failure: Option<String>,
}
impl Report {
fn ok(summary: Vec<String>) -> Self {
Self {
summary,
pending: false,
failure: None,
}
}
/// 0 = drained, 2 = failure, 3 = nothing lost but work remains.
pub fn exit_code(&self) -> i32 {
if self.failure.is_some() {
2
} else if self.pending {
3
} else {
0
}
}
}
/// `init`: bind a queue directory to one destination, producer, actor and scope.
pub fn init(
dir: &Path,
server_url: &str,
producer: &str,
actor: &str,
workspace: &str,
project: &str,
) -> Result<Report> {
identity::check_producer(producer).map_err(anyhow::Error::msg)?;
identity::check_actor(actor).map_err(anyhow::Error::msg)?;
identity::check_scope("workspace", workspace).map_err(anyhow::Error::msg)?;
identity::check_scope("project", project).map_err(anyhow::Error::msg)?;
let binding = Binding {
server_url: http::normalize_server_url(server_url)?,
producer: producer.to_owned(),
actor: actor.to_owned(),
workspace: workspace.to_owned(),
project: project.to_owned(),
};
let dir = fsguard::prepare_queue_dir(dir)?;
let mut queue = Queue::open(&dir)?;
let outcome = queue.bind(&binding, crate::now_ms())?;
let verb = match outcome {
BindOutcome::Created => "bound",
BindOutcome::Recognized => "already bound (identical)",
};
Ok(Report::ok(vec![
format!("queue {verb}: {}", queue.path().display()),
format!("server: {}", binding.server_url),
format!("producer: {}", binding.producer),
format!("actor: {}", binding.actor),
format!("scope: {}/{}", binding.workspace, binding.project),
]))
}
/// `enqueue`: validate a whole file, then persist it all-or-nothing.
pub fn enqueue(dir: &Path, file: &Path) -> Result<Report> {
let dir = fsguard::prepare_queue_dir(dir)?;
let mut queue = Queue::open(&dir)?;
let binding = queue.binding()?;
let events = read_input(file, &binding)?;
let report = queue.enqueue(&events, crate::now_ms())?;
let stats = queue.stats(crate::now_ms())?;
Ok(Report::ok(vec![
format!(
"enqueued {} event(s), recognized {} duplicate(s)",
report.accepted, report.recognized
),
format!(
"pending now: {} event(s) across {} session(s), {} byte(s)",
stats.pending_items, stats.pending_sessions, stats.pending_bytes
),
]))
}
/// Read and validate an input file without ever quoting its content back.
fn read_input(file: &Path, binding: &Binding) -> Result<Vec<ValidEvent>> {
let meta =
std::fs::metadata(file).with_context(|| format!("read --file {}", file.display()))?;
if meta.len() > MAX_INPUT_BYTES {
bail!(
"--file {} is {} bytes, over the {MAX_INPUT_BYTES}-byte input limit",
file.display(),
meta.len()
);
}
let mut raw = Vec::new();
std::fs::File::open(file)
.with_context(|| format!("open --file {}", file.display()))?
.take(MAX_INPUT_BYTES + 1)
.read_to_end(&mut raw)
.with_context(|| format!("read --file {}", file.display()))?;
// The metadata check above is an early out; this is the one that counts,
// because the file can grow between `stat` and `read`.
if raw.len() as u64 > MAX_INPUT_BYTES {
bail!(
"--file {} is over the {MAX_INPUT_BYTES}-byte input limit",
file.display()
);
}
let input: Vec<InputEvent> = serde_json::from_slice(&raw).map_err(|e| {
// Category plus position only. `Display` on a serde error can quote the
// offending value ("invalid type: string \"...\""), and an event body is
// exactly the thing that must not reach a terminal or a log.
anyhow::anyhow!(
"--file {} is not a valid event array ({} error at line {}, column {}): expected \
[{{\"event_id\",\"agent\",\"event\",\"body\"}}, ...]",
file.display(),
match e.classify() {
serde_json::error::Category::Io => "io",
serde_json::error::Category::Syntax => "syntax",
serde_json::error::Category::Data => "schema",
serde_json::error::Category::Eof => "truncated input",
},
e.line(),
e.column()
)
})?;
if input.is_empty() {
bail!("--file {} contains no events", file.display());
}
if input.len() > MAX_INPUT_ITEMS {
bail!(
"--file {} carries {} events, over the {MAX_INPUT_ITEMS}-event input limit",
file.display(),
input.len()
);
}
input
.into_iter()
.enumerate()
.map(|(index, event)| {
identity::validate(index, event, &binding.producer, &binding.actor)
.map_err(anyhow::Error::msg)
})
.collect()
}
/// `status`: payload-free counters for the queue, as one stable JSON object.
///
/// Machine-readable on purpose: an operator script (and the end-to-end smoke
/// test) reads `pending_items` from it. Every field here is a count, a name the
/// operator chose, or a path. No event body, no cwd, no token.
pub fn status(dir: &Path) -> Result<Report> {
let dir = fsguard::prepare_queue_dir(dir)?;
let queue = Queue::open(&dir)?;
let binding = queue.binding()?;
let stats = queue.stats(crate::now_ms())?;
let document = serde_json::json!({
"queue": queue.path().display().to_string(),
"server_url": binding.server_url,
"producer": binding.producer,
"actor": binding.actor,
"workspace": binding.workspace,
"project": binding.project,
"pending_items": stats.pending_items,
"pending_bytes": stats.pending_bytes,
"pending_sessions": stats.pending_sessions,
"attempted_items": stats.attempted_items,
"expired_items": stats.expired_items,
"oldest_pending_age_ms": stats.oldest_pending_age_ms,
"max_attempts": stats.max_attempts,
"receipts": stats.receipts,
"known_sessions": stats.known_sessions,
"limits": {
"pending_items": queue::MAX_PENDING_ITEMS,
"pending_bytes": queue::MAX_PENDING_BYTES,
"receipts": queue::MAX_RECEIPTS,
"known_sessions": queue::MAX_SESSION_PINS,
"body_bytes": crate::identity::MAX_BODY_BYTES,
"batch_items": MAX_BATCH_ITEMS,
"batch_bytes": MAX_BATCH_BYTES,
},
"retry_window_ms": crate::identity::RETRY_WINDOW_MS,
"ack_means": "delivered or dropped by server policy, not that an observation was written",
});
Ok(Report::ok(vec![serde_json::to_string_pretty(&document)?]))
}
/// `flush`: deliver pending events, oldest first, one head per session per batch.
pub fn flush(dir: &Path, options: &FlushOptions) -> Result<Report> {
let dir = fsguard::prepare_queue_dir(dir)?;
let lock_path = dir.join("flush.lock");
fsguard::check_queue_file(&lock_path)?;
if !lock_path.exists() {
fsguard::create_private_file(&lock_path)?;
fsguard::check_queue_file(&lock_path)?;
}
let lock = std::fs::OpenOptions::new()
.truncate(false)
.write(true)
.open(&lock_path)
.with_context(|| format!("open {}", lock_path.display()))?;
// Serialize flushes across processes. The lock file itself is never removed:
// releasing the lock is dropping the handle.
if fs2::FileExt::try_lock_exclusive(&lock).is_err() {
bail!(
"another flush already holds {}; run one flush at a time",
lock_path.display()
);
}
let result = flush_locked(&dir, options);
let _ = fs2::FileExt::unlock(&lock);
result
}
fn flush_locked(dir: &Path, options: &FlushOptions) -> Result<Report> {
let mut queue = Queue::open(dir)?;
let binding = queue.binding()?;
queue.prune(crate::now_ms())?;
let token = Token::from_env()?;
let authenticated = token.is_some();
let sender = Sender::new(binding.server_url.clone(), token, CONNECT_TIMEOUT)?;
let started = Instant::now();
let mut deferred: HashSet<(String, String)> = HashSet::new();
let mut delivered = 0usize;
let mut batches = 0usize;
let mut blocked_sessions = 0usize;
let mut failure: Option<String> = None;
let mut backed_off = false;
for _ in 0..options.max_batches {
// Cheap pre-check, so an exhausted budget does not buy another round of
// selection and stamping. The binding one is taken just before the send.
if TOTAL_BUDGET.checked_sub(started.elapsed()).is_none() {
break;
}
let now = crate::now_ms();
let base = binding.server_url.clone();
let for_cost = binding.clone();
let batch: Batch =
queue.select_batch(MAX_BATCH_ITEMS, MAX_BATCH_BYTES, now, &deferred, |item| {
item_cost(&base, &for_cost, item)
})?;
blocked_sessions = batch.blocked_sessions;
if batch.items.is_empty() {
break;
}
batches += 1;
let keys: Vec<String> = batch.items.iter().map(|i| i.ingest_key.clone()).collect();
// Durable, and before the request: a lost ack must not restart the
// 30-day clock on the next run.
queue.stamp_attempt(&keys, now)?;
let items: Vec<BatchItem> = batch
.items
.iter()
.map(|item| {
Ok(BatchItem {
url: http::item_url(&binding.server_url, &binding, item),
body: serde_json::from_str(&item.body_json)
.context("stored body is no longer valid JSON")?,
})
})
.collect::<Result<Vec<_>>>()?;
// Selection checked the retry window once; a retry inside this batch can
// still cross it. Carry the batch's tightest expiry into the sender so
// every attempt is re-checked against it.
//
// Both clocks are read here, after stamping and JSON building: measuring
// from a `now` captured before that preparation would hand the batch
// however long the preparation took as extra allowance.
let send_now = crate::now_ms();
let ttl_left = batch
.items
.iter()
.map(|item| ttl_remaining(item.first_attempt_ms.unwrap_or(now), send_now))
.min()
.unwrap_or(Duration::ZERO);
let Some(remaining) = TOTAL_BUDGET.checked_sub(started.elapsed()) else {
break;
};
let delivery = match send_with_retries(&sender, &items, remaining, ttl_left) {
Ok(delivery) => delivery,
Err(e) => {
// Transport failure: everything stays pending, by definition.
for item in &batch.items {
queue.record_failure(&item.ingest_key, "transport")?;
}
failure = Some(format!("{e} (all {} item(s) kept)", batch.items.len()));
break;
}
};
// An ack is only meaningful on 200 and 429. Every other status keeps the
// whole batch, whatever its body claims.
if !matches!(delivery.status, 200 | 429) {
let detail = match delivery.status {
401 | 403 => {
if authenticated {
"server rejected the bearer token (AI_MEMORY_AUTH_TOKEN)"
} else {
"server requires authentication; set AI_MEMORY_AUTH_TOKEN"
}
}
413 => {
"the server refused the batch as too large. The relay's own batch bounds \
are fixed; record smaller event bodies at the producer, or raise the \
server's body limit"
}
_ => "unexpected status",
};
for item in &batch.items {
queue.record_failure(&item.ingest_key, &format!("http {}", delivery.status))?;
}
failure = Some(format!(
"HTTP {} from /hook/batch: {detail}; all {} item(s) kept pending",
delivery.status,
batch.items.len()
));
break;
}
let parsed: BatchAck = match serde_json::from_slice(&delivery.body) {
Ok(parsed) => parsed,
Err(_) => {
for item in &batch.items {
queue.record_failure(&item.ingest_key, "malformed ack")?;
}
failure = Some(format!(
"HTTP {} from /hook/batch with an unreadable ack ({} byte(s)); \
all {} item(s) kept pending",
delivery.status,
delivery.body.len(),
batch.items.len()
));
break;
}
};
let accepted = match ack::validate(batch.items.len(), &parsed) {
Ok(accepted) => accepted,
Err(rejected) => {
for item in &batch.items {
queue.record_failure(&item.ingest_key, "inconsistent ack")?;
}
failure = Some(format!(
"HTTP {} from /hook/batch with an inconsistent ack ({rejected}); \
all {} item(s) kept pending",
delivery.status,
batch.items.len()
));
break;
}
};
// Only now, after the whole response was validated, does anything leave
// the queue.
let confirmed: Vec<String> = accepted
.iter()
.filter_map(|idx| batch.items.get(*idx))
.map(|item| item.ingest_key.clone())
.collect();
delivered += queue.confirm(&confirmed, crate::now_ms())?;
if let Some(failed_index) = parsed.failed_index
&& let Some(item) = batch.items.get(failed_index)
{
queue.record_failure(&item.ingest_key, "server reported failed_index")?;
}
// The server stops at failed_index. Keep later, untried sessions
// eligible; otherwise one failed head could block them on every flush.
// Defer the failed head and earlier rate-limited items for this flush.
for (position, item) in batch.items.iter().enumerate() {
if accepted.contains(&position) {
continue;
}
let untried = parsed.failed_index.is_some_and(|failed| position > failed);
if !untried {
deferred.insert(item.session());
}
}
if delivery.status == 429 {
backed_off = true;
break;
}
}
let stats = queue.stats(crate::now_ms())?;
let mut summary = vec![format!(
"acknowledged {delivered} event(s) in {batches} batch(es); {} still pending across {} session(s)",
stats.pending_items, stats.pending_sessions
)];
if backed_off {
summary.push(
"server answered 429: the acknowledged items were applied and the flush stopped. \
Back off before the next run"
.to_owned(),
);
}
if !deferred.is_empty() {
summary.push(format!(
"{} session(s) deferred to a later flush (a failed or rate-limited head); \
their later events were never sent ahead of it",
deferred.len()
));
}
if blocked_sessions > 0 {
summary.push(format!(
"{blocked_sessions} session(s) held: their oldest event passed the 30-day retry \
window. The relay will not resend those automatically; they are retained for you"
));
}
summary.push(
"acknowledged means delivered or dropped by server policy, not that an observation \
was written"
.to_owned(),
);
Ok(Report {
summary,
pending: stats.pending_items > 0,
failure,
})
}
/// How long an event may still be auto-resent, measured from its first durable
/// attempt. Zero means the 30-day window is spent.
///
/// Accuracy depends on a stable host clock: the relay carries no time source of
/// its own, so a clock that jumps stretches or shortens this measurement.
pub fn ttl_remaining(first_attempt_ms: i64, now_ms: i64) -> Duration {
let spent = now_ms.saturating_sub(first_attempt_ms);
let left = identity::RETRY_WINDOW_MS.saturating_sub(spent);
Duration::from_millis(u64::try_from(left).unwrap_or(0))
}
/// Retry a batch a finite number of times under one shared deadline.
///
/// Two clocks bound it. `remaining` is the flush budget left when the batch
/// started; `ttl_left` is how long the batch's oldest item may still be resent.
/// Every request *and* every backoff sleep is drawn from the tighter of the two
/// and recomputed, so three attempts can neither overrun the budget nor put an
/// event on the wire after its retry window closed mid-flush.
fn send_with_retries(
sender: &Sender,
items: &[BatchItem],
remaining: Duration,
ttl_left: Duration,
) -> Result<http::Delivery> {
let start = Instant::now();
let deadline = start + remaining.min(ttl_left);
let ttl_deadline = start + ttl_left;
let mut attempt = 0usize;
loop {
if Instant::now() >= ttl_deadline {
bail!(
"hook batch not sent: the batch's oldest event reached the 30-day retry window \
mid-flush; it is retained, not resent"
);
}
let left = deadline.saturating_duration_since(Instant::now());
if left.is_zero() {
bail!("hook batch transport failure: flush budget exhausted");
}
match sender.post_batch(items, batch_timeout(items.len(), left)) {
Ok(delivery) => return Ok(delivery),
Err(e) => {
if attempt >= TRANSPORT_RETRIES {
return Err(e);
}
let backoff = Duration::from_millis(250)
.saturating_mul(1u32 << attempt.min(8))
.min(MAX_BACKOFF)
.min(deadline.saturating_duration_since(Instant::now()));
if backoff.is_zero() {
return Err(e);
}
attempt += 1;
std::thread::sleep(backoff);
}
}
}
}
/// Wire cost of one item, used for the batch byte budget.
pub fn item_cost(base: &str, binding: &Binding, item: &PendingItem) -> usize {
http::item_url(base, binding, item).len() + item.body_json.len() + ITEM_FRAMING_BYTES
}
/// Per-event timeout scaled by item count, capped twice: by an absolute maximum
/// and by whatever is left of the flush budget. Mirrors core's spool drain.
pub fn batch_timeout(items: usize, remaining: Duration) -> Duration {
let items = u32::try_from(items.max(1)).unwrap_or(u32::MAX);
PER_EVENT_TIMEOUT
.checked_mul(items)
.unwrap_or(Duration::MAX)
.min(MAX_BATCH_TIMEOUT)
.min(remaining)
.max(Duration::from_millis(1))
}
companions/ai-memory-relay/tests/relay.rs
+1342
View File
File diff suppressed because it is too large Load Diff
crates/ai-memory-cli/Cargo.toml
+5
View File
@@ -38,6 +38,9 @@ sha2.workspace = true
base64.workspace = true
anyhow.workspace = true
axum.workspace = true
# Sets SO_KEEPALIVE on accepted `serve` sockets (#792); not a workspace dep,
# `ai-memory-cli` is the only crate that needs it.
socket2 = "0.6.3"
clap.workspace = true
clap_complete.workspace = true
crossterm.workspace = true
@@ -73,6 +76,7 @@ serde_json.workspace = true
toml_edit.workspace = true
sysinfo.workspace = true
tar.workspace = true
zip.workspace = true
jiff.workspace = true
jsonc-parser.workspace = true
tokio.workspace = true
@@ -90,6 +94,7 @@ tempfile.workspace = true
ai-memory-test-support.workspace = true
tempfile.workspace = true
rstest.workspace = true
rusqlite.workspace = true
tower.workspace = true
tokio = { workspace = true, features = ["test-util"] }
crates/ai-memory-cli/src/cli.rs
+587 -20
View File
@@ -42,6 +42,11 @@ pub enum Command {
/// mid-project doesn't start amnesiac. No-op once the store has any
/// sessions unless `--force`.
Backfill(BackfillArgs),
/// Correct `sessions.started_at`/`ended_at` for sessions that `backfill`
/// already imported before it carried the transcript's own event times,
/// by re-reading the local transcripts and matching them by session id.
/// Dry-run by default; `--confirm` applies.
RepairBackfillTimestamps(RepairBackfillTimestampsArgs),
/// Launch an agent in an opt-in, cross-harness managed workstream.
/// Native arguments are forwarded except exact wrapper flags such as
/// `--yolo` and `--fresh`.
@@ -96,6 +101,23 @@ pub enum Command {
/// database, so every write blocks until it finishes and it needs free
/// disk space of roughly the database's own size.
Compact(CompactArgs),
/// Drop the superseded ledger versions the pre-2.1.1 indexer left behind.
///
/// #660 stopped the indexer from rewriting the whole `log-YYYY-MM.md` row
/// on every hook append. The fix stopped new rows; it did not remove the
/// ones already written, and no other command reaches them — `compact`
/// deletes nothing, `forget-sweep` only hard-deletes decay tombstones, and
/// `reindex` loses the DB-only state. One reported store held 6,539
/// versions of 101 live pages.
///
/// Only paths whose *content* is a hook event ledger are considered, so a
/// real page a human happened to name `log-2026-09.md` keeps its whole
/// version chain. Each dropped version is a byte prefix of the one after
/// it, so nothing the file on disk does not already hold is lost.
///
/// Prints what it would remove and changes nothing unless `--confirm` is
/// passed. See also `ai-memory status` for the reclaimable figure.
ReclaimLedgerVersions(ReclaimLedgerVersionsArgs),
/// Snapshot wiki/, db/, and config.toml into a gzipped tarball.
Backup(BackupArgs),
/// Export one project's wiki as an OKF v0.2 bundle tarball.
@@ -211,6 +233,13 @@ pub enum Command {
/// Remove ai-memory's wiring (hooks, MCP, instructions, and default-root
/// managed skills) from all detected agents. Dry-run unless `--apply`.
Uninstall(UninstallArgs),
/// Upgrade a GitHub-release native install: download the matching
/// release asset, verify its `.sha256`, atomically replace this
/// binary (and sibling `hooks/` when present), then re-stage hooks
/// for agents already under the data-dir hooks tree. Docker-wrapper
/// installs keep using the shell wrapper's `upgrade` (image pull);
/// package-managed installs (Homebrew, AUR, …) are refused.
Upgrade(UpgradeArgs),
/// Manage optional upstream LLM provider authentication.
Auth(AuthArgs),
/// Manage human users and deprecated 1.x compatibility tokens. All
@@ -220,6 +249,14 @@ pub enum Command {
/// the root bearer token and `[auth].token_pepper`.
#[command(name = "api-key")]
ApiKey(ApiKeyArgs),
/// Project settings. `project access` sets a project `open` (any user)
/// or `restricted` (root and grant holders) (#708). Requires the root
/// bearer token.
Project(ProjectArgs),
/// Manage local server profiles, which a repository's `.ai-memory.toml`
/// selects with `server = "<name>"` to route its hook capture to a
/// different ai-memory server.
Server(ServerArgs),
/// Print a shell-completion script to stdout. Generated from this
/// binary's own command tree, so it never drifts from the real CLI
/// surface. See `docs/shell-completions.md` for install paths.
@@ -251,6 +288,16 @@ pub struct RunArgs {
/// equivalent dangerous-mode option.
#[arg(long)]
pub yolo: bool,
/// Everything `--yolo` does, plus — for Claude — forcing
/// `bypassPermissions` via `--settings` over any settings `defaultMode`.
/// Claude still honors your own explicit `ask` rules in every mode. For
/// every other harness it is interchangeable with `--yolo`; passing both
/// is redundant but fine. Off by default; `[claude_true_yolo]` in
/// config.toml applies the same Claude extra to an explicit `--yolo`
/// launch. Best paired with ai-jail — see
/// `docs/design-yolo-safety-ai-jail.md`.
#[arg(long = "true-yolo")]
pub true_yolo: bool,
/// Start a new native session in the selected workstream instead of
/// resuming or adopting an existing harness session.
#[arg(long)]
@@ -261,6 +308,20 @@ pub struct RunArgs {
/// `AI_MEMORY_RUN_AUTOWIRE=false`) to launch without touching harness config.
#[arg(long)]
pub no_autowire: bool,
/// Extra environment variable for the spawned harness, `KEY=VALUE`.
/// Repeatable; wrapper-owned like `--yolo`/`--executable`, so it must
/// precede `harness`. Reaches the spawned process, ai-memory's own
/// native-session resolution and first-launch auto-wire (e.g.
/// `CLAUDE_CONFIG_DIR`), so session store, hooks and MCP agree on one
/// config home. A later `--env` wins over an earlier one and over a
/// same-key `--env-file` entry.
#[arg(long = "env", value_parser = parse_env_kv, value_name = "KEY=VALUE")]
pub env: Vec<(String, String)>,
/// Read `KEY=VALUE` lines from this file (blank lines and `#` comments
/// skipped) and merge them into the launch environment (same reach as
/// `--env`) before `--env` entries, which override a same-key line here.
#[arg(long = "env-file", value_name = "PATH")]
pub env_file: Option<PathBuf>,
/// Agent harness to launch. When omitted, continue the newest managed or
/// checkout-local session among the auto-detected harnesses. Any value
/// starting with `claude` (e.g. `claude-corp`, `claude-personal`) also
@@ -351,6 +412,23 @@ fn parse_run_harness_choice(value: &str) -> Result<RunHarnessChoice, String> {
))
}
/// Parse one `KEY=VALUE` entry for `--env` (also reused for `--env-file`
/// lines). The value is taken literally — no expansion, no interpretation —
/// so a caller-supplied value reaches the harness exactly as written.
pub(crate) fn parse_env_kv(value: &str) -> Result<(String, String), String> {
let Some((key, value)) = value.split_once('=') else {
return Err(format!(
"invalid value '{value}' for --env; expected KEY=VALUE"
));
};
if key.is_empty() {
return Err(format!(
"invalid value '{key}={value}' for --env; KEY must not be empty"
));
}
Ok((key.to_string(), value.to_string()))
}
/// Arguments for `show`.
#[derive(Debug, Args)]
pub struct ShowArgs {
@@ -368,6 +446,9 @@ pub struct ShowArgs {
/// equivalent dangerous-mode option. Forwarded to `run`.
#[arg(long)]
pub yolo: bool,
/// Claude-only true-yolo (see `RunArgs::true_yolo`). Forwarded to `run`.
#[arg(long = "true-yolo")]
pub true_yolo: bool,
/// Start a new native session instead of resuming or adopting an existing
/// harness session. Forwarded to `run`.
#[arg(long)]
@@ -391,6 +472,9 @@ pub struct ContinueArgs {
/// equivalent dangerous-mode option. Forwarded to `run`.
#[arg(long)]
pub yolo: bool,
/// Claude-only true-yolo (see `RunArgs::true_yolo`). Forwarded to `run`.
#[arg(long = "true-yolo")]
pub true_yolo: bool,
/// Start a new native session instead of resuming the linked one.
/// Forwarded to `run`.
#[arg(long)]
@@ -414,6 +498,9 @@ pub struct ResumeArgs {
/// equivalent dangerous-mode option. Forwarded to `run`.
#[arg(long)]
pub yolo: bool,
/// Claude-only true-yolo (see `RunArgs::true_yolo`). Forwarded to `run`.
#[arg(long = "true-yolo")]
pub true_yolo: bool,
/// Start a new native session instead of resuming the linked one.
/// Forwarded to `run`.
#[arg(long)]
@@ -630,6 +717,94 @@ pub struct UserArgs {
pub command: UserCommand,
}
/// Arguments for `user grant`.
#[derive(Debug, Args)]
pub struct UserGrantArgs {
/// The user, by username.
#[arg(long)]
pub user: String,
/// The workspace the project lives in.
#[arg(long, default_value = "default")]
pub workspace: String,
/// The project, by name.
#[arg(long)]
pub project: String,
/// `read` or `write`. Required: a level left unsaid is not guessed at.
#[arg(long)]
pub level: String,
}
/// Arguments for `user revoke`.
#[derive(Debug, Args)]
pub struct UserRevokeArgs {
/// The user, by username.
#[arg(long)]
pub user: String,
/// The workspace the project lives in.
#[arg(long, default_value = "default")]
pub workspace: String,
/// The project, by name.
#[arg(long)]
pub project: String,
}
/// Arguments for `user grants`.
#[derive(Debug, Args)]
pub struct UserGrantsArgs {
/// Only this user's grants. Omit for every grant on the server.
#[arg(long)]
pub user: Option<String>,
/// Emit the response as JSON instead of a table.
#[arg(long)]
pub json: bool,
}
/// Arguments for `project`.
#[derive(Debug, Args)]
pub struct ProjectArgs {
/// Project action to run.
#[command(subcommand)]
pub command: ProjectCommand,
}
/// `project` subcommands.
#[derive(Debug, Subcommand)]
pub enum ProjectCommand {
/// Set a project `open` or `restricted`.
Access(ProjectAccessArgs),
/// List who holds a grant on one project.
Grants(ProjectGrantsArgs),
}
/// Arguments for `project grants`.
#[derive(Debug, Args)]
pub struct ProjectGrantsArgs {
/// The workspace the project lives in.
#[arg(long, default_value = "default")]
pub workspace: String,
/// The project, by name.
#[arg(long)]
pub project: String,
/// Emit the response as JSON instead of a table.
#[arg(long)]
pub json: bool,
}
/// Arguments for `project access`.
#[derive(Debug, Args)]
pub struct ProjectAccessArgs {
/// The workspace the project lives in.
#[arg(long, default_value = "default")]
pub workspace: String,
/// The project, by name.
#[arg(long)]
pub project: String,
/// `open` (any authenticated user) or `restricted` (root and grant
/// holders). Required: a mode left unsaid is not guessed at.
#[arg(long)]
pub mode: String,
}
/// Arguments for `completions`.
#[derive(Debug, Args)]
pub struct CompletionsArgs {
@@ -669,6 +844,14 @@ pub enum UserCommand {
Enable(UserEnableArgs),
/// Update display name, email, and/or role (`root` or `user`).
Patch(UserPatchArgs),
/// Grant a user `read` or `write` on a project, or change the level they
/// hold (#708). Grants decide access to `restricted` projects; an `open`
/// project admits every user — see `ai-memory project access`.
Grant(UserGrantArgs),
/// Take away whatever a user holds on a project.
Revoke(UserRevokeArgs),
/// List grants: one user's with `--user`, else every grant on the server.
Grants(UserGrantsArgs),
}
/// Arguments for `user add`.
@@ -803,6 +986,61 @@ pub struct ApiKeyArgs {
pub command: ApiKeyCommand,
}
/// Arguments for `server`.
#[derive(Debug, Args)]
pub struct ServerArgs {
/// Server-profile action to run.
#[command(subcommand)]
pub command: ServerCommand,
}
/// Subcommands for `server`.
#[derive(Debug, Subcommand)]
pub enum ServerCommand {
/// Register a server profile, or replace one with the same name.
Add(ServerAddArgs),
/// List server profiles (never prints a token).
List(ServerListArgs),
/// Remove a server profile and its stored token.
Remove(ServerRemoveArgs),
}
/// Arguments for `server add`.
#[derive(Debug, Args)]
pub struct ServerAddArgs {
/// Profile name a marker selects: lowercase letters, digits, `-`, `_`.
pub name: String,
/// Server URL, e.g. `https://memory.example.com`.
#[arg(long)]
pub url: String,
/// Directory allowed to select this profile (repeatable; absolute or
/// `~/`). Required once more than one profile is registered.
#[arg(long = "root")]
pub roots: Vec<String>,
/// Bearer token for this server. Prefer `--auth-token-stdin`, which
/// keeps it out of shell history and the process table.
#[arg(long, hide_env_values = true, conflicts_with = "auth_token_stdin")]
pub auth_token: Option<String>,
/// Read the bearer token from the first line of stdin.
#[arg(long)]
pub auth_token_stdin: bool,
}
/// Arguments for `server list`.
#[derive(Debug, Args)]
pub struct ServerListArgs {
/// Emit the list as JSON.
#[arg(long)]
pub json: bool,
}
/// Arguments for `server remove`.
#[derive(Debug, Args)]
pub struct ServerRemoveArgs {
/// Profile name to remove.
pub name: String,
}
/// Subcommands for `api-key`.
#[derive(Debug, Subcommand)]
pub enum ApiKeyCommand {
@@ -964,12 +1202,26 @@ pub struct UninstallArgs {
/// Skip the interactive confirmation when a TTY is attached.
#[arg(long)]
pub yes: bool,
/// Profile to use for OMP extensions, which relocates the path to
/// `~/.omp/profiles/<profile>/agent/extensions/`.
/// OMP profile whose extension and MCP entry to remove, as `omp
/// --profile` names it. Beats `OMP_PROFILE` and `PI_PROFILE`; the default
/// profile's files are swept as well.
#[arg(long)]
pub profile: Option<String>,
}
/// Arguments for `upgrade`.
#[derive(Debug, Args)]
pub struct UpgradeArgs {
/// Pin a specific release tag (with or without a leading `v`). Defaults
/// to the latest GitHub Release for akitaonrails/ai-memory.
#[arg(long)]
pub version: Option<String>,
/// Re-download and replace even when the installed version already
/// matches the resolved release tag.
#[arg(long)]
pub force: bool,
}
/// Arguments for `reorg`.
#[derive(Debug, Args)]
pub struct ReorgArgs {
@@ -987,6 +1239,34 @@ pub struct CompactArgs {
pub confirm: bool,
}
/// Arguments for `reclaim-ledger-versions`.
#[derive(Debug, Args)]
pub struct ReclaimLedgerVersionsArgs {
/// Actually delete. Without this the command reports what it would remove
/// — the ledger paths, the row count and the bytes — and changes nothing.
#[arg(long)]
pub confirm: bool,
/// Also drop each ledger's *live* row, not just its superseded versions.
///
/// Since #660 the indexer skips ledgers, so each ledger's live row is also
/// left over from before the fix. It is kept by default: it is the
/// version the file on disk corresponds to, and dropping it is the
/// operator's call. Drop it only once the ledger file itself has been
/// removed, or the next hook append starts a fresh chain.
#[arg(long)]
pub drop_latest: bool,
/// Rebuild the FTS index and VACUUM afterwards, returning the freed bytes
/// to the filesystem.
///
/// Without this the reclaim is a logical delete — the rows are gone and
/// nothing can reach them, but their bytes stay in free pages of the
/// database file until it is next rewritten. `VACUUM` rewrites the whole
/// file under an exclusive lock and needs free disk of roughly the
/// database's own size, so it is opt-in.
#[arg(long)]
pub compact: bool,
}
/// Arguments for `purge-project`.
#[derive(Debug, Args)]
pub struct PurgeProjectArgs {
@@ -1005,6 +1285,18 @@ pub struct PurgeProjectArgs {
pub project: Option<String>,
/// REQUIRED for the purge to run. Without this flag the CLI errors
/// out — purging is destructive and irreversible.
///
/// Before erroring, the CLI asks the server for a preview (bounded to a
/// few seconds, auth refresh included): the reported counts (pages,
/// sessions, observations, handoffs, embeddings, workstreams, managed
/// runs, plus any collateral rows a purge of this project would delete
/// or orphan in *another* project) come from the same queries a
/// confirmed purge itself uses to decide what to delete. The preview is
/// best-effort and never changes the outcome, only what gets printed
/// before it: a 404/409/403 (or anything else unexpected) prints the
/// server's own error first; a timeout, an unreachable server, or an
/// older server that predates this preview just gets the plain refusal,
/// same as before this existed.
#[arg(long)]
pub confirm: bool,
/// Also reclaim the freed bytes: rebuild the FTS indexes and VACUUM the
@@ -1038,6 +1330,18 @@ pub struct PurgeSessionArgs {
pub project: Option<String>,
/// REQUIRED for the purge to run. Without this flag the CLI errors
/// out — purging is destructive and irreversible.
///
/// Before erroring, the CLI asks the server for a preview (bounded to a
/// few seconds, auth refresh included): the reported counts
/// (observations, handoffs, pages, auto-improve runs, plus any
/// collateral rows a purge of this session would delete or orphan in
/// *another* project) come from the same queries a confirmed purge
/// itself uses to decide what to delete. The preview is best-effort and
/// never changes the outcome, only what gets printed before it: a
/// 404/403 (or anything else unexpected) prints the server's own error
/// first; a timeout, an unreachable server, or an older server that
/// predates this preview just gets the plain refusal, same as before
/// this existed.
#[arg(long)]
pub confirm: bool,
/// Also reclaim the freed bytes: rebuild the affected FTS indexes and
@@ -1214,6 +1518,8 @@ pub enum InstallSkillsAgent {
Devin,
/// Grok Build CLI's `.grok/skills` directory.
Grok,
/// Hermes Agent's `.hermes/skills` directory.
Hermes,
/// Install into both Claude Code and `.agents` skill directories.
Both,
}
@@ -1269,11 +1575,14 @@ pub struct BootstrapArgs {
/// Maximum total tokens of source text sent to the LLM in one
/// run. When the collected sources exceed this, lower-priority
/// inputs (older git commits, then code module headers, then
/// docs) are dropped first. Default is 150K — comfortably under
/// Haiku/Sonnet 4.5's 200K context (leaves room for the ~64K
/// output budget) — so the model sees as much of your project
/// as possible. Lower it explicitly only if you're cost-
/// sensitive or running against a smaller-context provider.
/// docs) are dropped first. Default is 150K, so the model sees as
/// much of your project as possible; with chunking on (the
/// default), each call carries at most `--chunk-input-tokens` of
/// it plus up to 16K output tokens. Lower it if you're
/// cost-sensitive. With `--chunk-input-tokens 0`, this whole
/// budget goes into one call that also asks for up to 64K output
/// tokens, so budget + 64K must fit the model's context window
/// (the 150K default does not fit a 200K window).
#[arg(long, default_value_t = 150_000)]
pub max_input_tokens: usize,
/// Max estimated input tokens per LLM call. When pruned sources exceed
@@ -1414,6 +1723,43 @@ pub struct BackfillArgs {
/// been attempted for this checkout. Not for interactive use.
#[arg(long, hide = true)]
pub auto: bool,
/// Internal: deliver to the server the spawning hook is installed
/// against, instead of the configured one. Set by the SessionStart
/// trigger so the backfill cannot depend on the agent's environment.
#[arg(long, hide = true, conflicts_with = "server_profile")]
pub server_url: Option<String>,
/// Internal: deliver to this registered server profile (#992), with the
/// profile's own stored token. Set by the SessionStart trigger when the
/// repository's marker selects a profile.
#[arg(long, hide = true)]
pub server_profile: Option<String>,
}
/// Arguments for `repair-backfill-timestamps`.
///
/// A thin client like every other lifecycle command: it reads the local
/// transcripts (read-only, reusing `backfill`'s own discovery) to compute
/// candidate `started_at`/`ended_at` values, then posts them to
/// `POST /admin/repair-session-times`, which validates each one against the
/// scope and the backfill bug's own signature, and applies (or, without
/// `--confirm`, only reports) the change.
#[derive(Debug, Args)]
pub struct RepairBackfillTimestampsArgs {
/// Workspace name. Defaults to the current project's resolved scope.
#[arg(long)]
pub workspace: Option<String>,
/// Project name. Defaults to the current project's resolved scope.
#[arg(long)]
pub project: Option<String>,
/// Apply the computed times. Without this flag the command only reports
/// what would change (sessions repaired/skipped, by reason, plus the
/// before/after date range) — the server runs the write inside a
/// rolled-back transaction, so this is a real dry run, not an estimate.
#[arg(long)]
pub confirm: bool,
/// Emit the server's report(s) as JSON instead of the human summary.
#[arg(long)]
pub json: bool,
}
/// Arguments for `doctor`.
@@ -1618,8 +1964,9 @@ pub enum AgentChoice {
/// lifecycle capture and bridges ai-memory's HTTP MCP tools into Pi.
Pi,
/// Oh My Pi (`omp`) — TypeScript extension
/// under `~/.omp/agent/extensions/`. `--apply` writes the extension
/// file directly; restart `omp` for it to load.
/// under `~/.omp/agent/extensions/`, or the active profile's agent dir.
/// `--apply` writes the extension file directly; restart `omp` for it to
/// load.
#[value(alias = "oh-my-pi")]
Omp,
/// OpenClaw personal AI gateway — native plugin package with
@@ -1686,6 +2033,19 @@ pub enum AgentChoice {
/// so close sessions with `ai-memory finalize-session --agent zcode`.
#[value(alias = "zai")]
Zcode,
/// Hermes Agent (Nous Research) — lifecycle hooks declared in the `hooks:`
/// block of `~/.hermes/config.yaml`. Hermes runs each `command` through
/// `shlex.split` with the event JSON on stdin and **no shell**, so
/// ai-memory's native `hook` subcommand is invoked directly (exec form,
/// like Zero and ZCode). ai-memory wires the two events that give Hermes
/// tool observations — `pre_tool_call` / `post_tool_call`, whose payload
/// carries `tool_name` / `tool_input`, the envelope the router already
/// recognises for `agent=hermes`. `~/.hermes/config.yaml` is NOT written:
/// it is YAML the installer would have to splice, and Hermes gates user
/// hooks behind its own acceptance prompt (`hooks_auto_accept`), so
/// `install-hooks --agent hermes` prints the ready-to-paste block.
#[value(alias = "hermes-agent")]
Hermes,
}
impl AgentChoice {
@@ -1715,6 +2075,7 @@ impl AgentChoice {
Self::CommandCode => AgentKind::CommandCode,
Self::Pool => AgentKind::Pool,
Self::Zcode => AgentKind::Zcode,
Self::Hermes => AgentKind::Hermes,
}
}
@@ -1732,7 +2093,8 @@ impl AgentChoice {
| Self::Omp
| Self::Openclaw
| Self::Zero
| Self::Zcode => None,
| Self::Zcode
| Self::Hermes => None,
_ => Some(self.kind().as_str()),
}
}
@@ -1796,6 +2158,19 @@ pub struct FinalizeSessionArgs {
/// (still-active) session.
#[arg(long, conflicts_with = "all")]
pub session_id: Option<ai_memory_core::SessionId>,
/// Re-finalize a session that already ended (requires `--session-id`).
///
/// Use this when the conversation continued after a first finalize and
/// landed new observations: agents without a true session-end event
/// (Antigravity CLI, Kiro, ZCode, Pool) keep capturing under the same
/// session id, but the plain discovery step only sees open sessions, so
/// a second finalize would silently find nothing. With `--reopen` the
/// lookup also matches the ended session and the normal session-end
/// path re-runs (updated summary page, handoff, opt-in consolidation).
/// When nothing new landed since the first end, the re-run is a
/// harmless no-op.
#[arg(long, requires = "session_id")]
pub reopen: bool,
/// Emit a JSON summary.
#[arg(long)]
pub json: bool,
@@ -1844,7 +2219,7 @@ impl SchemaFlavor {
pub enum McpClient {
/// Anthropic Claude Code — `claude mcp add`.
ClaudeCode,
/// OpenAI Codex CLI — `~/.codex/config.toml`.
/// OpenAI Codex CLI — `$CODEX_HOME/config.toml` (default `~/.codex/config.toml`).
Codex,
/// OpenCode — `opencode.json`. Accepts `opencode` (no hyphen) as
/// an alias for symmetry with `AgentChoice` and the on-disk
@@ -1870,7 +2245,8 @@ pub enum McpClient {
/// Real Pi coding agent. Uses ai-memory's generated bridge extension
/// because Pi has no native MCP config.
Pi,
/// Oh My Pi (`omp`) — `~/.omp/agent/mcp.json`.
/// Oh My Pi (`omp`) — `~/.omp/agent/mcp.json`, or the active profile's
/// agent dir.
#[value(alias = "oh-my-pi")]
Omp,
/// Google Antigravity CLI (`agy`) — `~/.gemini/config/mcp_config.json`.
@@ -2270,7 +2646,7 @@ pub struct HookArgs {
/// the local spool or the wire.
#[arg(long, value_enum)]
pub capture_mode: Option<CaptureModeArg>,
/// Opt in to assistant/Stop capture: on a Claude Code `stop` event, attach a
/// Opt in to assistant/Stop capture: on a supported agent's `stop` event, attach a
/// sanitized, capped excerpt of the assistant's final turn as the Stop body.
/// Baked onto the native `stop` command by
/// `install-hooks --capture-assistant`; the server must also enable
@@ -2348,11 +2724,10 @@ pub struct InstallHooksArgs {
/// silently revert `repo-root` back to `basename`.
#[arg(long, value_enum)]
pub project_strategy: Option<ProjectStrategyArg>,
/// Bake `--capture-assistant` onto the installed native `stop` command so a
/// Claude Code `stop` event carries a sanitized excerpt of the assistant's
/// final turn (#196). Only valid for `--agent claude-code` on a native
/// platform; the server must also set `capture_assistant = true`. Re-running
/// without this flag removes it (idempotent). Default off.
/// Capture a sanitized excerpt of the assistant's final turn through the
/// native Stop hook. Supported for Claude Code, Codex and OpenCode 2 on a
/// native platform; the server must also set `capture_assistant = true`.
/// A bare re-apply preserves an existing opt-in. Default off.
#[arg(long)]
pub capture_assistant: bool,
/// Persist the capture failure mode for this install (#446). Under
@@ -2373,8 +2748,9 @@ pub struct InstallHooksArgs {
/// `--no-capture-prompts` install. Only valid for Claude Code.
#[arg(long, conflicts_with = "no_capture_prompts")]
pub capture_prompts: bool,
/// Profile to use for OMP extensions, which relocates the path to
/// `~/.omp/profiles/<profile>/agent/extensions/`.
/// OMP profile to install into, as `omp --profile` names it:
/// `~/.omp/profiles/<profile>/agent/extensions/`. Beats `OMP_PROFILE` and
/// `PI_PROFILE`; `default` selects the default profile.
#[arg(long)]
pub profile: Option<String>,
}
@@ -2569,6 +2945,73 @@ mod tests {
use clap::{CommandFactory, Parser};
use std::collections::BTreeSet;
/// The management surface uses the design's spelling (#708):
/// `user grant --user … --workspace … --project … --level …`, `user revoke`,
/// and listings under `user grants` / `project grants`. A level is never
/// defaulted, and the former top-level `grant` command is gone.
#[test]
fn grant_commands_use_the_designs_spelling() {
let parsed = Cli::try_parse_from([
"ai-memory",
"user",
"grant",
"--user",
"alice",
"--workspace",
"acme",
"--project",
"api",
"--level",
"write",
])
.expect("user grant parses");
let Command::User(UserArgs {
command: UserCommand::Grant(args),
}) = parsed.command
else {
panic!("expected user grant");
};
assert_eq!(
(
args.user.as_str(),
args.workspace.as_str(),
args.project.as_str(),
args.level.as_str()
),
("alice", "acme", "api", "write")
);
assert!(
Cli::try_parse_from([
"ai-memory",
"user",
"grant",
"--user",
"alice",
"--project",
"api"
])
.is_err(),
"a level left unsaid is not guessed at"
);
for argv in [
&[
"ai-memory",
"user",
"revoke",
"--user",
"alice",
"--project",
"api",
][..],
&["ai-memory", "user", "grants"][..],
&["ai-memory", "user", "grants", "--user", "alice"][..],
&["ai-memory", "project", "grants", "--project", "api"][..],
] {
Cli::try_parse_from(argv).unwrap_or_else(|e| panic!("{argv:?}: {e}"));
}
assert!(Cli::try_parse_from(["ai-memory", "grant", "list"]).is_err());
}
#[test]
fn serve_parses_insecure_no_auth_override() {
let parsed = Cli::try_parse_from([
@@ -3199,6 +3642,88 @@ mod tests {
);
}
#[test]
fn run_env_flag_parses_repeatable_key_value_pairs() {
let cli = Cli::try_parse_from([
"ai-memory",
"run",
"--env",
"CLAUDE_CONFIG_DIR=/accounts/work",
"--env",
"FOO=bar=baz",
"claude",
])
.expect("valid --env pairs parse");
let Command::Run(args) = cli.command else {
panic!("expected run command");
};
assert_eq!(
args.env,
vec![
(
"CLAUDE_CONFIG_DIR".to_string(),
"/accounts/work".to_string()
),
("FOO".to_string(), "bar=baz".to_string()),
]
);
}
#[test]
fn run_env_flag_rejects_a_pair_without_equals() {
let error = Cli::try_parse_from(["ai-memory", "run", "--env", "NOEQUALS", "claude"])
.expect_err("a value without '=' must be rejected");
assert!(
error.to_string().contains("expected KEY=VALUE"),
"unexpected error: {error}"
);
}
#[test]
fn run_env_flag_rejects_an_empty_key() {
let error = Cli::try_parse_from(["ai-memory", "run", "--env", "=value", "claude"])
.expect_err("an empty key must be rejected");
assert!(
error.to_string().contains("KEY must not be empty"),
"unexpected error: {error}"
);
}
#[test]
fn run_env_file_flag_parses_as_a_path() {
let cli = Cli::try_parse_from([
"ai-memory",
"run",
"--env-file",
"/tmp/ai-memory-env-example.env",
"claude",
])
.expect("--env-file parses");
let Command::Run(args) = cli.command else {
panic!("expected run command");
};
assert_eq!(
args.env_file,
Some(PathBuf::from("/tmp/ai-memory-env-example.env"))
);
}
#[test]
fn parse_env_kv_accepts_pairs_and_rejects_malformed_entries() {
assert_eq!(
parse_env_kv("KEY=VALUE"),
Ok(("KEY".to_string(), "VALUE".to_string()))
);
// The value is taken literally, including any further '=' signs.
assert_eq!(
parse_env_kv("KEY=a=b=c"),
Ok(("KEY".to_string(), "a=b=c".to_string()))
);
assert_eq!(parse_env_kv("KEY="), Ok(("KEY".to_string(), String::new())));
assert!(parse_env_kv("NOEQUALS").is_err());
assert!(parse_env_kv("=value").is_err());
}
#[test]
fn devin_hook_agent_parses() {
let hook_cli = Cli::try_parse_from([
@@ -3306,6 +3831,37 @@ mod tests {
assert_eq!(args.agent, ai_memory_core::AgentKind::Zcode);
}
/// Hermes is the second no-shell harness (after ZCode): it splits the
/// configured `command` into argv itself, so the generated block invokes
/// the native `hook` subcommand instead of a `.sh` bundle.
#[test]
fn hermes_hook_and_finalize_aliases_parse() {
for alias in ["hermes", "hermes-agent"] {
let cli = Cli::try_parse_from([
"ai-memory",
"install-hooks",
"--agent",
alias,
"--server-url",
"http://127.0.0.1:49374",
])
.unwrap_or_else(|error| panic!("failed to parse Hermes alias {alias}: {error}"));
let Command::InstallHooks(args) = cli.command else {
panic!("expected install-hooks for Hermes alias {alias}");
};
assert_eq!(args.agent, AgentChoice::Hermes);
assert_eq!(args.agent.kind(), ai_memory_core::AgentKind::Hermes);
// Native exec-form integration: no script bundle to stage.
assert_eq!(args.agent.script_hook_subdir(), None);
}
let cli = Cli::try_parse_from(["ai-memory", "finalize-session", "--agent", "hermes"])
.expect("failed to parse finalize-session --agent hermes");
let Command::FinalizeSession(args) = cli.command else {
panic!("expected finalize-session for hermes");
};
assert_eq!(args.agent, ai_memory_core::AgentKind::Hermes);
}
#[test]
fn command_code_mcp_and_hook_aliases_parse() {
for alias in ["command-code", "commandcode", "cmdc", "cmd"] {
@@ -3712,4 +4268,15 @@ mod tests {
fn completions_requires_a_shell() {
assert!(Cli::try_parse_from(["ai-memory", "completions"]).is_err());
}
#[test]
fn upgrade_parses_version_and_force() {
let cli = Cli::try_parse_from(["ai-memory", "upgrade", "--version", "v2.3.2", "--force"])
.unwrap();
let Command::Upgrade(args) = cli.command else {
panic!("expected upgrade command");
};
assert_eq!(args.version.as_deref(), Some("v2.3.2"));
assert!(args.force);
}
}
crates/ai-memory-cli/src/commands/auto_improve.rs
+2
View File
@@ -20,6 +20,7 @@ struct AutoImproveRequest {
max_input_tokens: usize,
max_proposals_per_run: usize,
max_patchable_pages: usize,
patchable_page_prefixes: Vec<String>,
max_patchable_body_chars: usize,
max_edits_per_proposal: usize,
max_edit_content_chars: usize,
@@ -83,6 +84,7 @@ pub async fn run(config: &Config, args: AutoImproveArgs) -> Result<()> {
max_input_tokens: args.max_input_tokens.unwrap_or(settings.max_input_tokens),
max_proposals_per_run: args.max_proposals.unwrap_or(settings.max_proposals_per_run),
max_patchable_pages: settings.max_patchable_pages,
patchable_page_prefixes: settings.patchable_page_prefixes.clone(),
max_patchable_body_chars: settings.max_patchable_body_chars,
max_edits_per_proposal: settings.max_edits_per_proposal,
max_edit_content_chars: settings.max_edit_content_chars,
crates/ai-memory-cli/src/commands/backfill.rs
+381 -51
View File
@@ -40,11 +40,11 @@ use serde::{Deserialize, Serialize};
use ai_memory_core::{NewWorkstreamEvent, WorkstreamEventKind};
use ai_memory_workstream::{
ManagedHarness, build_launch_plan, export_transcript, list_native_sessions,
wait_for_transcript_flush,
LaunchRoots, ManagedHarness, build_launch_plan_with_env, export_transcript,
list_native_sessions, wait_for_transcript_flush,
};
use super::doctor::SCANNED_HARNESSES;
use super::doctor::{SCANNED_HARNESSES, relocated_session_dir};
use super::run;
use crate::config::Config;
use crate::http_client::{ServerEndpoint, get_json, post_json};
@@ -68,9 +68,9 @@ const BACKFILL_EXTENSION: &str = "ai-memory-backfill";
/// A local native session eligible for import.
#[derive(Debug, Clone)]
pub(crate) struct SessionRef {
harness: ManagedHarness,
native_session_id: String,
updated_at: SystemTime,
pub(crate) harness: ManagedHarness,
pub(crate) native_session_id: String,
pub(crate) updated_at: SystemTime,
}
/// The outcome of a backfill run, and the JSON output shape.
@@ -145,11 +145,48 @@ fn write_sentinel(data_dir: &Path, cwd: &Path) {
let _ = std::fs::write(&path, b"");
}
/// The server this backfill delivers to.
///
/// A SessionStart-spawned run names the hook's own target: a server profile
/// (#992), whose stored token is the only credential it will present, or the
/// install-time hook URL, authenticated exactly like the hook's own events to
/// it: the persisted hook token, then OIDC. The config/env bearer is never
/// used there, because it may belong to a different server than the one the
/// hook is installed against. A manual run keeps resolving from config.
async fn backfill_endpoint(
config: &Config,
args: &crate::cli::BackfillArgs,
) -> Result<ServerEndpoint> {
if let Some(raw) = args.server_profile.as_deref() {
let name = crate::server_profiles::ProfileName::parse(raw)
.with_context(|| format!("`{raw}` is not a valid server profile name"))?;
let profile = crate::server_profiles::lookup(&config.data_dir, &name)
.map_err(|r| anyhow::anyhow!("server profile `{name}` was refused ({})", r.as_str()))?;
return Ok(ServerEndpoint::for_hook_target(
profile.url,
Some(profile.token),
));
}
if let Some(url) = args.server_url.as_deref() {
let static_token = crate::config::read_hook_auth_token(&config.data_dir);
let token = super::hook_spool::resolve_bearer(
&reqwest::Client::new(),
&config.data_dir,
static_token.as_deref(),
)
.await;
return Ok(ServerEndpoint::for_hook_target(url.to_owned(), token));
}
Ok(ServerEndpoint::from_config_resolving_auth(config).await)
}
/// Run the backfill.
///
/// # Errors
/// Returns an error when the scope cannot be resolved, the working directory
/// cannot be read, or the server is unreachable for the emptiness check.
/// cannot be read, the server is unreachable for the emptiness check, or any
/// selected session fails to import. The report is emitted before import errors
/// are returned, including in JSON mode.
pub async fn run(config: &Config, args: crate::cli::BackfillArgs) -> Result<()> {
let cwd = std::env::current_dir().context("resolving the current working directory")?;
@@ -157,14 +194,16 @@ pub async fn run(config: &Config, args: crate::cli::BackfillArgs) -> Result<()>
// that it has been attempted for this checkout, so it runs at most once per
// machine regardless of outcome. Manual runs ignore both.
if args.auto && !config.backfill_on_start {
write_sentinel(&config.data_dir, &cwd);
if !args.dry_run {
write_sentinel(&config.data_dir, &cwd);
}
return Ok(());
}
let (workspace, project) =
super::resolve_scope(config, args.workspace.as_deref(), args.project.as_deref())?;
let home = run::native_home(config).context("locating the local harness session stores")?;
let endpoint = ServerEndpoint::from_config_resolving_auth(config).await;
let endpoint = backfill_endpoint(config, &args).await?;
let mut report = BackfillReport {
workspace: workspace.clone(),
@@ -180,18 +219,20 @@ pub async fn run(config: &Config, args: crate::cli::BackfillArgs) -> Result<()>
if !empty {
report.skipped_non_empty = true;
// Mark attempted so the auto-trigger stops probing this checkout.
write_sentinel(&config.data_dir, &cwd);
if !args.dry_run {
write_sentinel(&config.data_dir, &cwd);
}
return finish(&args, &report);
}
}
// Enumerate local sessions for this cwd across every supported harness.
let candidates = collect_local_sessions(&home, &cwd, args.session.as_deref()).await;
let (candidates, _limit_hit) =
collect_local_sessions(&home, &cwd, args.session.as_deref()).await;
let selected = select_sessions(candidates, args.max_sessions.max(1));
report.selected = selected.len();
if args.dry_run {
write_sentinel(&config.data_dir, &cwd);
return finish(&args, &report);
}
@@ -209,19 +250,27 @@ pub async fn run(config: &Config, args: crate::cli::BackfillArgs) -> Result<()>
}
Err(error) => {
report.failed_sessions += 1;
if !args.quiet {
eprintln!(
"ai-memory: backfill of {} session {} failed: {error:#}",
session.harness.as_str(),
display_id(&session.native_session_id)
);
}
// Quiet suppresses the success summary, not failures: the
// detached worker's stderr is the operator's diagnostic log.
eprintln!(
"ai-memory: backfill of {} session {} failed: {error:#}",
session.harness.as_str(),
display_id(&session.native_session_id)
);
}
}
}
write_sentinel(&config.data_dir, &cwd);
finish(&args, &report)
finish(&args, &report)?;
if report.failed_sessions > 0 {
bail!(
"backfill failed to import {} of {} selected session(s)",
report.failed_sessions,
report.selected
);
}
Ok(())
}
/// Sum the server's per-agent session counts for this scope; zero means empty.
@@ -250,16 +299,33 @@ async fn project_is_empty(
/// Enumerate local native sessions for `cwd` across every scanned harness.
/// Read-only; a harness whose store is absent/unreadable contributes nothing.
async fn collect_local_sessions(
///
/// Returns, alongside the sessions, the harnesses whose scan came back at
/// exactly [`PER_HARNESS_SCAN_LIMIT`] — a caller that cares about missing
/// older sessions (as opposed to `backfill`'s own newest-first + cap
/// selection, which does not) can surface that.
pub(crate) async fn collect_local_sessions(
home: &Path,
cwd: &Path,
only_session: Option<&str>,
) -> Vec<SessionRef> {
) -> (Vec<SessionRef>, Vec<ManagedHarness>) {
collect_local_sessions_with(home, cwd, only_session, relocated_session_dir).await
}
/// [`collect_local_sessions`] with the relocation lookup passed in, for the
/// same reason as `doctor::scan_local_with`: tests pass `|_| None` so a
/// developer's `CLAUDE_CONFIG_DIR` cannot hide a fixture planted under a
/// temporary `$HOME`.
pub(crate) async fn collect_local_sessions_with(
home: &Path,
cwd: &Path,
only_session: Option<&str>,
session_dir_for: impl Fn(ManagedHarness) -> Option<PathBuf>,
) -> (Vec<SessionRef>, Vec<ManagedHarness>) {
let mut out = Vec::new();
let mut limit_hit = Vec::new();
for &harness in SCANNED_HARNESSES {
let session_dir = build_launch_plan(harness, None, Vec::new(), None)
.ok()
.and_then(|plan| plan.session_dir);
let session_dir = session_dir_for(harness);
let Ok(sessions) = list_native_sessions(
harness,
home,
@@ -271,6 +337,9 @@ async fn collect_local_sessions(
else {
continue;
};
if sessions.len() >= PER_HARNESS_SCAN_LIMIT {
limit_hit.push(harness);
}
for session in sessions {
if only_session.is_some_and(|want| want != session.native_session_id) {
continue;
@@ -282,7 +351,7 @@ async fn collect_local_sessions(
});
}
}
out
(out, limit_hit)
}
/// One item in a `POST /hook/batch` request: the full hook URL (whose query the
@@ -315,9 +384,11 @@ async fn import_one(
cwd: &Path,
session: &SessionRef,
) -> Result<usize> {
let session_dir = build_launch_plan(session.harness, None, Vec::new(), None)
.ok()
.and_then(|plan| plan.session_dir);
let roots = LaunchRoots { home, cwd };
let session_dir =
build_launch_plan_with_env(session.harness, None, Vec::new(), None, &[], Some(roots))
.ok()
.and_then(|plan| plan.session_dir);
// These are historical sessions, so the flush wait is a quick no-op; ignore
// its result and read whatever is on disk.
let _ = wait_for_transcript_flush(
@@ -341,6 +412,7 @@ async fn import_one(
let sid = &session.native_session_id;
let agent = session.harness.agent_kind().as_str();
let resolved = resolve_occurred_at(&transcript.events);
let mut items = Vec::with_capacity(transcript.events.len() + 2);
items.push(hook_item(
endpoint,
@@ -351,11 +423,11 @@ async fn import_one(
sid,
&format!("{sid}:session-start"),
None,
serde_json::json!({ "session_id": sid }),
serde_json::json!({ "session_id": sid, "occurred_at": resolved.earliest }),
)?);
let mut content = 0usize;
for event in &transcript.events {
if let Some(mapped) = map_event(sid, event) {
for (event, occurred_at) in transcript.events.iter().zip(&resolved.per_event) {
if let Some(mapped) = map_event(sid, event, occurred_at.as_deref()) {
items.push(hook_item(
endpoint,
workspace,
@@ -379,13 +451,85 @@ async fn import_one(
sid,
&format!("{sid}:session-end"),
None,
serde_json::json!({ "session_id": sid }),
serde_json::json!({ "session_id": sid, "occurred_at": resolved.latest }),
)?);
post_hook_items(endpoint, &items).await?;
Ok(content)
}
/// Per-event `occurred_at` resolution for one transcript, plus the session's
/// overall boundary times.
struct ResolvedOccurredAt {
/// Effective `occurred_at` for each event, in transcript order (RFC 3339).
per_event: Vec<Option<String>>,
/// The earliest valid event time — the session-start's `occurred_at`.
/// A transcript is not guaranteed to be time-sorted (a reordered or
/// clock-skewed import), so the first *entry* is not reliably the
/// earliest *time*.
earliest: Option<String>,
/// The latest valid event time — the session-end's `occurred_at`, same
/// reasoning as `earliest`.
latest: Option<String>,
}
/// Resolve each event's effective `occurred_at`, letting one missing or
/// unparsable own timestamp inherit the nearest preceding *valid* one; an
/// event before the first valid timestamp inherits that first one instead of
/// staying unresolved (there is nothing earlier to inherit from). A value
/// that fails to parse as RFC 3339 is treated exactly like a missing one — it
/// never reaches the hook body, so a malformed transcript timestamp cannot
/// masquerade as a validated one downstream.
fn resolve_occurred_at(events: &[NewWorkstreamEvent]) -> ResolvedOccurredAt {
let valid: Vec<Option<jiff::Timestamp>> = events
.iter()
.map(|event| {
event
.occurred_at
.as_deref()
.and_then(|s| s.parse::<jiff::Timestamp>().ok())
})
.collect();
let mut per_event: Vec<Option<jiff::Timestamp>> = Vec::with_capacity(events.len());
let mut last_valid: Option<jiff::Timestamp> = None;
for ts in &valid {
if ts.is_some() {
last_valid = *ts;
}
per_event.push(last_valid);
}
// Backward-fill the leading gap: events before the first valid timestamp
// had nothing preceding them to inherit above.
if let Some(first_valid) = valid.iter().copied().flatten().next() {
for slot in per_event.iter_mut() {
match slot {
Some(_) => break,
None => *slot = Some(first_valid),
}
}
}
let (earliest, latest) = valid.into_iter().flatten().fold(
(None, None),
|(min, max): (Option<jiff::Timestamp>, Option<jiff::Timestamp>), ts| {
(
Some(min.map_or(ts, |m| m.min(ts))),
Some(max.map_or(ts, |m| m.max(ts))),
)
},
);
ResolvedOccurredAt {
per_event: per_event
.into_iter()
.map(|ts| ts.map(|t| t.to_string()))
.collect(),
earliest: earliest.map(|t| t.to_string()),
latest: latest.map(|t| t.to_string()),
}
}
/// A transcript event mapped to its `/hook` shape.
struct MappedEvent {
event: String,
@@ -398,7 +542,15 @@ struct MappedEvent {
/// `user-prompt` observation; every other content-bearing event is recorded as
/// a backfill extension observation. Non-content boundary events (compaction,
/// checkpoint, annotation) are dropped — they are not session content.
fn map_event(session_id: &str, event: &NewWorkstreamEvent) -> Option<MappedEvent> {
///
/// `occurred_at` (RFC 3339), already resolved by the caller, rides in the hook
/// body so the imported observation is dated at the transcript's own event
/// time rather than at import time.
fn map_event(
session_id: &str,
event: &NewWorkstreamEvent,
occurred_at: Option<&str>,
) -> Option<MappedEvent> {
let content = event.content.trim();
if content.is_empty() {
return None;
@@ -410,7 +562,11 @@ fn map_event(session_id: &str, event: &NewWorkstreamEvent) -> Option<MappedEvent
event: "user-prompt".to_string(),
source_event: None,
ingest_key,
body: serde_json::json!({ "session_id": session_id, "prompt": event.content }),
body: serde_json::json!({
"session_id": session_id,
"prompt": event.content,
"occurred_at": occurred_at,
}),
}),
WorkstreamEventKind::Message
| WorkstreamEventKind::ToolCall
@@ -429,6 +585,7 @@ fn map_event(session_id: &str, event: &NewWorkstreamEvent) -> Option<MappedEvent
"session_id": session_id,
"title": first_line(content),
"message": event.content,
"occurred_at": occurred_at,
}),
})
}
@@ -552,6 +709,9 @@ fn finish(args: &crate::cli::BackfillArgs, report: &BackfillReport) -> Result<()
"📼 ai-memory imported {} prior local session(s) (~{} events) for {}/{}.",
report.imported_sessions, report.imported_events, report.workspace, report.project
);
if report.failed_sessions > 0 {
line.push_str(&format!(" {} session(s) failed.", report.failed_sessions));
}
if report.skipped_for_cap > 0 {
line.push_str(&format!(
" {} older session(s) skipped (import cap).",
@@ -644,6 +804,7 @@ mod tests {
let m = map_event(
"sid",
&event(WorkstreamEventKind::Message, Some("user"), "do the thing"),
Some("2026-09-10T12:00:00Z"),
)
.expect("user message maps");
assert_eq!(m.event, "user-prompt");
@@ -652,6 +813,7 @@ mod tests {
"user-prompt is a lifecycle event, not an extension"
);
assert_eq!(m.body["prompt"], "do the thing");
assert_eq!(m.body["occurred_at"], "2026-09-10T12:00:00Z");
assert_eq!(
m.ingest_key, "sid:evt-1",
"ingest key is stable per source event"
@@ -667,6 +829,7 @@ mod tests {
Some("assistant"),
"here is the plan\nline2",
),
None,
)
.expect("assistant maps");
assert_eq!(a.event, "backfill.assistant-message");
@@ -676,22 +839,134 @@ mod tests {
"title is the first line"
);
assert_eq!(a.body["message"], "here is the plan\nline2");
assert!(
a.body["occurred_at"].is_null(),
"no occurred_at was supplied"
);
let t = map_event(
"sid",
&event(WorkstreamEventKind::ToolCall, None, "grep foo"),
None,
)
.expect("tool call maps");
assert_eq!(t.event, "backfill.tool_call");
assert_eq!(t.source_event.as_deref(), Some("tool_call"));
}
/// Round-trips a literal RFC 3339 string through `jiff::Timestamp` so
/// expectations match `resolve_occurred_at`'s own parse-then-format
/// output rather than assuming it echoes the input string verbatim.
fn ts(literal: &str) -> String {
literal.parse::<jiff::Timestamp>().unwrap().to_string()
}
#[test]
fn resolve_occurred_at_fills_gaps_from_the_preceding_event() {
let mut e1 = event(WorkstreamEventKind::Message, Some("user"), "one");
e1.occurred_at = Some("2026-09-10T12:00:00Z".to_string());
let e2 = event(WorkstreamEventKind::Message, Some("assistant"), "two"); // no timestamp
let mut e3 = event(WorkstreamEventKind::ToolCall, None, "three");
e3.occurred_at = Some("2026-09-10T12:05:00Z".to_string());
let e4 = event(WorkstreamEventKind::ToolResult, None, "four"); // no timestamp
let resolved = resolve_occurred_at(&[e1, e2, e3, e4]);
assert_eq!(
resolved.per_event,
vec![
Some(ts("2026-09-10T12:00:00Z")),
Some(ts("2026-09-10T12:00:00Z")),
Some(ts("2026-09-10T12:05:00Z")),
Some(ts("2026-09-10T12:05:00Z")),
],
"an event without its own timestamp inherits the nearest preceding one"
);
assert_eq!(resolved.earliest, Some(ts("2026-09-10T12:00:00Z")));
assert_eq!(resolved.latest, Some(ts("2026-09-10T12:05:00Z")));
}
#[test]
fn resolve_occurred_at_backfills_the_leading_gap_from_the_first_valid_timestamp() {
let e1 = event(WorkstreamEventKind::Message, Some("user"), "one"); // no timestamp
let e2 = event(WorkstreamEventKind::Message, Some("assistant"), "two"); // no timestamp
let mut e3 = event(WorkstreamEventKind::ToolCall, None, "three");
e3.occurred_at = Some("2026-09-10T12:05:00Z".to_string());
let resolved = resolve_occurred_at(&[e1, e2, e3]);
assert_eq!(
resolved.per_event,
vec![
Some(ts("2026-09-10T12:05:00Z")),
Some(ts("2026-09-10T12:05:00Z")),
Some(ts("2026-09-10T12:05:00Z")),
],
"events before the first valid timestamp inherit it backward, \
not just the ones after"
);
}
#[test]
fn resolve_occurred_at_treats_an_unparsable_timestamp_as_missing() {
let mut e1 = event(WorkstreamEventKind::Message, Some("user"), "one");
e1.occurred_at = Some("2026-09-10T12:00:00Z".to_string());
let mut e2 = event(WorkstreamEventKind::Message, Some("assistant"), "two");
e2.occurred_at = Some("not-a-timestamp".to_string());
let resolved = resolve_occurred_at(&[e1, e2]);
assert_eq!(
resolved.per_event,
vec![
Some(ts("2026-09-10T12:00:00Z")),
Some(ts("2026-09-10T12:00:00Z"))
],
"an unparsable timestamp must not reach the hook body; the \
preceding valid one is inherited instead"
);
assert_eq!(resolved.latest, Some(ts("2026-09-10T12:00:00Z")));
}
#[test]
fn resolve_occurred_at_uses_min_and_max_not_first_and_last_when_out_of_order() {
// A transcript is not guaranteed to be time-sorted (clock skew,
// reordering); the session boundary must reflect the actual extremes,
// not just the first/last entries.
let mut e1 = event(WorkstreamEventKind::Message, Some("user"), "one");
e1.occurred_at = Some("2026-09-10T12:05:00Z".to_string());
let mut e2 = event(WorkstreamEventKind::Message, Some("assistant"), "two");
e2.occurred_at = Some("2026-09-10T12:00:00Z".to_string());
let resolved = resolve_occurred_at(&[e1, e2]);
assert_eq!(
resolved.earliest,
Some(ts("2026-09-10T12:00:00Z")),
"earliest must be the minimum valid time, not the first entry"
);
assert_eq!(
resolved.latest,
Some(ts("2026-09-10T12:05:00Z")),
"latest must be the maximum valid time, not the last entry"
);
}
#[test]
fn resolve_occurred_at_stays_none_when_nothing_has_a_timestamp() {
let events = vec![
event(WorkstreamEventKind::Message, Some("user"), "one"),
event(WorkstreamEventKind::Message, Some("assistant"), "two"),
];
let resolved = resolve_occurred_at(&events);
assert_eq!(resolved.per_event, vec![None, None]);
assert_eq!(resolved.earliest, None);
assert_eq!(resolved.latest, None);
}
#[test]
fn empty_and_boundary_events_are_dropped() {
assert!(
map_event(
"sid",
&event(WorkstreamEventKind::Message, Some("user"), " ")
&event(WorkstreamEventKind::Message, Some("user"), " "),
None,
)
.is_none(),
"whitespace-only content is not an observation"
@@ -702,7 +977,7 @@ mod tests {
WorkstreamEventKind::Annotation,
] {
assert!(
map_event("sid", &event(kind, None, "boundary")).is_none(),
map_event("sid", &event(kind, None, "boundary"), None).is_none(),
"{kind:?} is not session content"
);
}
@@ -740,7 +1015,8 @@ mod tests {
});
std::fs::write(session_dir.join("foreign.jsonl"), format!("{foreign}\n")).unwrap();
let found = collect_local_sessions(home.path(), cwd.path(), None).await;
let (found, _limit_hit) =
collect_local_sessions_with(home.path(), cwd.path(), None, |_| None).await;
let claude: Vec<_> = found
.iter()
.filter(|s| s.harness == ManagedHarness::Claude)
@@ -752,10 +1028,75 @@ mod tests {
);
// `--session` narrows to one id.
let only = collect_local_sessions(home.path(), cwd.path(), Some("nope")).await;
let (only, _limit_hit) =
collect_local_sessions_with(home.path(), cwd.path(), Some("nope"), |_| None).await;
assert!(only.is_empty(), "no session matches the filter: {only:?}");
}
fn spawned_args(
server_url: Option<&str>,
server_profile: Option<&str>,
) -> crate::cli::BackfillArgs {
crate::cli::BackfillArgs {
workspace: None,
project: None,
session: None,
force: false,
dry_run: false,
max_sessions: 25,
json: false,
quiet: true,
auto: true,
server_url: server_url.map(str::to_owned),
server_profile: server_profile.map(str::to_owned),
}
}
/// #992: a SessionStart-spawned backfill presents only the credential the
/// hook itself uses for that server — never the config/env bearer, which
/// may belong to another server entirely.
#[tokio::test]
async fn a_spawned_backfill_authenticates_like_the_hook_that_spawned_it() {
let home = tempfile::tempdir().unwrap();
let data_dir = tempfile::tempdir().unwrap();
let mut config =
crate::config::Config::load(None, Some(home.path().to_path_buf())).unwrap();
config.data_dir = data_dir.path().to_path_buf();
config.auth.bearer_token = Some("CONFIG-SERVER-TOKEN".into());
let url = "https://hook.example/wiki";
let endpoint = backfill_endpoint(&config, &spawned_args(Some(url), None))
.await
.unwrap();
assert_eq!(
endpoint.auth_token, None,
"no hook token: nothing, not config's"
);
assert_eq!(endpoint.url, "https://hook.example");
assert_eq!(endpoint.base_path, "/wiki");
crate::config::store_hook_auth_token(data_dir.path(), "HOOK-TOKEN").unwrap();
let endpoint = backfill_endpoint(&config, &spawned_args(Some(url), None))
.await
.unwrap();
assert_eq!(endpoint.auth_token.as_deref(), Some("HOOK-TOKEN"));
let name = crate::server_profiles::ProfileName::parse("team-b").unwrap();
crate::server_profiles::add(data_dir.path(), &name, "https://b.example", &[], Some("B"))
.unwrap();
let endpoint = backfill_endpoint(&config, &spawned_args(None, Some("team-b")))
.await
.unwrap();
assert_eq!(endpoint.url, "https://b.example");
assert_eq!(endpoint.auth_token.as_deref(), Some("B"));
assert!(
backfill_endpoint(&config, &spawned_args(None, Some("nobody")))
.await
.is_err(),
"an unregistered profile is refused, not replaced by config"
);
}
/// The automatic path must honor the `backfill_on_start` opt-out: it records
/// the attempt (so it never re-spawns) and returns without contacting the
/// server at all.
@@ -771,18 +1112,7 @@ mod tests {
// must return before we ever build the endpoint.
config.server_url = "http://127.0.0.1:9".to_string();
let args = crate::cli::BackfillArgs {
workspace: None,
project: None,
session: None,
force: false,
dry_run: false,
max_sessions: 25,
json: false,
quiet: true,
auto: true,
};
run(&config, args)
run(&config, spawned_args(None, None))
.await
.expect("opted-out auto run must succeed without contacting the server");
crates/ai-memory-cli/src/commands/continue_session.rs
+3
View File
@@ -72,8 +72,11 @@ pub async fn run(config: &Config, args: ContinueArgs) -> Result<i32> {
new_workstream: None,
executable: None,
yolo: args.yolo,
true_yolo: args.true_yolo,
fresh: args.fresh,
no_autowire: false,
env: Vec::new(),
env_file: None,
// Bare mode: `run` resolves the harness that owns the newest
// usable session for this workstream.
harness: None,
crates/ai-memory-cli/src/commands/doctor.rs
+29 -9
View File
@@ -20,7 +20,7 @@
//! for the captured side, and the local side is read-only.
use std::collections::BTreeMap;
use std::path::Path;
use std::path::{Path, PathBuf};
use std::time::{Duration, SystemTime};
use anyhow::{Context, Result};
@@ -175,15 +175,35 @@ pub(crate) fn build_rows(
/// Read-only. A harness whose store is unreadable, absent, or unsupported
/// simply contributes nothing — the command never invents a gap it cannot see.
pub(crate) async fn scan_local(home: &Path, cwd: &Path, since_days: u32) -> Vec<LocalScan> {
scan_local_with(home, cwd, since_days, relocated_session_dir).await
}
/// Where `harness` keeps its sessions when the environment relocates its home
/// (`CLAUDE_CONFIG_DIR`, `CODEX_HOME`, `KIMI_CODE_HOME`, …), via the same
/// launch-plan resolver `ai-memory run` uses. `None` means the default
/// `$HOME`-relative store.
pub(crate) fn relocated_session_dir(harness: ManagedHarness) -> Option<PathBuf> {
build_launch_plan(harness, None, Vec::new(), None)
.ok()
.and_then(|plan| plan.session_dir)
}
/// [`scan_local`] with the relocation lookup passed in. The lookup reads the
/// process environment, so a test that plants a fixture under a temporary
/// `$HOME` passes `|_| None`: otherwise a developer's `CLAUDE_CONFIG_DIR`
/// wins over that `$HOME` and the fixture is never found.
async fn scan_local_with(
home: &Path,
cwd: &Path,
since_days: u32,
session_dir_for: impl Fn(ManagedHarness) -> Option<PathBuf>,
) -> Vec<LocalScan> {
let recent_cutoff = recent_cutoff(SystemTime::now(), since_days);
let mut scans = Vec::new();
for &harness in SCANNED_HARNESSES {
// Honor harness home relocations (CODEX_HOME, KIMI_CODE_HOME, …) via the
// same launch-plan resolver `ai-memory run` uses; fall back to the
// default $HOME-relative store when a probe plan cannot be built.
let session_dir = build_launch_plan(harness, None, Vec::new(), None)
.ok()
.and_then(|plan| plan.session_dir);
// Honor harness home relocations (see `relocated_session_dir`); fall
// back to the default $HOME-relative store when there is none.
let session_dir = session_dir_for(harness);
let Ok(sessions) =
list_native_sessions(harness, home, cwd, session_dir.as_deref(), SCAN_LIMIT).await
else {
@@ -478,7 +498,7 @@ mod tests {
});
std::fs::write(session_dir.join("foreign.jsonl"), format!("{foreign}\n")).unwrap();
let scans = scan_local(home.path(), cwd.path(), 0).await;
let scans = scan_local_with(home.path(), cwd.path(), 0, |_| None).await;
let claude = scans
.iter()
.find(|s| s.agent == AgentKind::ClaudeCode)
@@ -488,7 +508,7 @@ mod tests {
// And a project with no local stores yields no scans at all.
let empty_home = tempfile::tempdir().unwrap();
let none = scan_local(empty_home.path(), cwd.path(), 0).await;
let none = scan_local_with(empty_home.path(), cwd.path(), 0, |_| None).await;
assert!(none.is_empty(), "no stores should mean no scans: {none:?}");
}
}
crates/ai-memory-cli/src/commands/finalize_session.rs
+100 -2
View File
@@ -52,6 +52,18 @@ struct FinalizeSessionReport {
/// Returns an error if the configured server cannot list the scope's open
/// sessions or rejects a synthetic `session-end` hook.
pub async fn run(config: &Config, args: FinalizeSessionArgs) -> Result<()> {
let (workspace, project, finalized) = finalize(config, &args).await?;
let agent = args.agent;
print_report(args, workspace, project, agent, finalized)
}
/// Finalizes the sessions `args` selects (open ones, plus the exact ended
/// session with `reopen`) and returns the scope and the ids it sent a
/// session-end for, without printing; `ai-memory run` reports on its own.
pub(crate) async fn finalize(
config: &Config,
args: &FinalizeSessionArgs,
) -> Result<(String, String, Vec<String>)> {
let agent = args.agent;
let (workspace, project) =
super::resolve_scope(config, args.workspace.as_deref(), args.project.as_deref())?;
@@ -64,10 +76,11 @@ pub async fn run(config: &Config, args: FinalizeSessionArgs) -> Result<()> {
args.all,
args.all_owners,
args.session_id,
args.reopen,
)
.await?;
if sessions.is_empty() {
return print_report(args, workspace, project, agent, Vec::new());
return Ok((workspace, project, Vec::new()));
}
let client = reqwest::Client::new();
@@ -94,13 +107,16 @@ pub async fn run(config: &Config, args: FinalizeSessionArgs) -> Result<()> {
// agent session would inherit the closed id.
super::hook::clear_session_id(&config.data_dir, agent);
print_report(args, workspace, project, agent, finalized)
Ok((workspace, project, finalized))
}
/// List open sessions for the scope + agent via the server. An unknown
/// workspace/project fails closed server-side with a 404; that maps to
/// "nothing to finalize" here, matching the previous direct-DB behavior
/// for a missing scope.
// Eight arguments is the full request shape (scope + agent + the four
// selection flags + exact id); cf. `post_session_end_batch` below.
#[allow(clippy::too_many_arguments)]
async fn fetch_open_sessions(
endpoint: &ServerEndpoint,
workspace: &str,
@@ -109,15 +125,18 @@ async fn fetch_open_sessions(
all: bool,
all_owners: bool,
session_id: Option<SessionId>,
include_ended: bool,
) -> Result<Vec<OpenSessionEntry>> {
let all = if all { "true" } else { "false" };
let all_owners = if all_owners { "true" } else { "false" };
let include_ended = if include_ended { "true" } else { "false" };
let mut query = vec![
("workspace", workspace),
("project", project),
("agent", agent.as_str()),
("all", all),
("all_owners", all_owners),
("include_ended", include_ended),
];
let session_id = session_id.map(|sid| sid.to_string());
if let Some(sid) = session_id.as_deref() {
@@ -290,6 +309,7 @@ mod tests {
store
.writer
.begin_session(NewSession {
occurred_at: None,
id,
workspace_id: ws,
project_id,
@@ -344,6 +364,7 @@ mod tests {
store
.writer
.begin_session(NewSession {
occurred_at: None,
id,
workspace_id: ws,
project_id: proj,
@@ -363,6 +384,7 @@ mod tests {
AgentKind::KiroCli,
ai_memory_core::OwnerFilter::Any,
older,
false,
)
.await
.unwrap();
@@ -384,12 +406,88 @@ mod tests {
AgentKind::KiroCli,
ai_memory_core::OwnerFilter::Any,
latest,
false,
)
.await
.unwrap();
assert_eq!(still_open.map(|session| session.session_id), Some(latest));
}
/// Regression for the Antigravity manual-finalize gap: after a first
/// `finalize-session` closes the session, the conversation may continue
/// and land new observations under the same id. The default discovery
/// must keep excluding the ended session (closing something already
/// closed stays a silent no-op), while the `--reopen` lookup
/// (`include_ended = true`) must find it so a second finalize re-runs
/// the session-end path.
#[tokio::test]
async fn ended_session_matches_only_with_include_ended() {
let tmp = TempDir::new().unwrap();
let store = Store::open(tmp.path()).unwrap();
let ws = store
.writer
.get_or_create_workspace("default".to_string())
.await
.unwrap();
let proj = store
.writer
.get_or_create_project(ws, "target".to_string(), None)
.await
.unwrap();
let ended = SessionId::new();
store
.writer
.begin_session(NewSession {
id: ended,
workspace_id: ws,
project_id: proj,
agent_kind: AgentKind::AntigravityCli,
cwd: Some(std::path::PathBuf::from("/tmp/target")),
actor_user: None,
occurred_at: None,
})
.await
.unwrap();
store.writer.end_session(ended, None).await.unwrap();
let default_lookup = store
.reader
.open_session_for_scope_agent_by_id(
ws,
proj,
AgentKind::AntigravityCli,
ai_memory_core::OwnerFilter::Any,
ended,
false,
)
.await
.unwrap();
assert_eq!(
default_lookup.map(|session| session.session_id),
None,
"an ended session must stay invisible to the default finalize discovery"
);
let reopen_lookup = store
.reader
.open_session_for_scope_agent_by_id(
ws,
proj,
AgentKind::AntigravityCli,
ai_memory_core::OwnerFilter::Any,
ended,
true,
)
.await
.unwrap();
assert_eq!(
reopen_lookup.map(|session| session.session_id),
Some(ended),
"--reopen must reach the ended session for a second finalize"
);
}
#[test]
fn synthetic_end_url_propagates_all_owners_once_when_requested() {
let endpoint =
crates/ai-memory-cli/src/commands/grant.rs
+185
View File
@@ -0,0 +1,185 @@
//! Grant management for `ai-memory user grant|revoke|grants` and `ai-memory
//! project grants` (#708).
//!
//! Thin HTTP client over `/admin/users/{username}/grant|revoke|grants` and
//! `/admin/projects/grants`. The caller's bearer token must authenticate as
//! root: grants are an operator action, and the server is usually somewhere the
//! operator's laptop cannot open the database directly.
use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use crate::cli::{ProjectGrantsArgs, UserGrantArgs, UserGrantsArgs, UserRevokeArgs};
use crate::commands::user::url_encode;
use crate::http_client::{ServerEndpoint, get_json, post_json};
#[derive(Debug, Serialize)]
struct GrantRequest<'a> {
workspace: &'a str,
project: &'a str,
#[serde(skip_serializing_if = "Option::is_none")]
level: Option<&'a str>,
}
#[derive(Debug, Deserialize, Serialize)]
struct GrantRow {
username: String,
workspace: String,
project: String,
level: String,
}
#[derive(Debug, Deserialize)]
struct GrantList {
grants: Vec<GrantRow>,
}
fn print_grants(grants: &[GrantRow], json: bool) -> Result<()> {
if json {
println!("{}", serde_json::to_string_pretty(grants)?);
return Ok(());
}
if grants.is_empty() {
// Say what empty means: nothing is lost on an open project.
println!("(no grants)");
println!(
"Open projects admit every user regardless; a restricted project with no \
grants admits only root."
);
return Ok(());
}
let project_w = grants
.iter()
.map(|g| g.workspace.len() + 1 + g.project.len())
.max()
.unwrap_or(7)
.max(7);
let user_w = grants
.iter()
.map(|g| g.username.len())
.max()
.unwrap_or(8)
.max(8);
println!("{:<project_w$} {:<user_w$} LEVEL", "PROJECT", "USERNAME");
for g in grants {
let project = format!("{}/{}", g.workspace, g.project);
println!(
"{project:<project_w$} {:<user_w$} {}",
g.username, g.level
);
}
Ok(())
}
/// `ai-memory user grants [--user NAME]`.
///
/// # Errors
/// Transport failures, a non-root token, or an unknown user.
pub async fn list_for_user(ep: &ServerEndpoint, args: &UserGrantsArgs) -> Result<()> {
let resp: GrantList = match &args.user {
Some(user) => get_json(
ep,
&format!("/admin/users/{}/grants", url_encode(user)),
&[],
)
.await
.with_context(|| format!("listing {user}'s grants"))?,
None => get_json(ep, "/admin/projects/grants", &[])
.await
.context("listing grants")?,
};
print_grants(&resp.grants, args.json)
}
/// `ai-memory project grants --workspace W --project P`.
///
/// # Errors
/// Transport failures, a non-root token, or an unknown project.
pub async fn list_for_project(ep: &ServerEndpoint, args: &ProjectGrantsArgs) -> Result<()> {
let resp: GrantList = get_json(
ep,
"/admin/projects/grants",
&[("workspace", &args.workspace), ("project", &args.project)],
)
.await
.with_context(|| format!("listing grants on {}/{}", args.workspace, args.project))?;
print_grants(&resp.grants, args.json)
}
#[derive(Debug, Deserialize)]
struct GrantResponse {
level: String,
changed: bool,
previous: Option<String>,
}
/// `ai-memory user grant --user U --workspace W --project P --level L`.
///
/// # Errors
/// Transport failures, a non-root token, an unknown user, project or level.
pub async fn grant(ep: &ServerEndpoint, args: &UserGrantArgs) -> Result<()> {
let body = GrantRequest {
workspace: &args.workspace,
project: &args.project,
level: Some(&args.level),
};
let resp: GrantResponse = post_json(
ep,
&format!("/admin/users/{}/grant", url_encode(&args.user)),
&body,
)
.await
.with_context(|| {
format!(
"granting {} on {}/{}",
args.user, args.workspace, args.project
)
})?;
let who = &args.user;
let project = format!("{}/{}", args.workspace, args.project);
match (resp.changed, resp.previous) {
(false, _) => println!(
"{who} already holds {} on {project}; nothing changed.",
resp.level
),
(true, Some(previous)) => println!("{who}: {previous} -> {} on {project}.", resp.level),
(true, None) => println!("{who} now holds {} on {project}.", resp.level),
}
Ok(())
}
#[derive(Debug, Deserialize)]
struct RevokeResponse {
revoked: bool,
}
/// `ai-memory user revoke --user U --workspace W --project P`.
///
/// # Errors
/// Transport failures, a non-root token, an unknown user or project.
pub async fn revoke(ep: &ServerEndpoint, args: &UserRevokeArgs) -> Result<()> {
let body = GrantRequest {
workspace: &args.workspace,
project: &args.project,
level: None,
};
let resp: RevokeResponse = post_json(
ep,
&format!("/admin/users/{}/revoke", url_encode(&args.user)),
&body,
)
.await
.with_context(|| {
format!(
"revoking {} on {}/{}",
args.user, args.workspace, args.project
)
})?;
let project = format!("{}/{}", args.workspace, args.project);
if resp.revoked {
println!("{} no longer holds anything on {project}.", args.user);
} else {
println!("{} held nothing on {project}; nothing changed.", args.user);
}
Ok(())
}
crates/ai-memory-cli/src/commands/hook.rs
+1149 -74
View File
File diff suppressed because it is too large Load Diff
crates/ai-memory-cli/src/commands/hook_capture.rs
+199
View File
@@ -260,12 +260,14 @@ fn marker_query_suffix_impl(
let (mut workspace, mut project, mut strategy, mut drop_subagent, mut default_global) =
(None, None, None, None, None);
let (mut briefing, mut briefing_budget) = (None, None);
let mut explicit_identity = None;
// The nearest marker that declares more than `[capture]` (#668): a
// nested capture-only marker (e.g. one that only sets `ignore_paths`)
// must not shadow an outer marker's workspace/project/briefing/etc.
if let Some(marker) = find_settings_marker(cwd) {
workspace = parse_toml_key(&marker, "workspace");
project = parse_toml_key(&marker, "project");
explicit_identity = parse_toml_key(&marker, "identity");
strategy = parse_toml_key(&marker, "project_strategy");
drop_subagent = parse_toml_key(&marker, "drop_subagent_captures");
// `[recall] default_global = true` (or top-level; quoted or bare) —
@@ -281,6 +283,10 @@ fn marker_query_suffix_impl(
// tell a deliberate marker rescope from a host-derived repo-root name.
// Only the latter may yield to session-sticky attribution (#394).
let mut project_src = project.as_ref().map(|_| "marker");
// Resolved before repo-root can fill `project` below: a repo-root name is
// an inference, while the chain's `manifest` rung means a name somebody
// wrote in the marker.
let identity = repository_identity(cwd, explicit_identity.as_deref(), project.as_deref());
if strategy.is_none() {
strategy = default_strategy.map(str::to_owned);
}
@@ -300,6 +306,13 @@ fn marker_query_suffix_impl(
if let Some(val) = strategy {
qs.push_str(&format!("&project_strategy={}", url_encode(&val)));
}
if let Some(identity) = identity {
qs.push_str(&format!(
"&identity={}&identity_src={}",
url_encode(&identity.identity),
identity.source.as_str()
));
}
// Per-project `drop_subagent_captures` opt-in: forward the marker's value as
// the `drop_subagent` flag so the server scopes the drop to this project.
// The server interprets truthiness (`1`/`true`/…).
@@ -329,6 +342,38 @@ fn marker_query_suffix_impl(
qs
}
/// The repository identity to send with this checkout's events (#708), or
/// `None` to let the server route by project name as it always has.
///
/// Only the rungs that route by identity are sent: an explicit `identity`
/// from the marker, or — when the marker declares no `project` — the
/// `upstream`/`origin` remote. A declared `project` outranks the remote and
/// routes by name, so git is not consulted at all when one is present; that
/// also keeps the lookup off the hot path for every repository that declares
/// itself. The remote is normalised here, so credentials embedded in its URL
/// never leave the machine.
fn repository_identity(
cwd: &str,
explicit_identity: Option<&str>,
declared_project: Option<&str>,
) -> Option<ai_memory_core::repository_identity::RepositoryIdentity> {
use ai_memory_core::repository_identity::{IdentityInputs, resolve};
let declared = |value: Option<&str>| value.is_some_and(|v| !v.trim().is_empty());
let (upstream, origin) = if declared(explicit_identity) || declared(declared_project) {
(None, None)
} else {
ai_memory_consolidate::read_identity_remotes(std::path::Path::new(cwd))
};
resolve(&IdentityInputs {
explicit_identity,
manifest_name: declared_project,
upstream_remote: upstream.as_deref(),
origin_remote: origin.as_deref(),
folder_name: None,
})
.filter(|identity| identity.source.routes_by_identity())
}
/// Build a reqwest client for the hook's one-shot requests. `no_proxy`
/// skips Windows proxy auto-detection (registry / WinINET lookups), which
/// is pure overhead for a loopback/LAN POST. Built once per invocation and
@@ -363,6 +408,17 @@ pub enum PostOutcome {
/// so the drain can skip past them instead of stopping at the first one
/// (#493).
Unreachable,
/// `403` — the server understood the request and will never accept it.
/// Today that means the author may not write the project
/// the event belongs to.
///
/// Terminal, and that is the whole point of separating it from
/// [`Self::Failed`]. A refusal that counts as a failure gets re-sent until
/// it exhausts `MAX_ATTEMPTS`, and every one of those attempts is
/// guaranteed to be refused for the same reason. Retrying something that
/// cannot succeed is how a parse failure once cost this project 10.7M
/// tokens in a day. The entry is dropped on the spot.
Refused,
/// `401` — the server rejected the bearer. Distinguished from
/// [`Self::Failed`] because it says something about the *credential*
/// rather than the entry: a spooled event carries the token frozen at
@@ -398,6 +454,7 @@ pub async fn post_hook(
PostOutcome::Saturated
}
Ok(resp) if resp.status() == reqwest::StatusCode::UNAUTHORIZED => PostOutcome::Unauthorized,
Ok(resp) if resp.status() == reqwest::StatusCode::FORBIDDEN => PostOutcome::Refused,
Ok(_) => PostOutcome::Failed,
Err(_) => PostOutcome::Unreachable,
}
@@ -555,6 +612,19 @@ pub async fn get_handoff(
return None;
}
};
if resp.status() == reqwest::StatusCode::FORBIDDEN {
// Not authorized for this repository (#708). Same reasoning as the
// transport warning above: silence here would read as "nothing was
// handed off", when the truth is "you cannot see what was". The body
// is the server's reason and goes to stderr, never into context.
let reason = resp.text().await.unwrap_or_default();
eprintln!(
"ai-memory hook warning: not authorized for this project's memory ({}); \
nothing was injected",
reason.trim()
);
return None;
}
if !resp.status().is_success() {
return None;
}
@@ -679,6 +749,36 @@ mod tests {
assert_eq!(outcome, PostOutcome::Delivered);
}
#[tokio::test]
async fn post_hook_refused_on_403_is_terminal_not_a_failure() {
// 403 means the server will never accept this event. Classifying it as
// `Failed` would re-send it until it burnt `MAX_ATTEMPTS`, and every
// attempt would be refused identically — the shape of retry loop that
// once cost this project 10.7M tokens in a day. It must be its own
// outcome so the drain can drop it on the spot.
let url = serve_once("403 Forbidden", "capture not authorized").await;
let outcome = post_hook(&build_client(), &url, "{}", None, Duration::from_secs(1)).await;
assert_eq!(outcome, PostOutcome::Refused);
assert_ne!(
outcome,
PostOutcome::Failed,
"a refusal must never be charged a retry attempt"
);
}
#[tokio::test]
async fn get_handoff_never_injects_a_refusal() {
// A 403 carries the server's reason (#708). It is reported on stderr
// and must not reach the agent as context.
let url = serve_once(
"403 Forbidden",
"not authorized for scratch. This is an access problem, not an empty memory",
)
.await;
let got = get_handoff(&build_client(), &url, None, Duration::from_secs(1)).await;
assert!(got.is_none(), "a refusal must not become context");
}
#[tokio::test]
async fn get_handoff_ignores_non_success_status() {
let url = serve_once("401 Unauthorized", "unauthorized").await;
@@ -797,6 +897,105 @@ mod tests {
);
}
/// A throwaway repository with the given remotes, or `None` when there is
/// no `git` binary to build one with. The host is not a real forge, so a
/// developer's global `url.<base>.insteadOf` rules cannot rewrite it.
fn repo_with_remotes(remotes: &[(&str, &str)]) -> Option<tempfile::TempDir> {
let git = |args: &[&str], dir: &std::path::Path| {
std::process::Command::new("git")
.arg("-C")
.arg(dir)
.args(args)
.status()
.ok()
.filter(std::process::ExitStatus::success)
};
let tmp = tempfile::TempDir::new().unwrap();
git(&["init", "-q"], tmp.path())?;
for (name, url) in remotes {
git(&["remote", "add", name, url], tmp.path())?;
}
Some(tmp)
}
/// An undeclared checkout sends the identity of its remote — `upstream`
/// over `origin` — normalised on this side, so the credentials in the URL
/// never reach the query string.
#[test]
fn marker_query_suffix_sends_the_remote_identity_of_an_undeclared_checkout() {
let Some(repo) = repo_with_remotes(&[
(
"origin",
"https://someone:s3cret-token@git.example.test/Fork/API.git",
),
("upstream", "git@git.example.test:Acme/API.git"),
]) else {
return;
};
let qs = marker_query_suffix(repo.path().to_str().unwrap(), None);
assert!(
qs.contains("&identity=git.example.test%2Facme%2Fapi&identity_src=git_remote"),
"{qs}"
);
assert!(
!qs.contains("s3cret"),
"credentials must not leave the machine: {qs}"
);
let Some(origin_only) = repo_with_remotes(&[(
"origin",
"https://someone:s3cret-token@git.example.test/Fork/API.git",
)]) else {
return;
};
let qs = marker_query_suffix(origin_only.path().to_str().unwrap(), None);
assert!(
qs.contains("&identity=git.example.test%2Ffork%2Fapi"),
"{qs}"
);
assert!(!qs.contains("s3cret"), "{qs}");
}
/// A declared `project` outranks the remote and routes by name, so nothing
/// is sent for it; an explicit `identity` outranks both.
#[test]
fn marker_query_suffix_lets_declarations_decide_the_identity() {
let Some(repo) = repo_with_remotes(&[("upstream", "git@git.example.test:acme/tool.git")])
else {
return;
};
let cwd = repo.path().to_str().unwrap();
let marker = repo.path().join(".ai-memory.toml");
std::fs::write(&marker, "project = \"my-fork\"\n").unwrap();
let qs = marker_query_suffix(cwd, None);
assert!(qs.contains("&project=my-fork"), "{qs}");
assert!(
!qs.contains("&identity="),
"a declared project routes by name: {qs}"
);
std::fs::write(
&marker,
"project = \"my-fork\"\nidentity = \"Acme/Platform\"\n",
)
.unwrap();
let qs = marker_query_suffix(cwd, None);
assert!(
qs.contains("&identity=acme%2Fplatform&identity_src=explicit"),
"{qs}"
);
}
/// No repository, no remote, no declaration: nothing to send, and the
/// server routes by folder name exactly as before.
#[test]
fn marker_query_suffix_sends_no_identity_outside_a_repository() {
let tmp = tempfile::TempDir::new().unwrap();
let qs = marker_query_suffix(tmp.path().to_str().unwrap(), None);
assert!(!qs.contains("identity"), "{qs}");
}
#[test]
fn marker_query_suffix_appends_marker_fields() {
let tmp = tempfile::TempDir::new().unwrap();
crates/ai-memory-cli/src/commands/hook_drain_process.rs
+59 -7
View File
@@ -92,15 +92,37 @@ pub fn command_spec(data_dir: &Path, live_token: Option<&str>) -> io::Result<Dra
})
}
/// Build `ai-memory --data-dir <dir> backfill --auto --quiet`.
/// Which server the spawned backfill delivers to: the one the spawning hook
/// resolved for this event, never whatever the environment points at.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BackfillTarget<'a> {
/// The hook's install-time `--server-url`.
ServerUrl(&'a str),
/// A server profile the repository's marker selected (#992).
Profile(&'a str),
}
/// Build `ai-memory --data-dir <dir> backfill --auto --quiet` plus the hook's
/// resolved target.
///
/// The one-time boot backfill of pre-hook local history is potentially long
/// (it reads native transcripts and imports them), so it runs detached exactly
/// like the drainer rather than inline in the SessionStart hook's tight budget.
/// It carries no bearer on its argv or in its environment: `backfill` resolves
/// auth from config / the persisted hook token itself, and the automatic path
/// only runs in the loopback single-operator posture where no token is needed.
pub fn backfill_command_spec(data_dir: &Path) -> io::Result<DrainCommandSpec> {
///
/// The target is passed explicitly because, left to itself, `backfill`
/// resolves the server from config and the environment: a hook pointed at a
/// remote server then backfilled into `127.0.0.1` whenever the agent was
/// launched without `AI_MEMORY_SERVER_URL`. It carries no bearer on its argv
/// or in its environment: `backfill` reads the persisted hook token, or the
/// profile's own token, from the data dir.
pub fn backfill_command_spec(
data_dir: &Path,
target: BackfillTarget<'_>,
) -> io::Result<DrainCommandSpec> {
let (flag, value) = match target {
BackfillTarget::ServerUrl(url) => ("--server-url", url),
BackfillTarget::Profile(name) => ("--server-profile", name),
};
Ok(DrainCommandSpec {
exe: std::env::current_exe()?,
args: vec![
@@ -109,6 +131,8 @@ pub fn backfill_command_spec(data_dir: &Path) -> io::Result<DrainCommandSpec> {
OsString::from("backfill"),
OsString::from("--auto"),
OsString::from("--quiet"),
OsString::from(flag),
OsString::from(value),
],
stderr_log: data_dir.join("logs").join("backfill.log"),
live_token: None,
@@ -117,8 +141,8 @@ pub fn backfill_command_spec(data_dir: &Path) -> io::Result<DrainCommandSpec> {
/// Spawn the detached one-time boot backfill without inheriting hook stdio.
/// Best-effort: a spawn failure must never break session start.
pub fn spawn_backfill(data_dir: &Path) -> io::Result<()> {
let spec = backfill_command_spec(data_dir)?;
pub fn spawn_backfill(data_dir: &Path, target: BackfillTarget<'_>) -> io::Result<()> {
let spec = backfill_command_spec(data_dir, target)?;
spawn_spec(&spec)
}
@@ -358,6 +382,34 @@ mod tests {
);
}
/// The backfill goes where the hook resolved, and a profile travels as a
/// name only — the worker reads its token from the data dir.
#[test]
fn backfill_carries_the_hooks_target_and_no_token() {
let tmp = tempfile::tempdir().unwrap();
let default =
backfill_command_spec(tmp.path(), BackfillTarget::ServerUrl("https://a.example"))
.unwrap();
assert_eq!(
&default.args[2..],
[
"backfill",
"--auto",
"--quiet",
"--server-url",
"https://a.example"
]
.map(OsString::from)
);
let profile = backfill_command_spec(tmp.path(), BackfillTarget::Profile("team-b")).unwrap();
assert_eq!(
&profile.args[5..],
["--server-profile", "team-b"].map(OsString::from)
);
assert!(profile.live_token.is_none());
assert!(default.live_token.is_none());
}
/// An absent or empty token carries nothing, so a no-auth deployment does
/// not export an empty bearer into the child.
#[test]
crates/ai-memory-cli/src/commands/hook_spool.rs
+410 -47
View File
@@ -87,6 +87,15 @@ pub struct SpoolEntry {
/// (with `created_ms`) to drop a permanently-undeliverable event.
#[serde(default)]
pub attempts: u32,
/// Server profile this event was routed to by its repository's marker
/// (#992), `None` for the install default. Absent from the file when
/// `None`, so entries without a profile are byte-identical to before.
///
/// Marks the entry as bound to exactly one server: the drain never
/// retries it with another server's credential, and never re-points it
/// at the configured default address.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub profile: Option<String>,
}
/// `<data_dir>/hook-spool` — the spool directory.
@@ -229,36 +238,9 @@ pub fn enqueue(spool: &Path, entry: &SpoolEntry) -> std::io::Result<()> {
/// Create the spool directory `0700` on Unix so the excerpt bodies that reside
/// there (up to `MAX_AGE_MS`) — and even the timestamp+pid metadata in the file
/// names — are only reachable by the owner (#196). The spool holds private
/// capture until it drains; a world-readable directory would leak that. On
/// non-Unix the mode is a no-op (falls back to `create_dir_all`). Idempotent:
/// an existing directory's mode is left untouched (never widened, never
/// narrowed) to avoid churning a path an operator may have set deliberately.
/// capture until it drains; a world-readable directory would leak that.
fn create_spool_dir(spool: &Path) -> std::io::Result<()> {
#[cfg(unix)]
{
use std::os::unix::fs::DirBuilderExt as _;
if spool.is_dir() {
return Ok(());
}
match std::fs::DirBuilder::new()
.recursive(true)
.mode(0o700)
.create(spool)
{
Ok(()) => Ok(()),
// A concurrent drainer/enqueue may have created it between the check
// and the call; treat an existing directory as success.
Err(e) if spool.is_dir() => {
let _ = e;
Ok(())
}
Err(e) => Err(e),
}
}
#[cfg(not(unix))]
{
std::fs::create_dir_all(spool)
}
super::path_util::create_private_dir(spool)
}
fn write_private(path: &Path, bytes: &[u8]) -> std::io::Result<()> {
@@ -298,6 +280,16 @@ pub fn entry_for(
auth_mode,
token,
attempts: 0,
profile: None,
}
}
impl SpoolEntry {
/// Tag the entry with the server profile its marker routed it to (#992).
#[must_use]
pub fn routed_to(mut self, profile: Option<&str>) -> Self {
self.profile = profile.map(str::to_owned);
self
}
}
@@ -557,6 +549,7 @@ pub async fn drain_with_live_token(
let client = build_client();
let started = Instant::now();
let mut oidc_cache: Option<Option<String>> = None; // outer None = not yet resolved
let mut profile_tokens = ProfileTokens::new(data_dir);
let mut result = DrainResult::default();
let mut idx = 0;
@@ -622,7 +615,12 @@ pub async fn drain_with_live_token(
bump_or_drop(&next_path, &next_entry, &mut result);
continue;
}
if batch_endpoint(&next_entry.url) != base {
// A profile entry never shares a request with another route,
// even at the same address: the retry and reroot decisions
// below are made once per chunk, from its first entry.
if batch_endpoint(&next_entry.url) != base
|| next_entry.profile != chunk[0].1.profile
{
idx -= 1;
break;
}
@@ -738,6 +736,13 @@ pub async fn drain_with_live_token(
PostOutcome::Failed => {
bump_or_drop(path, entry, &mut result);
}
PostOutcome::Refused => {
// Never retried: see `PostOutcome::Refused`.
// Dropped rather than charged an attempt, so
// it cannot sit in the spool being re-sent.
let _ = std::fs::remove_file(path);
result.dropped += 1;
}
PostOutcome::Unreachable => {
// Same reasoning as the batch arm: the address
// is dead, not the entry. Charge it once, then
@@ -748,15 +753,19 @@ pub async fn drain_with_live_token(
PostOutcome::Unauthorized => {
// Credential rejected, not the entry. Retry once
// with this drain's live token (#542).
let retry =
static_retry_token(entry, item_bearer.as_deref(), live_token);
let retry = static_retry_token(
entry,
item_bearer.as_deref(),
live_token,
&mut profile_tokens,
);
let recovered = match retry {
Some(token) => matches!(
post_hook(
&client,
&entry.url,
&entry.body,
Some(token),
Some(&token),
per_event_timeout,
)
.await,
@@ -793,8 +802,12 @@ pub async fn drain_with_live_token(
let configured = configured_server
.get_or_insert_with(|| configured_server_url(data_dir))
.clone();
// A profile entry is left alone too: `config.toml`'s
// address is the install default, and re-pointing a
// profile's capture there is the cross-server delivery
// profiles exist to prevent (#992).
let retry = configured.as_deref().and_then(|server| {
if !is_loopback_url(&chunk[0].1.url) {
if chunk[0].1.profile.is_some() || !is_loopback_url(&chunk[0].1.url) {
return None;
}
let rerooted = batch_endpoint(&reroot_url(&chunk[0].1.url, server)?);
@@ -850,9 +863,15 @@ pub async fn drain_with_live_token(
// current by construction. Runs only for a batch that has
// already been rejected, so a healthy drain never pays for
// it.
let retry = static_retry_token(&chunk[0].1, bearer.as_deref(), live_token);
let retry = static_retry_token(
&chunk[0].1,
bearer.as_deref(),
live_token,
&mut profile_tokens,
);
if let Some(token) = retry {
match post_batch(&client, &base, &payload, Some(token), batch_timeout).await
match post_batch(&client, &base, &payload, Some(&token), batch_timeout)
.await
{
BatchOutcome::Accepted(k) => {
let k = k.min(chunk.len());
@@ -916,6 +935,11 @@ pub async fn drain_with_live_token(
PostOutcome::Saturated => {
result.remaining += 1;
}
PostOutcome::Refused => {
// Terminal; see `PostOutcome::Refused`.
let _ = std::fs::remove_file(path);
result.dropped += 1;
}
PostOutcome::Failed => {
bump_or_drop(&path, &entry, &mut result);
}
@@ -926,14 +950,19 @@ pub async fn drain_with_live_token(
PostOutcome::Unauthorized => {
// Credential rejected, not the entry. Retry once with this
// drain's live token (#542).
let retry = static_retry_token(&entry, bearer.as_deref(), live_token);
let retry = static_retry_token(
&entry,
bearer.as_deref(),
live_token,
&mut profile_tokens,
);
let recovered = match retry {
Some(token) => matches!(
post_hook(
&client,
&entry.url,
&entry.body,
Some(token),
Some(&token),
per_event_timeout,
)
.await,
@@ -1073,15 +1102,61 @@ fn live_static_token() -> Option<String> {
/// refreshed once per pass, so a 401 there is a real rejection — and the live
/// token must actually differ from the one that just failed. Retrying an
/// identical credential would only repeat the failure at double the cost.
fn static_retry_token<'a>(
///
/// A profile-routed entry (#992) never sees the live token, which belongs to
/// whichever install-default hook spawned this drain: presenting it to the
/// entry's server would hand one server's credential to another. It retries
/// with its own profile's current token instead, re-read from the data dir,
/// and a profile that has since been removed means no retry at all.
fn static_retry_token(
entry: &SpoolEntry,
rejected: Option<&str>,
live: Option<&'a str>,
) -> Option<&'a str> {
live: Option<&str>,
profile_tokens: &mut ProfileTokens<'_>,
) -> Option<String> {
if entry.auth_mode != AuthMode::Static {
return None;
}
live.filter(|t| Some(*t) != rejected)
let current = match entry.profile.as_deref() {
Some(raw) => profile_tokens.current(raw, &entry.url)?,
None => live?.to_owned(),
};
(Some(current.as_str()) != rejected).then_some(current)
}
/// Each profile's current URL and token, read at most once per drain pass so
/// a backlog of rejected entries costs one registry read per profile, not per
/// entry.
struct ProfileTokens<'a> {
data_dir: &'a Path,
cache: std::collections::HashMap<String, Option<crate::server_profiles::ResolvedServer>>,
}
impl<'a> ProfileTokens<'a> {
fn new(data_dir: &'a Path) -> Self {
Self {
data_dir,
cache: std::collections::HashMap::new(),
}
}
/// The profile's current token, only while the profile still points at
/// the server `entry_url` was captured against. After a URL change the
/// current token belongs to the *new* server, and replaying it to the old
/// address would hand it to a server it was never issued for.
fn current(&mut self, raw: &str, entry_url: &str) -> Option<String> {
let data_dir = self.data_dir;
let resolved = self
.cache
.entry(raw.to_owned())
.or_insert_with(|| {
let name = crate::server_profiles::ProfileName::parse(raw)?;
crate::server_profiles::lookup(data_dir, &name).ok()
})
.as_ref()?;
let rest = entry_url.strip_prefix(resolved.url.as_str())?;
rest.starts_with('/').then(|| resolved.token.clone())
}
}
/// Is this URL's host a loopback address?
@@ -1934,6 +2009,91 @@ mod tests {
);
}
/// #992, adversarial: the same shape as the test above, but the entry
/// was routed to a profile. `config.toml`'s address is the install
/// default, so the reroot must not fire — the live default server gets
/// nothing even though it would accept the batch.
#[tokio::test]
async fn a_profile_entry_on_a_dead_loopback_port_is_not_rerouted_to_the_default() {
let default_hits = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
let live = serve_counting_hook(default_hits.clone(), "200 OK").await;
let dead = {
let l = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let a = l.local_addr().unwrap();
drop(l);
a
};
let tmp = tempfile::tempdir().unwrap();
std::fs::write(
tmp.path().join("config.toml"),
format!("server_url = \"http://{live}\"\n"),
)
.unwrap();
let spool = spool_dir(tmp.path());
enqueue(
&spool,
&entry_for(
format!("http://{dead}/hook?event=x"),
"{}".into(),
Some("b-token"),
false,
)
.routed_to(Some("team-b")),
)
.unwrap();
let r = drain(
&spool,
tmp.path(),
Duration::from_secs(5),
Duration::from_millis(500),
)
.await;
assert_eq!(r.sent, 0);
assert_eq!(
default_hits.load(std::sync::atomic::Ordering::SeqCst),
0,
"the install-default server must never receive a profile's capture"
);
assert_eq!(list_entries(&spool).unwrap().0.len(), 1);
}
/// A profile entry and an install-default entry never share a batch, even
/// when they share an address and a bearer: the retry and reroot
/// decisions are taken once per chunk.
#[tokio::test]
async fn profile_and_default_entries_at_one_address_ride_separate_batches() {
let hits = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
let addr = serve_counting_hook(hits.clone(), "200 OK").await;
let tmp = tempfile::tempdir().unwrap();
let spool = spool_dir(tmp.path());
for profile in [None, Some("team-b"), None] {
enqueue(
&spool,
&entry_for(
format!("http://{addr}/hook?event=x"),
"{}".into(),
Some("same"),
false,
)
.routed_to(profile),
)
.unwrap();
}
let r = drain(
&spool,
tmp.path(),
Duration::from_secs(5),
Duration::from_millis(500),
)
.await;
assert_eq!(r.sent, 3);
assert_eq!(hits.load(std::sync::atomic::Ordering::SeqCst), 3);
}
/// The retry must not re-point a spool captured against another host on
/// purpose. Only a loopback authority is unambiguous.
#[tokio::test]
@@ -2237,43 +2397,246 @@ mod tests {
assert_eq!(list_entries(&spool).unwrap().0.len(), 1, "entry survives");
}
/// #992, adversarial: server B accepts exactly the install default's live
/// token. Had the drain presented it to B, the entry would deliver; it
/// must stay queued instead, because that token belongs to another server.
#[tokio::test]
async fn a_profile_entry_is_never_retried_with_the_install_live_token() {
let tmp = tempfile::tempdir().unwrap();
let spool = spool_dir(tmp.path());
let count = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
let server_b = serve_token_gated_hook("A-LIVE", count.clone()).await;
let entry = entry_for(
format!("http://{server_b}/hook?event=e0"),
"{}".into(),
Some("b-old"),
false,
)
.routed_to(Some("team-b"));
enqueue(&spool, &entry).unwrap();
let r = drain_with_live_token(
&spool,
tmp.path(),
Duration::from_secs(5),
Duration::from_millis(500),
Some("A-LIVE"),
)
.await;
assert_eq!(r.sent, 0, "server B must never see the install's token");
assert_eq!(count.load(std::sync::atomic::Ordering::SeqCst), 1);
assert_eq!(list_entries(&spool).unwrap().0.len(), 1, "entry survives");
}
/// The control: a rotated *profile* token is recovered, from the profile
/// store rather than the drain's environment.
#[tokio::test]
async fn a_profile_entry_recovers_with_its_own_rotated_token() {
let tmp = tempfile::tempdir().unwrap();
let spool = spool_dir(tmp.path());
let count = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
let server_b = serve_token_gated_hook("b-new", count.clone()).await;
let name = crate::server_profiles::ProfileName::parse("team-b").unwrap();
crate::server_profiles::add(
tmp.path(),
&name,
&format!("http://{server_b}"),
&[],
Some("b-new"),
)
.unwrap();
let entry = entry_for(
format!("http://{server_b}/hook?event=e0"),
"{}".into(),
Some("b-old"),
false,
)
.routed_to(Some("team-b"));
enqueue(&spool, &entry).unwrap();
let r = drain_with_live_token(
&spool,
tmp.path(),
Duration::from_secs(5),
Duration::from_millis(500),
Some("A-LIVE"),
)
.await;
assert_eq!(r.sent, 1);
assert_eq!(list_entries(&spool).unwrap().0.len(), 0);
}
/// An identical token is not worth a second request, and a non-static
/// entry is not this mechanism's business: an OIDC bearer is already
/// resolved and refreshed once per pass, so a 401 there is a real
/// rejection rather than a stale freeze.
#[test]
fn static_retry_token_only_fires_when_it_can_change_the_outcome() {
let dd = tempfile::tempdir().unwrap();
let static_entry = entry_for("http://x/hook".into(), "{}".into(), Some("old"), false);
let oidc_entry = entry_for("http://x/hook".into(), "{}".into(), None, true);
let anon_entry = entry_for("http://x/hook".into(), "{}".into(), None, false);
assert_eq!(
static_retry_token(&static_entry, Some("old"), Some("new")),
static_retry_token(
&static_entry,
Some("old"),
Some("new"),
&mut ProfileTokens::new(dd.path())
)
.as_deref(),
Some("new"),
"a rotated token is worth one retry"
);
assert_eq!(
static_retry_token(&static_entry, Some("old"), Some("old")),
static_retry_token(
&static_entry,
Some("old"),
Some("old"),
&mut ProfileTokens::new(dd.path())
),
None,
"an identical token would only repeat the failure"
);
assert_eq!(
static_retry_token(&static_entry, Some("old"), None),
static_retry_token(
&static_entry,
Some("old"),
None,
&mut ProfileTokens::new(dd.path())
),
None,
"no live token, nothing to retry with"
);
assert_eq!(
static_retry_token(&oidc_entry, Some("resolved"), Some("new")),
static_retry_token(
&oidc_entry,
Some("resolved"),
Some("new"),
&mut ProfileTokens::new(dd.path())
),
None,
"OIDC is already re-resolved per pass; a 401 there is genuine"
);
assert_eq!(
static_retry_token(&anon_entry, None, Some("new")),
static_retry_token(
&anon_entry,
None,
Some("new"),
&mut ProfileTokens::new(dd.path())
),
None,
"an anonymous entry was never authenticated"
);
}
/// #992: a profile entry retries only with its own profile's current
/// token. The live token belongs to the install default; handing it to a
/// profile's server would leak one server's credential to another.
#[test]
fn a_profile_entry_never_retries_with_the_live_token() {
let dd = tempfile::tempdir().unwrap();
let entry = entry_for("http://b/hook".into(), "{}".into(), Some("b-old"), false)
.routed_to(Some("team-b"));
assert_eq!(
static_retry_token(
&entry,
Some("b-old"),
Some("A-LIVE"),
&mut ProfileTokens::new(dd.path())
),
None,
"an unregistered profile has nothing to retry with"
);
let name = crate::server_profiles::ProfileName::parse("team-b").unwrap();
crate::server_profiles::add(dd.path(), &name, "http://b", &[], Some("b-new")).unwrap();
assert_eq!(
static_retry_token(
&entry,
Some("b-old"),
Some("A-LIVE"),
&mut ProfileTokens::new(dd.path())
)
.as_deref(),
Some("b-new"),
"a rotated profile token is re-read from the store"
);
assert_eq!(
static_retry_token(
&entry,
Some("b-new"),
Some("A-LIVE"),
&mut ProfileTokens::new(dd.path())
),
None,
"the profile's current token already failed"
);
}
/// After the profile moves to another server, its current token belongs
/// to that server: an entry still addressed to the old one must not be
/// retried with it.
#[test]
fn a_profile_entry_is_not_retried_with_a_token_issued_for_a_new_url() {
let dd = tempfile::tempdir().unwrap();
let name = crate::server_profiles::ProfileName::parse("team-b").unwrap();
crate::server_profiles::add(dd.path(), &name, "https://new.example", &[], Some("NEW"))
.unwrap();
let stale = entry_for(
"https://old.example/hook?event=e".into(),
"{}".into(),
Some("OLD"),
false,
)
.routed_to(Some("team-b"));
let lookalike = entry_for(
"https://new.example.evil/hook?event=e".into(),
"{}".into(),
Some("OLD"),
false,
)
.routed_to(Some("team-b"));
let current = entry_for(
"https://new.example/hook?event=e".into(),
"{}".into(),
Some("OLD"),
false,
)
.routed_to(Some("team-b"));
let mut tokens = ProfileTokens::new(dd.path());
assert_eq!(
static_retry_token(&stale, Some("OLD"), None, &mut tokens),
None
);
assert_eq!(
static_retry_token(&lookalike, Some("OLD"), None, &mut tokens),
None
);
assert_eq!(
static_retry_token(&current, Some("OLD"), None, &mut tokens).as_deref(),
Some("NEW"),
"control: the same server still gets its rotated token"
);
}
/// Entries without a profile serialize exactly as they did before #992,
/// so an older binary draining a mixed spool reads them unchanged.
#[test]
fn an_install_default_entry_serializes_without_a_profile_field() {
let entry = entry_for("http://a/hook".into(), "{}".into(), Some("t"), false);
let json = serde_json::to_value(&entry).unwrap();
assert!(json.get("profile").is_none(), "{json}");
let routed = serde_json::to_value(entry.routed_to(Some("team-b"))).unwrap();
assert_eq!(routed["profile"], "team-b");
}
/// A `401` must be distinguishable from any other refusal, or the drain
/// cannot tell "this credential is stale" from "this event is bad".
#[tokio::test]
crates/ai-memory-cli/src/commands/install_hooks.rs
+1429 -345
View File
File diff suppressed because it is too large Load Diff
crates/ai-memory-cli/src/commands/install_mcp.rs
+241 -20
View File
@@ -11,8 +11,9 @@
//! community-standard `npx mcp-remote` stdio shim so the same HTTP
//! endpoint still works.
//!
//! OMP uses a native `~/.omp/agent/mcp.json` file with the same
//! `mcpServers` root as several other clients.
//! OMP uses a native `mcp.json` in its agent dir (`~/.omp/agent`, a named
//! profile's `~/.omp/profiles/<name>/agent`, or `$PI_CODING_AGENT_DIR`) with
//! the same `mcpServers` root as several other clients.
use std::path::{Path, PathBuf};
@@ -160,14 +161,27 @@ fn validate_kiro_remote_url(server_url: &str) -> Result<()> {
/// Returns an error for `Pi` (no MCP config), for Claude Desktop on
/// unsupported OSes, or when `$HOME` can't be resolved.
pub(crate) fn mcp_config_path(client: crate::cli::McpClient) -> Result<PathBuf> {
mcp_config_path_with(client, &|name| std::env::var_os(name))
}
/// [`mcp_config_path`] with the relocation variables (`CLAUDE_CONFIG_DIR`,
/// `CODEX_HOME`, `GROK_HOME`, `KIMI_CODE_HOME`, `KIRO_HOME`, and OMP's
/// `OMP_PROFILE`, `PI_PROFILE`, `PI_CODING_AGENT_DIR` and `PI_CONFIG_DIR`) read
/// through `env`,
/// so `ai-memory run --env` can point auto-wire at the same config home
/// it launches the harness with.
pub(crate) fn mcp_config_path_with(
client: crate::cli::McpClient,
env: &dyn Fn(&str) -> Option<std::ffi::OsString>,
) -> Result<PathBuf> {
use crate::cli::McpClient;
let home = || home_dir().context("could not locate $HOME for config-file auto-detect");
Ok(match client {
McpClient::ClaudeCode => claude_code_config_path_in(std::env::var_os("CLAUDE_CONFIG_DIR"))?,
McpClient::Codex => home()?.join(".codex").join("config.toml"),
McpClient::ClaudeCode => claude_code_config_path_in(env("CLAUDE_CONFIG_DIR"))?,
McpClient::Codex => codex_config_path_in(env("CODEX_HOME"))?,
// Project scope is `.grok/config.toml` under cwd/repo; pass
// --config-file for that case rather than inventing a second default.
McpClient::Grok => grok_home()?.join("config.toml"),
McpClient::Grok => grok_home_in(env("GROK_HOME"))?.join("config.toml"),
McpClient::OpenCode => home()?
.join(".config")
.join("opencode")
@@ -212,7 +226,11 @@ pub(crate) fn mcp_config_path(client: crate::cli::McpClient) -> Result<PathBuf>
McpClient::Pi => bail!(
"Pi has no native mcp.json; use `ai-memory install-hooks --agent pi --apply` to install the generated MCP bridge extension."
),
McpClient::Omp => home()?.join(".omp").join("agent").join("mcp.json"),
// OMP reads mcp.json from its agent dir, the one its extensions live
// in, so a profile, PI_CODING_AGENT_DIR or PI_CONFIG_DIR moves it too.
McpClient::Omp => {
ai_memory_workstream::omp_agent_dir(&home()?, None, env)?.join("mcp.json")
}
McpClient::AntigravityCli => home()?
.join(".gemini")
.join("config")
@@ -229,8 +247,8 @@ pub(crate) fn mcp_config_path(client: crate::cli::McpClient) -> Result<PathBuf>
// Kimi Code keeps its data dir at $KIMI_CODE_HOME when set,
// falling back to ~/.kimi-code; MCP servers live in mcp.json at
// that root.
McpClient::KimiCode => kimi_code_home(std::env::var_os("KIMI_CODE_HOME"))?.join("mcp.json"),
McpClient::KiroCli => kiro_home(std::env::var_os("KIRO_HOME"))?
McpClient::KimiCode => kimi_code_home(env("KIMI_CODE_HOME"))?.join("mcp.json"),
McpClient::KiroCli => kiro_home(env("KIRO_HOME"))?
.join("settings")
.join("mcp.json"),
McpClient::CommandCode => home()?.join(".commandcode").join("mcp.json"),
@@ -386,23 +404,41 @@ fn claude_code_config_path_in(env_override: Option<std::ffi::OsString>) -> Resul
.join(".claude.json"))
}
/// Kimi Code's data dir: `$KIMI_CODE_HOME` when set (non-empty), else
/// Codex's user config, where its MCP servers live: `$CODEX_HOME/config.toml`
/// when the var is set, else `~/.codex/config.toml`. Codex keeps its whole
/// config home under `CODEX_HOME`, which `install-hooks` already honors for
/// `hooks.json`; writing the MCP entry to the default path left a relocated
/// Codex with hooks but no ai-memory server. Blank counts as unset, as it does
/// for the hooks path. The env value comes in as a parameter so tests can
/// exercise both branches without mutating process env.
fn codex_config_path_in(env_override: Option<std::ffi::OsString>) -> Result<PathBuf> {
if let Some(dir) = crate::commands::path_util::agent_config_home(env_override) {
return Ok(dir.join("config.toml"));
}
Ok(home_dir()
.context("could not locate $HOME for ~/.codex/config.toml")?
.join(".codex")
.join("config.toml"))
}
/// Kimi Code's data dir: `$KIMI_CODE_HOME` when set and not blank, else
/// `~/.kimi-code`. The env value comes in as a parameter so tests can
/// exercise both branches without mutating process env.
fn kimi_code_home(env_override: Option<std::ffi::OsString>) -> Result<PathBuf> {
if let Some(dir) = env_override.filter(|value| !value.is_empty()) {
return Ok(PathBuf::from(dir));
if let Some(dir) = crate::commands::path_util::agent_config_home(env_override) {
return Ok(dir);
}
Ok(home_dir()
.context("could not locate $HOME for config-file auto-detect")?
.join(".kimi-code"))
}
/// Kiro CLI's global configuration root: `$KIRO_HOME` when set, otherwise
/// `~/.kiro`. The override is injected to keep path tests process-local.
/// Kiro CLI's global configuration root: `$KIRO_HOME` when set and not
/// blank, otherwise `~/.kiro`. The override is injected to keep path tests
/// process-local.
fn kiro_home(env_override: Option<std::ffi::OsString>) -> Result<PathBuf> {
if let Some(dir) = env_override.filter(|value| !value.is_empty()) {
return Ok(PathBuf::from(dir));
if let Some(dir) = crate::commands::path_util::agent_config_home(env_override) {
return Ok(dir);
}
Ok(home_dir()
.context("could not locate $HOME for Kiro configuration")?
@@ -412,8 +448,14 @@ fn kiro_home(env_override: Option<std::ffi::OsString>) -> Result<PathBuf> {
/// Resolve Grok Build CLI's user configuration root. Grok honours
/// `GROK_HOME`; otherwise it uses `~/.grok`.
pub(crate) fn grok_home() -> Result<PathBuf> {
if let Some(path) = std::env::var_os("GROK_HOME").filter(|path| !path.is_empty()) {
return Ok(PathBuf::from(path));
grok_home_in(std::env::var_os("GROK_HOME"))
}
/// [`grok_home`] with the `GROK_HOME` value passed in, so auto-wire can resolve
/// it from `ai-memory run --env` and tests stay process-local.
pub(crate) fn grok_home_in(env_override: Option<std::ffi::OsString>) -> Result<PathBuf> {
if let Some(dir) = crate::commands::path_util::agent_config_home(env_override) {
return Ok(dir);
}
Ok(home_dir()
.context("could not locate $HOME for Grok configuration")?
@@ -430,6 +472,47 @@ fn resolve_config_file(args: &InstallMcpArgs) -> Result<PathBuf> {
mcp_config_path(args.client)
}
/// True when the client's MCP config already holds a session-aware ai-memory
/// bridge under `name` — a Claude Code stdio entry whose args run `mcp-bridge`.
///
/// Auto-wire consults this before its version-keyed re-wire so it never
/// downgrades a deliberately-installed bridge to the static HTTP registration,
/// which would silently disable `[auto_scope] per_session` for MCP calls. Only
/// Claude Code has a session-aware variant, so every other client answers
/// `false`. A missing file or any parse error also answers `false`: the caller
/// then falls back to its normal install, which is the safe default.
pub(crate) fn existing_entry_is_session_aware(
client: McpClient,
config_file: Option<&Path>,
name: &str,
) -> bool {
if !matches!(client, McpClient::ClaudeCode) {
return false;
}
let path = match config_file {
Some(path) => path.to_path_buf(),
None => match mcp_config_path(client) {
Ok(path) => path,
Err(_) => return false,
},
};
let Ok(text) = std::fs::read_to_string(&path) else {
return false;
};
let Ok(root) = serde_json::from_str::<serde_json::Value>(&text) else {
return false;
};
let Some(entry) = root.get("mcpServers").and_then(|servers| servers.get(name)) else {
return false;
};
let is_stdio = entry.get("type").and_then(serde_json::Value::as_str) == Some("stdio");
let runs_bridge = entry
.get("args")
.and_then(serde_json::Value::as_array)
.is_some_and(|args| args.iter().any(|arg| arg.as_str() == Some("mcp-bridge")));
is_stdio && runs_bridge
}
/// Mutate the resolved client config file in place. Idempotent —
/// re-runs that produce the same content are reported as no-op.
fn apply_to_config_file(args: &InstallMcpArgs) -> Result<()> {
@@ -1113,7 +1196,7 @@ fn render_codex(args: &InstallMcpArgs) -> String {
// and NOT a `[mcp_servers.<name>.headers]` sub-table (the key
// is `http_headers`, with the `http_` prefix).
let mut out = format!(
"# Codex CLI — append to ~/.codex/config.toml\n\
"# Codex CLI — append to $CODEX_HOME/config.toml (default ~/.codex/config.toml)\n\
#\n\
[mcp_servers.{name}]\n\
url = \"{url}\"\n\
@@ -1335,7 +1418,10 @@ fn hook_server_url_from_mcp_url(url: &str) -> String {
fn render_omp(args: &InstallMcpArgs) -> Result<String> {
Ok(format!(
"# Oh My Pi / OMP — merge into ~/.omp/agent/mcp.json:\n\
"# Oh My Pi / OMP — merge into mcp.json in OMP's agent dir:\n\
# ~/.omp/agent/mcp.json by default, ~/.omp/profiles/<name>/agent/mcp.json\n\
# under a named profile (OMP_PROFILE), or $PI_CODING_AGENT_DIR/mcp.json\n\
# when that variable relocates the default profile.\n\
#\n\
# The current Oh My Pi package exposes the `omp` binary and native\n\
# `.omp` config directories. Restart `omp` after changing MCP config.\n\
@@ -1622,6 +1708,134 @@ mod tests {
}
}
#[test]
fn codex_config_path_honours_codex_home() {
let custom = if cfg!(windows) {
r"C:\custom\codex"
} else {
"/custom/codex"
};
let path = codex_config_path_in(Some(std::ffi::OsString::from(custom))).unwrap();
assert_eq!(path, std::path::Path::new(custom).join("config.toml"));
// Unset, empty and blank all fall back to ~/.codex/config.toml, the
// same reading `install-hooks` gives CODEX_HOME for hooks.json.
for env in [
None,
Some(std::ffi::OsString::new()),
Some(std::ffi::OsString::from(" ")),
] {
let path = codex_config_path_in(env).unwrap();
assert!(
path.ends_with(std::path::Path::new(".codex").join("config.toml")),
"default must be ~/.codex/config.toml, got {}",
path.display()
);
}
}
/// Auto-wire resolves MCP targets through `run --env`; every relocatable
/// client must read its variable from the supplied lookup, not the process.
#[test]
fn mcp_config_path_with_reads_relocation_from_the_supplied_env() {
let root = if cfg!(windows) {
std::path::Path::new(r"C:\relocated")
} else {
std::path::Path::new("/relocated")
};
for (client, var, relative) in [
(McpClient::ClaudeCode, "CLAUDE_CONFIG_DIR", ".claude.json"),
(McpClient::Codex, "CODEX_HOME", "config.toml"),
(McpClient::Grok, "GROK_HOME", "config.toml"),
(McpClient::KimiCode, "KIMI_CODE_HOME", "mcp.json"),
(McpClient::KiroCli, "KIRO_HOME", "settings/mcp.json"),
(McpClient::Omp, "PI_CODING_AGENT_DIR", "mcp.json"),
] {
let env = |name: &str| (name == var).then(|| root.as_os_str().to_owned());
assert_eq!(
mcp_config_path_with(client, &env).unwrap(),
root.join(relative),
"{client:?} must follow {var}"
);
}
}
/// OMP's MCP file lives in the same agent dir as its extensions: a named
/// profile owns it and ignores `PI_CODING_AGENT_DIR`, as OMP does.
#[test]
fn omp_mcp_config_follows_the_profile_agent_dir() {
let home = home_dir().unwrap();
let profile = home
.join(".omp")
.join("profiles")
.join("work")
.join("agent")
.join("mcp.json");
for pairs in [
&[("OMP_PROFILE", "work")][..],
&[("OMP_PROFILE", "work"), ("PI_CODING_AGENT_DIR", "/custom")][..],
&[("PI_PROFILE", "work")][..],
] {
let env = |name: &str| {
pairs
.iter()
.find(|(key, _)| *key == name)
.map(|(_, value)| std::ffi::OsString::from(value))
};
assert_eq!(
mcp_config_path_with(McpClient::Omp, &env).unwrap(),
profile,
"{pairs:?}"
);
}
let invalid = |name: &str| (name == "OMP_PROFILE").then(|| "Work".into());
assert!(mcp_config_path_with(McpClient::Omp, &invalid).is_err());
let renamed = |name: &str| (name == "PI_CONFIG_DIR").then(|| ".omp-alt".into());
assert_eq!(
mcp_config_path_with(McpClient::Omp, &renamed).unwrap(),
home.join(".omp-alt").join("agent").join("mcp.json"),
"PI_CONFIG_DIR renames OMP's root"
);
}
/// A blank relocation variable is unset, the rule every installer and
/// native session import share; otherwise the MCP entry, bearer token
/// included, landed in a whitespace-named directory under the cwd.
#[test]
fn mcp_config_path_with_treats_blank_relocation_as_unset() {
let home = home_dir().unwrap();
for (client, var, relative) in [
(McpClient::ClaudeCode, "CLAUDE_CONFIG_DIR", ".claude.json"),
(McpClient::Codex, "CODEX_HOME", ".codex/config.toml"),
(McpClient::Grok, "GROK_HOME", ".grok/config.toml"),
(McpClient::KimiCode, "KIMI_CODE_HOME", ".kimi-code/mcp.json"),
(McpClient::KiroCli, "KIRO_HOME", ".kiro/settings/mcp.json"),
(McpClient::Omp, "PI_CODING_AGENT_DIR", ".omp/agent/mcp.json"),
(McpClient::Omp, "OMP_PROFILE", ".omp/agent/mcp.json"),
(McpClient::Omp, "PI_CONFIG_DIR", ".omp/agent/mcp.json"),
] {
for blank in ["", " ", "\t", " \n"] {
let env = |name: &str| (name == var).then(|| std::ffi::OsString::from(blank));
assert_eq!(
mcp_config_path_with(client, &env).unwrap(),
home.join(relative),
"{client:?} with {var}={blank:?}"
);
}
}
}
/// The snippet names CODEX_HOME rather than a resolved path: it is often
/// rendered inside the Docker image, whose paths mean nothing on the host.
#[test]
fn codex_render_names_codex_home() {
let out = render_codex(&args_for(McpClient::Codex));
assert!(
out.contains("$CODEX_HOME/config.toml"),
"render must point a relocated Codex at its config home:\n{out}"
);
}
#[test]
fn zed_config_path_uses_platform_conventions() {
for (target_os, root, expected) in [
@@ -2188,7 +2402,14 @@ mod tests {
let grok_token = render_with_token(McpClient::Grok);
assert!(grok_token.contains("[mcp_servers.ai-memory.headers]"));
assert!(!grok_token.contains("http_headers"));
assert!(render_for_test(McpClient::Omp).contains("~/.omp/agent/mcp.json"));
let omp = render_for_test(McpClient::Omp);
for location in [
"~/.omp/agent/mcp.json",
"~/.omp/profiles/<name>/agent/mcp.json",
"$PI_CODING_AGENT_DIR/mcp.json",
] {
assert!(omp.contains(location), "{location} missing from:\n{omp}");
}
let pi = render_pi(&args_for(McpClient::Pi)).unwrap();
assert!(pi.contains("Pi has no native mcp.json"));
assert!(pi.contains("install-hooks --agent pi --apply"));
crates/ai-memory-cli/src/commands/install_skills.rs
+32 -2
View File
@@ -4,8 +4,8 @@ use std::fs;
use std::path::{Path, PathBuf};
use ai_memory_core::routing_skills::{
AGENTS_SKILL_DIR, CLAUDE_SKILL_DIR, DEVIN_SKILL_DIR, GROK_SKILL_DIR, MANAGED_MARKER,
MANAGED_SKILLS, ManagedSkill, SKILLS_DIR,
AGENTS_SKILL_DIR, CLAUDE_SKILL_DIR, DEVIN_SKILL_DIR, GROK_SKILL_DIR, HERMES_SKILL_DIR,
MANAGED_MARKER, MANAGED_SKILLS, ManagedSkill, SKILLS_DIR,
};
use anyhow::{Context, Result, bail};
@@ -176,6 +176,18 @@ fn resolve_target_roots_for_platform(
platform,
)?]
}
InstallSkillsAgent::Hermes => {
vec![agent_root(
args.scope,
SkillRootKind::Hermes,
cwd,
home,
appdata,
grok_home,
claude_config_dir,
platform,
)?]
}
InstallSkillsAgent::Both => vec![
agent_root(
args.scope,
@@ -225,6 +237,7 @@ enum SkillRootKind {
Agents,
Devin,
Grok,
Hermes,
}
#[allow(clippy::too_many_arguments)]
@@ -275,6 +288,7 @@ fn agent_root(
SkillRootKind::Agents => AGENTS_SKILL_DIR,
SkillRootKind::Devin => DEVIN_SKILL_DIR,
SkillRootKind::Grok => GROK_SKILL_DIR,
SkillRootKind::Hermes => HERMES_SKILL_DIR,
};
Ok(base.join(agent_dir).join(SKILLS_DIR))
}
@@ -420,6 +434,22 @@ mod tests {
.unwrap();
assert_eq!(root_names(&project_grok), ["/repo/.grok/skills"]);
let project_hermes = resolve_target_roots(
&args(InstallSkillsScope::Project, InstallSkillsAgent::Hermes),
cwd,
Some(home),
)
.unwrap();
assert_eq!(root_names(&project_hermes), ["/repo/.hermes/skills"]);
let global_hermes = resolve_target_roots(
&args(InstallSkillsScope::Global, InstallSkillsAgent::Hermes),
cwd,
Some(home),
)
.unwrap();
assert_eq!(root_names(&global_hermes), ["/home/alice/.hermes/skills"]);
let project_both = resolve_target_roots(
&args(InstallSkillsScope::Project, InstallSkillsAgent::Both),
cwd,
crates/ai-memory-cli/src/commands/mcp_bridge.rs
+1 -1
View File
@@ -198,7 +198,7 @@ mod tests {
use std::sync::{Arc, Mutex};
use axum::Router;
use rmcp::model::{Content, ServerCapabilities, Tool};
use rmcp::model::{ContentBlock as Content, ServerCapabilities, Tool};
use rmcp::transport::streamable_http_server::session::local::LocalSessionManager;
use rmcp::transport::streamable_http_server::{
StreamableHttpServerConfig, StreamableHttpService,
crates/ai-memory-cli/src/commands/mod.rs
+7
View File
@@ -40,6 +40,7 @@ pub mod export_okf;
pub mod finalize_session;
pub mod forget_sweep;
pub mod generate_auth_token;
pub mod grant;
pub mod handoffs;
pub mod hook;
pub mod hook_capture;
@@ -59,15 +60,19 @@ pub mod move_session;
pub mod openclaw_plugin;
pub mod path_util;
pub mod pending_writes;
pub mod project;
pub mod project_registry;
pub mod purge_preview;
pub mod purge_project;
pub mod purge_session;
pub mod read_page;
pub mod reclaim_ledger_versions;
pub mod reindex;
pub mod rename_project;
pub mod rename_workstream;
pub mod render_shared;
pub mod reorg;
pub mod repair_backfill_timestamps;
pub mod reset;
pub mod restore;
pub mod restore_page;
@@ -76,10 +81,12 @@ pub mod run;
pub mod run_autowire;
pub mod search;
pub mod serve;
pub mod server;
pub mod setup_agent;
pub mod show;
pub mod status;
pub mod uninstall;
pub mod upgrade;
pub mod user;
pub mod workstream_search;
pub mod workstreams;
crates/ai-memory-cli/src/commands/openclaw_plugin.rs
+27 -26
View File
@@ -10,6 +10,7 @@ use crate::cli::InstallHooksArgs;
use crate::commands::apply_shared::{ApplyOutcome, apply_atomic};
use crate::commands::render_shared::{
ts_capture_policy_v1, ts_resolve_token_fn, ts_spool_runtime, ts_string_literal,
ts_timeout_signal,
};
pub(crate) const PLUGIN_ID: &str = "ai-memory";
@@ -340,12 +341,7 @@ const AGENT = "openclaw";
{token_line}{resolve_fn}
{capture_policy}
function timeoutSignal(ms: number): AbortSignal | undefined {{
if (typeof AbortSignal === "undefined") return undefined;
const factory = (AbortSignal as unknown as {{ timeout?: (ms: number) => AbortSignal }}).timeout;
return factory ? factory(ms) : undefined;
}}
{timeout_signal}
function authHeaders(): Record<string, string> {{
const token = resolveToken();
return token ? {{ Authorization: `Bearer ${{token}}` }} : {{}};
@@ -388,26 +384,7 @@ function tomlKey(text: string, key: string): string | undefined {{
}}
function repoRootProject(cwd: string | undefined): string | undefined {{
if (!cwd) return undefined;
try {{
const inside = execFileSync("git", ["-C", cwd, "rev-parse", "--is-inside-work-tree"], {{
encoding: "utf8",
stdio: ["ignore", "pipe", "ignore"],
}}).trim();
if (inside !== "true") return undefined;
const common = execFileSync("git", ["-C", cwd, "rev-parse", "--path-format=absolute", "--git-common-dir"], {{
encoding: "utf8",
stdio: ["ignore", "pipe", "ignore"],
}}).trim();
if (!common) return undefined;
const root = dirname(common);
if (!root || root === dirname(root)) return undefined;
return basename(root);
}} catch (_e) {{
return undefined;
}}
}}
{repo_root_project}
{apply_marker_params}
function textFrom(value: unknown): string {{
@@ -523,6 +500,7 @@ function postPreCompact(event: any, ctx: any): void {{
async function fetchHandoff(event: any, ctx: any): Promise<string | undefined> {{
const currentCwd = cwd(event, ctx);
if (!currentCwd) return undefined;
if (captureServerRouted(currentCwd)) return undefined;
const url = new URL(`${{SERVER}}/handoff`);
url.searchParams.set("agent", AGENT);
applyMarkerParams(url, currentCwd);
@@ -603,6 +581,8 @@ export default definePluginEntry({{
"#,
server_literal = ts_string_literal(server_url),
token_line = token_line,
repo_root_project = super::install_hooks::TS_REPO_ROOT_PROJECT,
timeout_signal = ts_timeout_signal(),
spool_runtime = ts_spool_runtime(),
)
}
@@ -624,6 +604,7 @@ mod tests {
.contains("if (!resp || resp.status >= 500) spoolFailedHook(url, policy.payload);")
);
assert!(plugin.contains("else requestSpoolDrain();"));
crate::commands::render_shared::assert_shared_ts_delivery_runtime("openclaw", &plugin);
assert!(plugin.contains(r#"return join(env, "hook-spool");"#));
for f in [
"mkdirSync",
@@ -698,6 +679,8 @@ mod tests {
assert!(plugin.contains("if (existsSync(join(probe, \".git\")))"));
assert!(plugin.contains("boundary ??= dir;"));
assert!(plugin.contains("function repoRootProject"));
assert!(plugin.contains("repoProjectCache.set(cwd, project);"));
assert_eq!(plugin.matches("windowsHide: true").count(), 2);
assert!(plugin.contains("--git-common-dir"));
assert!(
plugin
@@ -776,6 +759,24 @@ mod tests {
);
}
/// #992: OpenClaw posts to `SERVER` directly and does not route `server`
/// profiles, so a routed repository must emit nothing and fetch no handoff.
#[test]
fn openclaw_plugin_fails_closed_on_a_server_profile_marker() {
let plugin = build_plugin("http://127.0.0.1:49374", None, None, "denylist");
assert!(
plugin.contains(
"if (captureServerRouted(cwd)) return { disposition: \"drop\", payload };"
),
"{plugin}"
);
let handoff = plugin.split_once("async function fetchHandoff(").unwrap().1;
let guard = handoff
.find("if (captureServerRouted(currentCwd)) return undefined;")
.expect("handoff fetch must be gated");
assert!(guard < handoff.find("/handoff`").unwrap());
}
#[test]
fn openclaw_plugin_denylist_bakes_inert_gate() {
let plugin = build_plugin("http://127.0.0.1:49374", Some("tok"), None, "denylist");
crates/ai-memory-cli/src/commands/path_util.rs
+33 -6
View File
@@ -3,6 +3,35 @@
use std::borrow::Cow;
use std::path::{Path, PathBuf};
/// Create `dir` (and its parents) `0700` on Unix. On non-Unix the mode is a
/// no-op (falls back to `create_dir_all`). Idempotent: an existing
/// directory's mode is left untouched (never widened, never narrowed) to
/// avoid churning a path an operator may have set deliberately.
pub(crate) fn create_private_dir(dir: &Path) -> std::io::Result<()> {
#[cfg(unix)]
{
use std::os::unix::fs::DirBuilderExt as _;
if dir.is_dir() {
return Ok(());
}
match std::fs::DirBuilder::new()
.recursive(true)
.mode(0o700)
.create(dir)
{
Ok(()) => Ok(()),
// A concurrent caller may have created it between the check and
// the call; treat an existing directory as success.
Err(_) if dir.is_dir() => Ok(()),
Err(e) => Err(e),
}
}
#[cfg(not(unix))]
{
std::fs::create_dir_all(dir)
}
}
/// Resolve the user home used for agent configuration paths.
///
/// Tests and scripted wrappers can set `AI_MEMORY_HOME` to exercise the
@@ -20,18 +49,16 @@ pub(crate) fn home_dir() -> Option<PathBuf> {
///
/// Blank is treated as unset on purpose: an exported-but-empty variable is far
/// more often an unset shell expansion than a deliberate request to install
/// into the filesystem root.
/// into the filesystem root. The rule lives in
/// [`ai_memory_workstream::env_dir_override`] so native session import applies
/// the same one and never reads a store the installers did not wire.
///
/// The env value comes in as a parameter so tests can exercise both branches
/// without mutating process env — which is not merely inconvenient here but
/// forbidden: `std::env::set_var` is `unsafe` under edition 2024 and this
/// workspace forbids `unsafe_code`.
pub(crate) fn agent_config_home(env_override: Option<std::ffi::OsString>) -> Option<PathBuf> {
let value = env_override?;
if value.to_str().is_some_and(|s| s.trim().is_empty()) {
return None;
}
Some(PathBuf::from(value))
ai_memory_workstream::env_dir_override(env_override)
}
/// Claude Code's relocated config root: `$CLAUDE_CONFIG_DIR` when set, else
crates/ai-memory-cli/src/commands/project.rs
+65
View File
@@ -0,0 +1,65 @@
//! `ai-memory project` — project settings (#708).
//!
//! Thin HTTP client over `/admin/projects/*`, root-only like `ai-memory
//! grant`: access is an operator decision, made against the server rather than
//! a database the operator's laptop usually cannot open.
use anyhow::{Context, Result};
use serde::Deserialize;
use crate::cli::{ProjectAccessArgs, ProjectArgs, ProjectCommand};
use crate::config::Config;
use crate::http_client::{ServerEndpoint, post_json};
/// Dispatch a `project` subcommand.
///
/// # Errors
/// Transport failures, a non-root token, an unknown project or mode.
pub async fn run(config: &Config, args: ProjectArgs) -> Result<()> {
let ep = ServerEndpoint::from_config_resolving_auth(config).await;
match args.command {
ProjectCommand::Access(args) => access(&ep, args).await,
ProjectCommand::Grants(args) => crate::commands::grant::list_for_project(&ep, &args).await,
}
}
#[derive(Debug, Deserialize)]
struct AccessResponse {
mode: String,
previous: String,
changed: bool,
without_access: Vec<String>,
}
async fn access(ep: &ServerEndpoint, args: ProjectAccessArgs) -> Result<()> {
let resp: AccessResponse = post_json(
ep,
"/admin/projects/access",
&serde_json::json!({
"workspace": args.workspace,
"project": args.project,
"mode": args.mode,
}),
)
.await
.context("setting project access")?;
let repo = format!("{}/{}", args.workspace, args.project);
if resp.changed {
println!("{repo} is now {} (was {}).", resp.mode, resp.previous);
} else {
println!("{repo} was already {}; nothing changed.", resp.mode);
}
if !resp.without_access.is_empty() {
println!();
println!("These users have written to {repo} and hold no grant, so they are now refused:");
for name in &resp.without_access {
println!(" {name}");
}
println!(
"Grant the ones who should keep access: ai-memory user grant --user <name> \
--workspace {} --project {} --level write",
args.workspace, args.project
);
}
Ok(())
}
crates/ai-memory-cli/src/commands/purge_preview.rs
+103
View File
@@ -0,0 +1,103 @@
//! Shared best-effort dry-run preview machinery for `purge-project` and
//! `purge-session`.
//!
//! Both subcommands refuse without `--confirm`, but first ask the server for
//! a bounded preview (`"dry_run": true`, which wins over `confirm`
//! server-side) of what the confirmed run would do. The HTTP round trip, its
//! timeout, and the classification of the result into "print this" vs "stay
//! silent" are identical between the two commands — only the request body's
//! shape, the endpoint path, and the summary line's wording differ, so those
//! stay in `purge_project.rs` / `purge_session.rs` and this module holds only
//! what would otherwise be copy-pasted between them.
use std::time::Duration;
use serde::Serialize;
use crate::config::Config;
use crate::http_client::{ServerEndpoint, ServerResponseError, post_json};
/// How long the preview request (auth resolution plus the HTTP round trip)
/// is allowed to take before this falls back to the plain refusal. The
/// preview is optional information layered on top of a refusal that must
/// still happen either way, so it must never be the reason `--confirm` takes
/// noticeably longer than it used to.
pub const PREVIEW_TIMEOUT: Duration = Duration::from_secs(5);
/// What became of the best-effort preview request, reduced to what deciding
/// whether — and what — to print needs. Kept separate from the network call
/// itself so the printing decision is a pure function and testable without a
/// server.
pub enum PreviewOutcome {
/// A 200 with `"dry_run": true`: the server understood the request and
/// ran the preview.
Previewed(serde_json::Value),
/// A 200 without `dry_run` set: an old-enough server both predates the
/// field AND happens to 200 an unrecognized shape. Treated the same as
/// not getting a preview at all.
Ignored,
/// A non-2xx response with a body worth showing: the scope resolved to
/// something the operator should know about before the refusal (a 404
/// naming the missing project/session, a 409 naming a live managed run,
/// a 403 naming the auth problem), or an unexpected status this command
/// has no specific handling for.
Refused { status: u16, body: String },
/// The request predates `dry_run` support (400, the pre-existing
/// "confirm=true" refusal body) — not worth repeating, since `run` below
/// prints its own version of exactly that message next regardless.
OlderServer,
/// Timed out or never reached a server at all (DNS/connect failure,
/// auth-refresh hang, etc). Indistinguishable from the operator's
/// perspective, and neither is this command's business to diagnose.
Unreachable,
}
/// Run the preview request under [`PREVIEW_TIMEOUT`] and classify the
/// result. Auth resolution (`ServerEndpoint::from_config_resolving_auth`,
/// which can itself refresh an OIDC token over the network) runs inside the
/// same timeout so a hung refresh cannot silently make a `--confirm`-less
/// purge command block far longer than the rest of it ever has.
pub async fn run_preview<Req: Serialize>(
config: &Config,
path: &str,
request: &Req,
) -> PreviewOutcome {
let attempt = tokio::time::timeout(PREVIEW_TIMEOUT, async {
let endpoint = ServerEndpoint::from_config_resolving_auth(config).await;
post_json::<_, serde_json::Value>(&endpoint, path, request).await
})
.await;
let Ok(result) = attempt else {
return PreviewOutcome::Unreachable;
};
match result {
Ok(report) if report["dry_run"].as_bool().unwrap_or(false) => {
PreviewOutcome::Previewed(report)
}
Ok(_) => PreviewOutcome::Ignored,
Err(e) => match e.downcast_ref::<ServerResponseError>() {
Some(resp) if resp.status().as_u16() == 400 => PreviewOutcome::OlderServer,
Some(resp) => PreviewOutcome::Refused {
status: resp.status().as_u16(),
body: resp.body().to_string(),
},
// Not an HTTP response at all: connect/DNS failure, request
// timeout already handled above, or a body that failed to
// deserialize as JSON.
None => PreviewOutcome::Unreachable,
},
}
}
/// Format a `PreviewOutcome::Refused` body as the one line every caller
/// prints before its refusal: `Preview refused (<status>): <message>`. The
/// message is the server's own `{"error": ...}` string when the body parses
/// that way, and the raw body verbatim otherwise.
pub fn refused_message(status: u16, body: &str) -> String {
let message = serde_json::from_str::<serde_json::Value>(body)
.ok()
.and_then(|v| v.get("error").and_then(|e| e.as_str()).map(str::to_string))
.unwrap_or_else(|| body.to_string());
format!("Preview refused ({status}): {message}")
}
crates/ai-memory-cli/src/commands/purge_project.rs
+360 -17
View File
@@ -4,6 +4,7 @@ use anyhow::{Result, bail};
use serde::Serialize;
use crate::cli::PurgeProjectArgs;
use crate::commands::purge_preview::{PreviewOutcome, refused_message, run_preview};
use crate::config::Config;
use crate::http_client::{ServerEndpoint, post_json};
@@ -17,6 +18,71 @@ struct PurgeProjectRequest {
force: bool,
/// Rebuild the FTS indexes and VACUUM after the delete commits.
compact: bool,
/// Preview only: wins over `confirm` on the server (mirrors
/// `reclaim-ledger-versions`), so this is always sent alongside
/// `confirm: false` here — never both true. Older servers that predate
/// this field simply never look at it (an unknown JSON field is not a
/// deserialize error), which is why the fallback below only has to
/// handle the request failing outright, never the field being silently
/// misread.
dry_run: bool,
}
/// One line naming what a purge (real or previewed) removed, in the fixed
/// order the operator can grep for: pages, sessions, observations, handoffs,
/// embeddings, workstreams, managed runs. `verb` is `"Purged"` for a
/// confirmed run or `"Would purge"` for a preview — everything else about
/// the line is identical, so a script matching one also matches the other.
fn purge_summary_line(verb: &str, label: &str, report: &serde_json::Value) -> String {
let pages = report["pages_deleted"].as_u64().unwrap_or(0);
let sessions = report["sessions_deleted"].as_u64().unwrap_or(0);
let observations = report["observations_deleted"].as_u64().unwrap_or(0);
let handoffs = report["handoffs_deleted"].as_u64().unwrap_or(0);
let embeddings = report["embeddings_deleted"].as_u64().unwrap_or(0);
// Workstreams cascade out of the project row, so a scope that looks empty
// by every other counter can still be carrying a managed workstream and
// its portable event ledger. Always name them.
let workstreams = report["workstreams_deleted"].as_u64().unwrap_or(0);
let managed_runs = report["managed_runs_deleted"].as_u64().unwrap_or(0);
let mut line = format!(
"{verb} {label}: {pages} pages, {sessions} sessions, \
{observations} observations, {handoffs} handoffs, {embeddings} embeddings, \
{workstreams} workstreams, {managed_runs} managed runs."
);
// The mirror of the incident this command guards against: purging this
// scope also collaterally deletes/orphans rows that live in *other*
// projects, via `sessions` cascading out of this one. Silent when zero
// so the common case reads exactly as before.
let collateral_observations = report["collateral_observations_deleted"]
.as_u64()
.unwrap_or(0);
if collateral_observations > 0 {
line.push_str(&format!(
" Plus {collateral_observations} observations in other projects via their sessions."
));
}
let collateral_handoffs = report["collateral_handoffs_denulled"].as_u64().unwrap_or(0);
if collateral_handoffs > 0 {
line.push_str(&format!(
" Plus {collateral_handoffs} handoffs in other projects that will lose their \
session reference (set to NULL, not deleted)."
));
}
line
}
/// Decide what to print, if anything, before the refusal — pure, so it is
/// unit-tested without a server. `fallback_label` is used only when the
/// server's own report has no `label` field.
fn preview_message(outcome: &PreviewOutcome, fallback_label: &str) -> Option<String> {
match outcome {
PreviewOutcome::Previewed(report) => {
let label = report["label"].as_str().unwrap_or(fallback_label);
Some(purge_summary_line("Would purge", label, report))
}
PreviewOutcome::Refused { status, body } => Some(refused_message(*status, body)),
PreviewOutcome::Ignored | PreviewOutcome::OlderServer | PreviewOutcome::Unreachable => None,
}
}
/// Run the `purge-project` subcommand.
@@ -25,14 +91,42 @@ struct PurgeProjectRequest {
/// `--project` is omitted), requires `--confirm` before sending the
/// destructive request, then prints the JSON summary.
///
/// Without `--confirm`, first asks the server for a preview (`dry_run:
/// true`, which wins over `confirm` server-side): the server reports the
/// counts a confirmed purge would produce without deleting anything. That
/// preview is best-effort, bounded by
/// [`purge_preview::PREVIEW_TIMEOUT`](crate::commands::purge_preview::PREVIEW_TIMEOUT),
/// and never changes the outcome — only what gets printed before it:
/// - a successful preview prints the "Would purge ..." line;
/// - a 404/409/403 (or any other unexpected status) prints the server's own
/// error first, since the operator asked what would happen and the server
/// has an answer, just not the one this command expected;
/// - a plain 400 (an older server that predates `dry_run`), a timeout, or an
/// unreachable server print nothing extra — the refusal below already
/// says everything a 400 would.
///
/// # Errors
/// Returns an error when `--confirm` is absent, the server is unreachable,
/// or the server returns a non-2xx response.
/// Returns an error when `--confirm` is absent (after printing whatever the
/// preview surfaced), the server is unreachable, or the server returns a
/// non-2xx response.
pub async fn run(config: &Config, args: PurgeProjectArgs) -> Result<()> {
let (workspace, project) =
super::resolve_scope(config, args.workspace.as_deref(), args.project.as_deref())?;
if !args.confirm {
let request = PurgeProjectRequest {
workspace: workspace.clone(),
project: project.clone(),
confirm: false,
force: args.force,
compact: args.compact,
dry_run: true,
};
let outcome = run_preview(config, "/admin/purge-project", &request).await;
let fallback_label = format!("{}/{}", workspace, project);
if let Some(line) = preview_message(&outcome, &fallback_label) {
println!("{line}");
}
bail!(
"purge-project is destructive and irreversible.\n\
Re-run with --confirm to proceed:\n\n \
@@ -52,6 +146,7 @@ pub async fn run(config: &Config, args: PurgeProjectArgs) -> Result<()> {
confirm: true,
force: args.force,
compact: args.compact,
dry_run: false,
},
)
.await?;
@@ -59,21 +154,7 @@ pub async fn run(config: &Config, args: PurgeProjectArgs) -> Result<()> {
// Human-friendly one-liner followed by the raw JSON for scripting.
let fallback_label = format!("{}/{}", workspace, project);
let label = report["label"].as_str().unwrap_or(&fallback_label);
let pages = report["pages_deleted"].as_u64().unwrap_or(0);
let sessions = report["sessions_deleted"].as_u64().unwrap_or(0);
let observations = report["observations_deleted"].as_u64().unwrap_or(0);
let handoffs = report["handoffs_deleted"].as_u64().unwrap_or(0);
let embeddings = report["embeddings_deleted"].as_u64().unwrap_or(0);
// Workstreams cascade out of the project row, so a scope that looks empty
// by every other counter can still be carrying a managed workstream and
// its portable event ledger. Always name them.
let workstreams = report["workstreams_deleted"].as_u64().unwrap_or(0);
let managed_runs = report["managed_runs_deleted"].as_u64().unwrap_or(0);
println!(
"Purged {label}: {pages} pages, {sessions} sessions, \
{observations} observations, {handoffs} handoffs, {embeddings} embeddings, \
{workstreams} workstreams, {managed_runs} managed runs."
);
println!("{}", purge_summary_line("Purged", label, &report));
if let Some(ids) = report["workstream_ids"].as_array()
&& !ids.is_empty()
{
@@ -105,3 +186,265 @@ pub async fn run(config: &Config, args: PurgeProjectArgs) -> Result<()> {
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
fn report(dry_run: bool) -> serde_json::Value {
serde_json::json!({
"label": "default/my-project",
"pages_deleted": 3,
"sessions_deleted": 1,
"observations_deleted": 1063,
"handoffs_deleted": 0,
"embeddings_deleted": 3,
"collateral_observations_deleted": 0,
"collateral_handoffs_denulled": 0,
"workstreams_deleted": 0,
"managed_runs_deleted": 0,
"workstream_ids": [],
"files_deleted": [],
"files_failed": [],
"compacted": false,
"dry_run": dry_run,
})
}
/// Pins the exact wording and field order a script would grep for.
/// `purge_summary_line` is the single source for both the confirmed
/// "Purged" line and the preview's "Would purge" line, so this also
/// proves the two can never drift apart.
#[test]
fn purge_summary_line_matches_the_documented_wording() {
let confirmed = report(false);
assert_eq!(
purge_summary_line("Purged", "default/my-project", &confirmed),
"Purged default/my-project: 3 pages, 1 sessions, 1063 observations, \
0 handoffs, 3 embeddings, 0 workstreams, 0 managed runs."
);
let preview = report(true);
assert_eq!(
purge_summary_line("Would purge", "default/my-project", &preview),
"Would purge default/my-project: 3 pages, 1 sessions, 1063 observations, \
0 handoffs, 3 embeddings, 0 workstreams, 0 managed runs."
);
}
/// Missing counters must not panic and must not silently show as
/// non-zero: a malformed or truncated reply reads as all-zero, which the
/// operator can visibly tell apart from a real "0 pages, 0 sessions".
#[test]
fn purge_summary_line_defaults_missing_counters_to_zero() {
let empty = serde_json::json!({});
assert_eq!(
purge_summary_line("Would purge", "default/x", &empty),
"Would purge default/x: 0 pages, 0 sessions, 0 observations, \
0 handoffs, 0 embeddings, 0 workstreams, 0 managed runs."
);
}
/// The mirror-of-the-incident collateral counts, when present, must be
/// visible in the same line the operator already reads — silent when
/// zero (the common case), spelled out when not.
#[test]
fn purge_summary_line_calls_out_collateral_damage_when_present() {
let mut r = report(true);
r["collateral_observations_deleted"] = serde_json::json!(7);
r["collateral_handoffs_denulled"] = serde_json::json!(2);
let line = purge_summary_line("Would purge", "default/looks-empty", &r);
assert!(
line.contains("Plus 7 observations in other projects via their sessions."),
"collateral observations must be called out: {line}"
);
assert!(
line.contains("Plus 2 handoffs in other projects"),
"collateral handoffs must be called out: {line}"
);
}
/// A successful preview prints the "Would purge" line.
#[test]
fn preview_message_prints_the_would_purge_line_on_success() {
let outcome = PreviewOutcome::Previewed(report(true));
let msg = preview_message(&outcome, "default/fallback").expect("must print a line");
assert!(msg.starts_with("Would purge default/my-project:"));
}
/// A 404 (unknown scope) is worth showing before the refusal: the
/// operator asked what would happen, and 404 is the server's answer.
#[test]
fn preview_message_surfaces_a_404_before_the_refusal() {
let outcome = PreviewOutcome::Refused {
status: 404,
body: r#"{"error":"project 'ghost' not found in workspace 'default'"}"#.to_string(),
};
let msg = preview_message(&outcome, "default/ghost").expect("must print a line");
assert_eq!(
msg,
"Preview refused (404): project 'ghost' not found in workspace 'default'"
);
}
/// A 409 (live managed run) is the same: surfaced, not swallowed.
#[test]
fn preview_message_surfaces_a_409_before_the_refusal() {
let outcome = PreviewOutcome::Refused {
status: 409,
body: r#"{"error":"managed run lease is active for 'main' (claude-code)"}"#.to_string(),
};
let msg = preview_message(&outcome, "default/x").expect("must print a line");
assert_eq!(
msg,
"Preview refused (409): managed run lease is active for 'main' (claude-code)"
);
}
/// A body that isn't the expected `{"error": ...}` shape still prints
/// something rather than nothing — the raw body, verbatim.
#[test]
fn preview_message_falls_back_to_the_raw_body_when_not_json() {
let outcome = PreviewOutcome::Refused {
status: 403,
body: "Forbidden".to_string(),
};
let msg = preview_message(&outcome, "default/x").expect("must print a line");
assert_eq!(msg, "Preview refused (403): Forbidden");
}
/// The three silent-fallback cases: an older server's plain 400, a
/// timeout/connect failure, and a 200 that oddly never set `dry_run`.
/// None of these should print anything — the refusal that follows in
/// `run` already says everything a 400 would, and there is nothing
/// useful to say about a request that never got an answer.
#[test]
fn preview_message_is_silent_for_older_server_unreachable_and_ignored() {
assert!(preview_message(&PreviewOutcome::OlderServer, "default/x").is_none());
assert!(preview_message(&PreviewOutcome::Unreachable, "default/x").is_none());
assert!(preview_message(&PreviewOutcome::Ignored, "default/x").is_none());
}
// -----------------------------------------------------------------
// `run_preview` against a real (local) server: proves the HTTP-status
// classification end to end, not just the pure `preview_message` mapping
// above.
// -----------------------------------------------------------------
fn config_for(tmp: &tempfile::TempDir, server_url: String) -> Config {
Config {
data_dir: tmp.path().to_path_buf(),
server_url,
..Config::default()
}
}
async fn spawn_fixed_response(status: u16, body: &'static str) -> String {
let app = axum::Router::new().route(
"/admin/purge-project",
axum::routing::post(move || async move {
(
axum::http::StatusCode::from_u16(status).unwrap(),
axum::Json(serde_json::from_str::<serde_json::Value>(body).unwrap()),
)
}),
);
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}")
}
fn preview_request() -> PurgeProjectRequest {
PurgeProjectRequest {
workspace: "default".into(),
project: "scratch".into(),
confirm: false,
force: false,
compact: false,
dry_run: true,
}
}
/// An older server that predates `dry_run` answers the plain 400
/// "confirm=true" refusal it always has — the CLI must fall back
/// silently, not print anything extra.
#[tokio::test]
async fn run_preview_classifies_a_400_as_older_server() {
let tmp = tempfile::TempDir::new().unwrap();
let url = spawn_fixed_response(
400,
r#"{"error": "destructive operation requires confirm=true"}"#,
)
.await;
let config = config_for(&tmp, url);
let outcome = run_preview(&config, "/admin/purge-project", &preview_request()).await;
assert!(matches!(outcome, PreviewOutcome::OlderServer));
assert!(preview_message(&outcome, "default/scratch").is_none());
}
/// A 404 is surfaced: the operator asked what a purge of this scope
/// would do, and "no such scope" is a real, useful answer.
#[tokio::test]
async fn run_preview_classifies_a_404_as_refused_and_surfaces_the_message() {
let tmp = tempfile::TempDir::new().unwrap();
let url = spawn_fixed_response(
404,
r#"{"error": "project 'scratch' not found in workspace 'default'"}"#,
)
.await;
let config = config_for(&tmp, url);
let outcome = run_preview(&config, "/admin/purge-project", &preview_request()).await;
match &outcome {
PreviewOutcome::Refused { status, body } => {
assert_eq!(*status, 404);
assert!(body.contains("not found"));
}
_ => panic!("expected Refused, got a different outcome"),
}
let msg = preview_message(&outcome, "default/scratch").expect("must print a line");
assert_eq!(
msg,
"Preview refused (404): project 'scratch' not found in workspace 'default'"
);
}
/// A successful preview (200, `dry_run: true`) is a `Previewed` outcome
/// carrying the report through untouched.
#[tokio::test]
async fn run_preview_classifies_a_successful_preview() {
let tmp = tempfile::TempDir::new().unwrap();
let url = spawn_fixed_response(
200,
r#"{"label": "default/scratch", "pages_deleted": 3, "sessions_deleted": 1,
"observations_deleted": 1063, "handoffs_deleted": 0, "embeddings_deleted": 3,
"collateral_observations_deleted": 0, "collateral_handoffs_denulled": 0,
"workstreams_deleted": 0, "managed_runs_deleted": 0, "workstream_ids": [],
"files_deleted": [], "files_failed": [], "compacted": false, "dry_run": true}"#,
)
.await;
let config = config_for(&tmp, url);
let outcome = run_preview(&config, "/admin/purge-project", &preview_request()).await;
let msg = preview_message(&outcome, "default/scratch").expect("must print a line");
assert!(msg.starts_with("Would purge default/scratch: 3 pages, 1 sessions, 1063"));
}
/// Nothing listening at all (connection refused) must classify as
/// `Unreachable`, the same as a timeout — both are silent fallbacks.
#[tokio::test]
async fn run_preview_classifies_a_connection_failure_as_unreachable() {
let tmp = tempfile::TempDir::new().unwrap();
// Bind then drop immediately: the port is very likely free again by
// the time the request lands, and nothing else can be listening on
// it inside this test's short lifetime.
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
drop(listener);
let config = config_for(&tmp, format!("http://{addr}"));
let outcome = run_preview(&config, "/admin/purge-project", &preview_request()).await;
assert!(matches!(outcome, PreviewOutcome::Unreachable));
assert!(preview_message(&outcome, "default/scratch").is_none());
}
}
crates/ai-memory-cli/src/commands/purge_session.rs
+348 -14
View File
@@ -4,6 +4,7 @@ use anyhow::{Result, bail};
use serde::Serialize;
use crate::cli::PurgeSessionArgs;
use crate::commands::purge_preview::{PreviewOutcome, refused_message, run_preview};
use crate::config::Config;
use crate::http_client::{ServerEndpoint, post_json};
@@ -16,6 +17,72 @@ struct PurgeSessionRequest {
confirm: bool,
/// Rebuild the FTS indexes and VACUUM after the delete commits.
compact: bool,
/// Preview only: wins over `confirm` on the server (mirrors
/// `purge-project`, which itself mirrors `reclaim-ledger-versions`), so
/// this is always sent alongside `confirm: false` here — never both
/// true. Older servers that predate this field simply never look at it
/// (an unknown JSON field is not a deserialize error), which is why the
/// fallback in [`purge_preview::run_preview`](crate::commands::purge_preview::run_preview)
/// only has to handle the request failing outright, never the field
/// being silently misread.
dry_run: bool,
}
/// One line naming what a session purge (real or previewed) removed, in the
/// fixed order the operator can grep for: observations, handoffs, pages,
/// auto-improve runs. `verb` is `"Purged"` for a confirmed run or `"Would
/// purge"` for a preview — everything else about the line is identical, so a
/// script matching one also matches the other.
///
/// The session id is deliberately never part of this line, confirmed or
/// previewed: the caller already has it (they passed `--session-id`), and
/// this command exists to make a session stop existing — echoing its id
/// into terminal scrollback and shell history leaves a pointer to the thing
/// just erased (or about to be). The scope (`label`) and the counts are what
/// confirm the operation did, or would do, what was asked.
fn purge_summary_line(verb: &str, label: &str, report: &serde_json::Value) -> String {
let observations = report["observations_deleted"].as_u64().unwrap_or(0);
let handoffs = report["handoffs_deleted"].as_u64().unwrap_or(0);
let pages = report["pages_deleted"].as_u64().unwrap_or(0);
let auto_improve_runs = report["auto_improve_runs_deleted"].as_u64().unwrap_or(0);
let mut line = format!(
"{verb} session from {label}: {observations} observations, {handoffs} handoffs, \
{pages} pages, {auto_improve_runs} auto-improve runs."
);
// The mirror of the incident `purge-project`'s own preview guards
// against, one level down at session granularity: purging this session
// also collaterally deletes/orphans rows that live in *other* projects,
// via this session's own id (not the project's rows) cascading. Silent
// when zero so the common case reads exactly as before.
let collateral_observations = report["collateral_observations_deleted"]
.as_u64()
.unwrap_or(0);
if collateral_observations > 0 {
line.push_str(&format!(
" Plus {collateral_observations} observations in other projects via this session."
));
}
let collateral_handoffs = report["collateral_handoffs_denulled"].as_u64().unwrap_or(0);
if collateral_handoffs > 0 {
line.push_str(&format!(
" Plus {collateral_handoffs} handoffs in other projects that will lose their \
session reference (set to NULL, not deleted)."
));
}
line
}
/// Decide what to print, if anything, before the refusal — pure, so it is
/// unit-tested without a server. `fallback_label` is used only when the
/// server's own report has no scope to build a label from (this endpoint has
/// no `label` field, unlike `purge-project`, so the caller always supplies
/// one built from `--workspace`/`--project`).
fn preview_message(outcome: &PreviewOutcome, label: &str) -> Option<String> {
match outcome {
PreviewOutcome::Previewed(report) => Some(purge_summary_line("Would purge", label, report)),
PreviewOutcome::Refused { status, body } => Some(refused_message(*status, body)),
PreviewOutcome::Ignored | PreviewOutcome::OlderServer | PreviewOutcome::Unreachable => None,
}
}
/// Run the `purge-session` subcommand.
@@ -25,6 +92,21 @@ struct PurgeSessionRequest {
/// session that does not belong to it, so a UUID alone is never authority
/// over another workspace or project.
///
/// Without `--confirm`, first asks the server for a preview (`dry_run:
/// true`, which wins over `confirm` server-side, exactly like
/// `purge-project`): the server reports the counts a confirmed purge would
/// produce without deleting anything. That preview is best-effort, bounded
/// by
/// [`purge_preview::PREVIEW_TIMEOUT`](crate::commands::purge_preview::PREVIEW_TIMEOUT),
/// and never changes the outcome — only what gets printed before it:
/// - a successful preview prints the "Would purge session from ..." line;
/// - a 404/403 (or any other unexpected status) prints the server's own
/// error first, since the operator asked what would happen and the server
/// has an answer, just not the one this command expected;
/// - a plain 400 (an older server that predates `dry_run`), a timeout, or an
/// unreachable server print nothing extra — the refusal below already
/// says everything a 400 would.
///
/// # Errors
/// Returns an error when `--confirm` is absent, the session id is not a
/// UUID, the server is unreachable, or the server returns a non-2xx
@@ -44,7 +126,21 @@ pub async fn run(config: &Config, args: PurgeSessionArgs) -> Result<()> {
);
}
let label = format!("{workspace}/{project}");
if !args.confirm {
let request = PurgeSessionRequest {
workspace: workspace.clone(),
project: project.clone(),
session_id: session_id.to_owned(),
confirm: false,
compact: args.compact,
dry_run: true,
};
let outcome = run_preview(config, "/admin/purge-session", &request).await;
if let Some(line) = preview_message(&outcome, &label) {
println!("{line}");
}
bail!(
"purge-session is destructive and irreversible.\n\
Re-run with --confirm to proceed:\n\n \
@@ -63,24 +159,12 @@ pub async fn run(config: &Config, args: PurgeSessionArgs) -> Result<()> {
session_id: session_id.to_owned(),
confirm: true,
compact: args.compact,
dry_run: false,
},
)
.await?;
let n = |key: &str| report[key].as_u64().unwrap_or(0);
// The session id is deliberately not echoed. The caller passed it, so
// repeating it adds nothing they do not have, and this command exists to
// make a session stop existing — writing its id into terminal scrollback
// and shell history leaves a pointer to the thing just erased. The scope
// and the counts are what confirm the operation did what was asked.
println!(
"Purged session from {workspace}/{project}: \
{} observations, {} handoffs, {} pages, {} auto-improve runs.",
n("observations_deleted"),
n("handoffs_deleted"),
n("pages_deleted"),
n("auto_improve_runs_deleted"),
);
println!("{}", purge_summary_line("Purged", &label, &report));
if let Some(paths) = report["files_deleted"].as_array()
&& !paths.is_empty()
{
@@ -112,3 +196,253 @@ pub async fn run(config: &Config, args: PurgeSessionArgs) -> Result<()> {
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
fn report(dry_run: bool) -> serde_json::Value {
serde_json::json!({
"session_id": "00000000-0000-0000-0000-000000000000",
"workspace": "default",
"project": "my-project",
"observations_deleted": 1063,
"handoffs_deleted": 0,
"pages_deleted": 1,
"auto_improve_runs_deleted": 2,
"removed_paths": [],
"collateral_observations_deleted": 0,
"collateral_handoffs_denulled": 0,
"files_deleted": [],
"files_failed": [],
"compacted": false,
"dry_run": dry_run,
})
}
/// Pins the exact wording and field order a script would grep for, and
/// that the session id never appears in it.
/// `purge_summary_line` is the single source for both the confirmed
/// "Purged" line and the preview's "Would purge" line, so this also
/// proves the two can never drift apart.
#[test]
fn purge_summary_line_matches_the_documented_wording() {
let confirmed = report(false);
let line = purge_summary_line("Purged", "default/my-project", &confirmed);
assert_eq!(
line,
"Purged session from default/my-project: 1063 observations, 0 handoffs, \
1 pages, 2 auto-improve runs."
);
assert!(
!line.contains("00000000"),
"the session id must never appear in the printed line: {line}"
);
let preview = report(true);
assert_eq!(
purge_summary_line("Would purge", "default/my-project", &preview),
"Would purge session from default/my-project: 1063 observations, 0 handoffs, \
1 pages, 2 auto-improve runs."
);
}
/// Missing counters must not panic and must not silently show as
/// non-zero.
#[test]
fn purge_summary_line_defaults_missing_counters_to_zero() {
let empty = serde_json::json!({});
assert_eq!(
purge_summary_line("Would purge", "default/x", &empty),
"Would purge session from default/x: 0 observations, 0 handoffs, \
0 pages, 0 auto-improve runs."
);
}
/// The mirror-of-the-incident collateral counts, when present, must be
/// visible in the same line the operator already reads — silent when
/// zero (the common case), spelled out when not.
#[test]
fn purge_summary_line_calls_out_collateral_damage_when_present() {
let mut r = report(true);
r["collateral_observations_deleted"] = serde_json::json!(7);
r["collateral_handoffs_denulled"] = serde_json::json!(2);
let line = purge_summary_line("Would purge", "default/looks-empty", &r);
assert!(
line.contains("Plus 7 observations in other projects via this session."),
"collateral observations must be called out: {line}"
);
assert!(
line.contains("Plus 2 handoffs in other projects"),
"collateral handoffs must be called out: {line}"
);
}
/// A successful preview prints the "Would purge" line.
#[test]
fn preview_message_prints_the_would_purge_line_on_success() {
let outcome = PreviewOutcome::Previewed(report(true));
let msg = preview_message(&outcome, "default/my-project").expect("must print a line");
assert!(msg.starts_with("Would purge session from default/my-project:"));
}
/// A 404 (session outside the named scope) is worth showing before the
/// refusal: the operator asked what would happen, and 404 is the
/// server's answer.
#[test]
fn preview_message_surfaces_a_404_before_the_refusal() {
let outcome = PreviewOutcome::Refused {
status: 404,
body: r#"{"error":"session ... not found in this workspace/project"}"#.to_string(),
};
let msg = preview_message(&outcome, "default/ghost").expect("must print a line");
assert_eq!(
msg,
"Preview refused (404): session ... not found in this workspace/project"
);
}
/// A body that isn't the expected `{"error": ...}` shape still prints
/// something rather than nothing — the raw body, verbatim.
#[test]
fn preview_message_falls_back_to_the_raw_body_when_not_json() {
let outcome = PreviewOutcome::Refused {
status: 403,
body: "Forbidden".to_string(),
};
let msg = preview_message(&outcome, "default/x").expect("must print a line");
assert_eq!(msg, "Preview refused (403): Forbidden");
}
/// The three silent-fallback cases: an older server's plain 400, a
/// timeout/connect failure, and a 200 that oddly never set `dry_run`.
#[test]
fn preview_message_is_silent_for_older_server_unreachable_and_ignored() {
assert!(preview_message(&PreviewOutcome::OlderServer, "default/x").is_none());
assert!(preview_message(&PreviewOutcome::Unreachable, "default/x").is_none());
assert!(preview_message(&PreviewOutcome::Ignored, "default/x").is_none());
}
// -----------------------------------------------------------------
// `run_preview` against a real (local) server: proves the HTTP-status
// classification end to end, not just the pure `preview_message` mapping
// above.
// -----------------------------------------------------------------
fn config_for(tmp: &tempfile::TempDir, server_url: String) -> Config {
Config {
data_dir: tmp.path().to_path_buf(),
server_url,
..Config::default()
}
}
async fn spawn_fixed_response(status: u16, body: &'static str) -> String {
let app = axum::Router::new().route(
"/admin/purge-session",
axum::routing::post(move || async move {
(
axum::http::StatusCode::from_u16(status).unwrap(),
axum::Json(serde_json::from_str::<serde_json::Value>(body).unwrap()),
)
}),
);
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}")
}
fn preview_request() -> PurgeSessionRequest {
PurgeSessionRequest {
workspace: "default".into(),
project: "scratch".into(),
session_id: "00000000-0000-0000-0000-000000000000".into(),
confirm: false,
compact: false,
dry_run: true,
}
}
/// An older server that predates `dry_run` answers the plain 400
/// "confirm=true" refusal it always has — the CLI must fall back
/// silently, not print anything extra.
#[tokio::test]
async fn run_preview_classifies_a_400_as_older_server() {
let tmp = tempfile::TempDir::new().unwrap();
let url = spawn_fixed_response(
400,
r#"{"error": "destructive operation requires confirm=true"}"#,
)
.await;
let config = config_for(&tmp, url);
let outcome = run_preview(&config, "/admin/purge-session", &preview_request()).await;
assert!(matches!(outcome, PreviewOutcome::OlderServer));
assert!(preview_message(&outcome, "default/scratch").is_none());
}
/// A 404 is surfaced: the operator asked what a purge of this session
/// would do, and "no such session" is a real, useful answer.
#[tokio::test]
async fn run_preview_classifies_a_404_as_refused_and_surfaces_the_message() {
let tmp = tempfile::TempDir::new().unwrap();
let url = spawn_fixed_response(
404,
r#"{"error": "session ... not found in this workspace/project"}"#,
)
.await;
let config = config_for(&tmp, url);
let outcome = run_preview(&config, "/admin/purge-session", &preview_request()).await;
match &outcome {
PreviewOutcome::Refused { status, body } => {
assert_eq!(*status, 404);
assert!(body.contains("not found"));
}
_ => panic!("expected Refused, got a different outcome"),
}
let msg = preview_message(&outcome, "default/scratch").expect("must print a line");
assert_eq!(
msg,
"Preview refused (404): session ... not found in this workspace/project"
);
}
/// A successful preview (200, `dry_run: true`) is a `Previewed` outcome
/// carrying the report through untouched.
#[tokio::test]
async fn run_preview_classifies_a_successful_preview() {
let tmp = tempfile::TempDir::new().unwrap();
let url = spawn_fixed_response(
200,
r#"{"session_id": "00000000-0000-0000-0000-000000000000", "workspace": "default",
"project": "scratch", "observations_deleted": 1063, "handoffs_deleted": 0,
"pages_deleted": 1, "auto_improve_runs_deleted": 2, "removed_paths": [],
"collateral_observations_deleted": 0, "collateral_handoffs_denulled": 0,
"files_deleted": [], "files_failed": [], "compacted": false, "dry_run": true}"#,
)
.await;
let config = config_for(&tmp, url);
let outcome = run_preview(&config, "/admin/purge-session", &preview_request()).await;
let msg = preview_message(&outcome, "default/scratch").expect("must print a line");
assert!(msg.starts_with("Would purge session from default/scratch: 1063 observations"));
}
/// Nothing listening at all (connection refused) must classify as
/// `Unreachable`, the same as a timeout — both are silent fallbacks.
#[tokio::test]
async fn run_preview_classifies_a_connection_failure_as_unreachable() {
let tmp = tempfile::TempDir::new().unwrap();
// Bind then drop immediately: the port is very likely free again by
// the time the request lands, and nothing else can be listening on
// it inside this test's short lifetime.
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
drop(listener);
let config = config_for(&tmp, format!("http://{addr}"));
let outcome = run_preview(&config, "/admin/purge-session", &preview_request()).await;
assert!(matches!(outcome, PreviewOutcome::Unreachable));
assert!(preview_message(&outcome, "default/scratch").is_none());
}
}
crates/ai-memory-cli/src/commands/reclaim_ledger_versions.rs
+105
View File
@@ -0,0 +1,105 @@
//! `ai-memory reclaim-ledger-versions` — thin HTTP client for dropping the
//! superseded versions of the raw hook event ledger.
use anyhow::Result;
use serde::Serialize;
use crate::cli::ReclaimLedgerVersionsArgs;
use crate::config::Config;
use crate::http_client::{ServerEndpoint, post_json};
use super::compact::human_bytes;
/// Request sent to `POST /admin/reclaim-ledger-versions`.
#[derive(Serialize)]
struct ReclaimLedgerVersionsRequest {
confirm: bool,
dry_run: bool,
drop_latest: bool,
compact: bool,
}
/// Run the `reclaim-ledger-versions` subcommand.
///
/// A dry run by default: the report says what would go, and `--confirm` is
/// what makes it go. `--confirm` is required for the deleting run because,
/// unlike `compact`, this one removes rows and nothing here can put them back.
///
/// # Errors
/// Returns an error when the server is unreachable or returns a non-2xx
/// response.
pub async fn run(config: &Config, args: ReclaimLedgerVersionsArgs) -> Result<()> {
let dry_run = !args.confirm;
let endpoint = ServerEndpoint::from_config_resolving_auth(config).await;
let report: serde_json::Value = post_json(
&endpoint,
"/admin/reclaim-ledger-versions",
&ReclaimLedgerVersionsRequest {
confirm: args.confirm,
dry_run,
drop_latest: args.drop_latest,
compact: args.compact,
},
)
.await?;
let n = |key: &str| report[key].as_u64().unwrap_or(0);
let paths = n("ledger_paths");
let rows = n("pages_deleted");
let body_bytes = n("bytes_deleted");
let dropped_latest = report["dropped_latest"].as_bool().unwrap_or(false);
let compacted = report["compacted"].as_bool().unwrap_or(false);
if paths == 0 || rows == 0 {
println!(
"No superseded ledger versions to reclaim. Nothing written by the \
pre-2.1.1 indexer is left in this store."
);
return Ok(());
}
if dry_run {
println!(
"Would delete {rows} superseded version(s) of the hook event ledger \
across {paths} ledger path(s), carrying {} of page body.",
human_bytes(body_bytes),
);
if dropped_latest {
println!(" --drop-latest: each ledger's live row would go too.");
}
println!("Nothing was changed. To apply:");
let mut flags = String::from(" ai-memory reclaim-ledger-versions --confirm");
if args.drop_latest {
flags.push_str(" --drop-latest");
}
if args.compact {
flags.push_str(" --compact");
}
println!("{flags}");
return Ok(());
}
println!(
"Deleted {rows} superseded ledger version(s) across {paths} ledger \
path(s), carrying {} of page body.",
human_bytes(body_bytes),
);
if dropped_latest {
println!(" Each ledger's live row was dropped as well (--drop-latest).");
}
if compacted {
println!(
"Reclaimed {} ({} → {}).",
human_bytes(n("bytes_reclaimed")),
human_bytes(n("bytes_before")),
human_bytes(n("bytes_after")),
);
} else {
println!(
"The bytes are free pages now. Run `ai-memory compact --confirm` \
to return them to the filesystem."
);
}
Ok(())
}
crates/ai-memory-cli/src/commands/reindex.rs
+7
View File
@@ -79,6 +79,13 @@ pub async fn run(config: &Config, _args: ReindexArgs) -> Result<()> {
summary.workspaces,
config.data_dir.join("wiki").display(),
);
if summary.skipped_collisions > 0 {
println!(
"skipped {} page(s) whose path differs from an indexed page only by case or Unicode \
normalization (they cannot coexist on macOS/Windows); the log names each pair",
summary.skipped_collisions,
);
}
Ok(())
}
crates/ai-memory-cli/src/commands/render_shared.rs
+457 -39
View File
@@ -188,6 +188,21 @@ pub(crate) fn ts_string_literal(s: &str) -> String {
/// repository with no `.ai-memory.toml` marker must emit nothing — mirroring
/// the native admit gate in `commands/hook.rs` (`repository_admits_capture`)
/// — so the check runs before any disposition logic, for every event kind.
///
/// Ahead of even that marker scan sits the external-ownership gate: when
/// `AI_MEMORY_CAPTURE_OWNER` holds any value that is non-empty after trimming,
/// `capturePolicy` returns `drop` immediately, so a generated consumer never
/// queues, spools, or POSTs a *new* capture event. It mirrors the native gate
/// in `commands/hook.rs` and leaves handoff fetching untouched, which is the
/// whole point: an external producer replaces capture without losing native
/// context delivery.
///
/// Two things it deliberately does not do. It does not suppress the consumer's
/// routing work around the call (`applyMarkerParams` reads the marker before
/// `capturePolicy` is even reached), and it does not erase an existing spool
/// backlog: a consumer that reaches `drainHookQueue` — on dispose, for
/// instance — still kicks off `requestSpoolDrain`, so events spooled before the
/// handover can still be delivered.
#[must_use]
pub(crate) fn ts_capture_policy_v1(capture_mode: &str) -> String {
const TEMPLATE: &str = r##"// capture-policy-v1 (generated; do not fork between adapters)
@@ -204,21 +219,35 @@ const CAPTURE_MAX_CALL_ID_CHARS = 128;
type CaptureDisposition = "keep" | "drop" | "metadata-only";
type CaptureProtocol = { version: 1; disposition: CaptureDisposition; policy_state: "inactive" | "active" | "invalid"; tool_family: "file" | "search-list" | "non-file" | "unknown"; path_count: number; extraction_state: "not-applicable" | "extracted" | "missing-or-malformed" | "unsupported-schema" };
type CaptureConfig = { state: "inactive" | "active" | "invalid"; patterns: { path: string; windows: boolean; directory?: string }[]; base: string };
type CaptureConfig = { state: "inactive" | "active" | "invalid"; patterns: { path: string; windows: boolean; directory?: string; prefix: string }[]; base: string; windowsHost: boolean };
function readFileSync(path: string, encoding?: "utf8"): any { if (encoding) return readMarkerText(path, encoding); const fd = openSync(path, "r"); try { const bytes = Buffer.allocUnsafe(CAPTURE_MARKER_MAX_BYTES + 1); const count = readSync(fd, bytes, 0, bytes.length, 0); if (count > CAPTURE_MARKER_MAX_BYTES) throw new Error("marker too large"); const result = bytes.subarray(0, count); new TextDecoder("utf-8", { fatal: true }).decode(result); return result; } finally { closeSync(fd); } }
function captureTrimComment(line: string): string { let quote = ""; let escaped = false; for (let i = 0; i < line.length; i++) { const c = line[i]; if (escaped) { escaped = false; continue; } if (c === "\\" && quote === '"') { escaped = true; continue; } if ((c === '"' || c === "'") && (!quote || quote === c)) quote = quote ? "" : c; else if (c === "#" && !quote) { line = line.slice(0, i); break; } } if (line.trimStart().startsWith("[") && !/^\s*\[[^\]]+\]\s*$/.test(line)) throw new Error("invalid table header"); if (quote) throw new Error("unterminated string"); return line; }
function captureNormalize(path: string): { path: string; windows: boolean } | undefined { const p = path.replace(/\\/g, "/"); let root: string; let tail: string[]; if (p.startsWith("//")) { const x = p.slice(2).split("/").filter(Boolean); if (x.length < 2) return undefined; root = `//${x.shift()}/${x.shift()}`; tail = x; } else if (/^[A-Za-z]:\//.test(p)) { root = `${p[0].toUpperCase()}:/`; tail = p.slice(3).split("/"); } else if (p.startsWith("/")) { root = "/"; tail = p.slice(1).split("/"); } else return undefined; const out: string[] = []; for (const x of tail) { if (!x || x === ".") continue; if (x === "..") out.pop(); else out.push(x); } return { path: root + (out.length ? (root.endsWith("/") ? "" : "/") + out.join("/") : ""), windows: root !== "/" }; }
// `windowsHost` is omitted for self-determining callers (the cwd/marker base,
// and ignore_paths patterns, which are allowed an explicit UNC/drive form
// regardless of host). It is passed explicitly, from the already-resolved
// host base, only when normalizing an untrusted tool-argument candidate: on a
// POSIX host a leading `//` there is an ordinary doubled separator, not a UNC
// root, and must collapse before flavor detection or it escapes every POSIX
// `ignore_paths` pattern via a flavor mismatch (GHSA-vh98).
function captureNormalize(path: string, windowsHost?: boolean): { path: string; windows: boolean } | undefined { const raw = windowsHost === false && path.startsWith("//") ? `/${path.replace(/^\/+/, "")}` : path; const p = raw.replace(/\\/g, "/"); let root: string; let tail: string[]; if (p.startsWith("//")) { const x = p.slice(2).split("/").filter(Boolean); if (x.length < 2) return undefined; root = `//${x.shift()}/${x.shift()}`; tail = x; } else if (/^[A-Za-z]:\//.test(p)) { root = `${p[0].toUpperCase()}:/`; tail = p.slice(3).split("/"); } else if (p.startsWith("/")) { root = "/"; tail = p.slice(1).split("/"); } else return undefined; const out: string[] = []; for (const x of tail) { if (!x || x === ".") continue; if (x === "..") out.pop(); else out.push(x); } return { path: root + (out.length ? (root.endsWith("/") ? "" : "/") + out.join("/") : ""), windows: root !== "/" }; }
function captureJoin(base: string, child: string): string { if (/^[^A-Za-z]?:|^[A-Za-z]:[^/\\]/.test(child)) return child; return `${base.replace(/[\\/]+$/, "")}/${child}`; }
function captureValidGlob(p: string): boolean { return !!p && [...p].length <= CAPTURE_MAX_PATTERN_CHARS && !/[!{}\[\]()|^$%]/.test(p) && !p.includes("${") && !p.includes("***") && !p.replace(/\\/g, "/").split("/").includes("..") && (!p.startsWith("~") || p.startsWith("~/")) && !/^[^A-Za-z]?:/.test(p) && !/^[A-Za-z]:[^/\\]/.test(p); }
function captureParseArray(value: string): string[] | undefined { let i = 0; const out: string[] = []; const ws = () => { while (/\s/.test(value[i] ?? "")) i++; }; const basic = { b: "\b", t: "\t", n: "\n", f: "\f", r: "\r", '"': '"', "\\": "\\" } as Record<string, string>; ws(); if (value[i++] !== "[") return undefined; for (;;) { ws(); if (value[i] === "]") { i++; ws(); return i === value.length ? out : undefined; } const quote = value[i++]; if (quote !== '"' && quote !== "'") return undefined; let s = ""; for (;;) { if (i >= value.length) return undefined; const c = value[i++]; if (c === quote) break; if (c === "\\" && quote === '"') { const e = value[i++]; if (e in basic) s += basic[e]; else if (e === "u" || e === "U") { const count = e === "u" ? 4 : 8; const hex = value.slice(i, i + count); if (!new RegExp(`^[0-9A-Fa-f]{${count}}$`).test(hex)) return undefined; const n = Number.parseInt(hex, 16); if (n > 0x10ffff || (n >= 0xd800 && n <= 0xdfff)) return undefined; s += String.fromCodePoint(n); i += count; } else return undefined; } else if (c === "\n" || c === "\r") return undefined; else s += c; } out.push(s); ws(); if (value[i] === ",") { i++; continue; } if (value[i] === "]") continue; return undefined; } }
function captureHostWindows(path: string): boolean { return path.startsWith("\\\\") || path.startsWith("//") || /^[A-Za-z]:/.test(path); }
function captureConfig(cwd: string | undefined): CaptureConfig {
const marker = findMarker(cwd);
const candidateBase = captureNormalize(cwd ? resolve(cwd) : "")?.path ?? "";
if (!marker) return { state: "inactive", patterns: [], base: candidateBase };
// Host flavor comes from the cwd STRING itself, never from `resolve(cwd)`:
// `resolve` re-roots a drive/UNC-shaped cwd through the real OS's own path
// module, which only agrees on the OS that actually owns that path. The
// cwd the hook reports is already absolute and already names its own host,
// exactly like the native hook's pure-string `flavor_of(cwd)`.
const windowsHost = captureHostWindows(cwd ?? "");
if (!marker) return { state: "inactive", patterns: [], base: candidateBase, windowsHost };
try {
const bytes = readFileSync(marker);
const markerBase = captureNormalize(dirname(marker))?.path ?? candidateBase;
if (bytes.byteLength > CAPTURE_MARKER_MAX_BYTES) return { state: "invalid", patterns: [], base: candidateBase };
if (bytes.byteLength > CAPTURE_MARKER_MAX_BYTES) return { state: "invalid", patterns: [], base: candidateBase, windowsHost };
let section = "";
let value = "";
let collecting = false;
@@ -228,26 +257,26 @@ function captureConfig(cwd: string | undefined): CaptureConfig {
if (!line) continue;
const table = /^\[([^\]]+)\]$/.exec(line);
if (table) {
if (collecting) return { state: "invalid", patterns: [], base: candidateBase };
if (collecting) return { state: "invalid", patterns: [], base: candidateBase, windowsHost };
section = table[1];
continue;
}
if (section !== "capture") continue;
if (!seen) {
const kv = /^([A-Za-z0-9_-]+)\s*=\s*(.*)$/.exec(line);
if (!kv || kv[1] !== "ignore_paths") return { state: "invalid", patterns: [], base: candidateBase };
if (!kv || kv[1] !== "ignore_paths") return { state: "invalid", patterns: [], base: candidateBase, windowsHost };
seen = true;
value = kv[2];
collecting = !value.includes("]");
} else if (collecting) {
value += ` ${line}`;
collecting = !value.includes("]");
} else return { state: "invalid", patterns: [], base: candidateBase };
} else return { state: "invalid", patterns: [], base: candidateBase, windowsHost };
}
if (collecting) return { state: "invalid", patterns: [], base: candidateBase };
if (!seen) return { state: "inactive", patterns: [], base: candidateBase };
if (collecting) return { state: "invalid", patterns: [], base: candidateBase, windowsHost };
if (!seen) return { state: "inactive", patterns: [], base: candidateBase, windowsHost };
const strings = captureParseArray(value);
if (!strings || strings.length > CAPTURE_MAX_PATTERNS) return { state: "invalid", patterns: [], base: candidateBase };
if (!strings || strings.length > CAPTURE_MAX_PATTERNS) return { state: "invalid", patterns: [], base: candidateBase, windowsHost };
const home = homedir();
const patterns = strings.map((source) => {
if (!captureValidGlob(source)) return undefined;
@@ -258,19 +287,50 @@ function captureConfig(cwd: string | undefined): CaptureConfig {
: captureJoin(markerBase, source);
const normalized = captureNormalize(expanded);
if (!normalized) return undefined;
return { path: normalized.path, windows: normalized.windows, directory: normalized.path.endsWith("/**") ? (normalized.path.slice(0, -3) || "/") : undefined };
return { path: normalized.path, windows: normalized.windows, directory: normalized.path.endsWith("/**") ? (normalized.path.slice(0, -3) || "/") : undefined, prefix: captureLiteralPrefix(normalized.path) };
});
if (patterns.some((p) => !p)) return { state: "invalid", patterns: [], base: candidateBase };
if (patterns.some((p) => !p)) return { state: "invalid", patterns: [], base: candidateBase, windowsHost };
return patterns.length
? { state: "active", patterns: patterns as CaptureConfig["patterns"], base: candidateBase }
: { state: "inactive", patterns: [], base: candidateBase };
? { state: "active", patterns: patterns as CaptureConfig["patterns"], base: candidateBase, windowsHost }
: { state: "inactive", patterns: [], base: candidateBase, windowsHost };
} catch (_e) {
return { state: "invalid", patterns: [], base: candidateBase };
return { state: "invalid", patterns: [], base: candidateBase, windowsHost };
}
}
function captureGlob(pattern: string, candidate: string, insensitive: boolean, budget: { work: number }): boolean | undefined { const p = [...pattern]; const c = [...candidate]; const eq = (a: string, b: string) => insensitive && a.charCodeAt(0) < 128 && b.charCodeAt(0) < 128 ? a.toLowerCase() === b.toLowerCase() : a === b; const previous = new Array<boolean>(p.length + 1).fill(false); previous[0] = true; for (let j = 1; j <= p.length; j++) previous[j] = p[j - 1] === "*" && p[j] !== "*" && previous[j - 1]; for (const ch of c) { const current = new Array<boolean>(p.length + 1).fill(false); for (let j = 1; j <= p.length; j++) { if (++budget.work > CAPTURE_MAX_WORK) return undefined; const x = p[j - 1]; current[j] = x === "*" && p[j] === "*" ? false : x === "*" && j >= 2 && p[j - 2] === "*" ? current[j - 2] || previous[j] : x === "*" ? current[j - 1] || (ch !== "/" && previous[j]) : x === "?" ? ch !== "/" && previous[j - 1] : eq(x, ch) && previous[j - 1]; } for (let j = 0; j <= p.length; j++) previous[j] = current[j]; } return previous[p.length]; }
function captureTool(payload: Record<string, unknown>): { family: CaptureProtocol["tool_family"]; paths?: string[]; extraction: CaptureProtocol["extraction_state"]; callID?: string } { const name = typeof payload.tool === "string" ? payload.tool.toLowerCase() : ""; const args = payload.args as Record<string, unknown> | undefined; const call = ["tool_use_id","toolUseId","tool_call_id","toolCallId","call_id","callId","callID"].map((k) => payload[k]).find((v): v is string => typeof v === "string" && /^[A-Za-z0-9_.-]{1,128}$/.test(v)); if (["search","grep","glob","find","list","ls","list_files","read_dir"].includes(name)) return { family: "search-list", extraction: "not-applicable", callID: call }; if (["bash","shell","execute","run_command","web_search"].includes(name)) return { family: "non-file", extraction: "extracted", callID: call }; if (!["read","write","edit","apply_patch","notebookedit","notebook_edit","create_file","delete_file","rename_file","move_file","multi_edit","multiedit","replace","replace_all"].includes(name)) return { family: "unknown", extraction: "extracted", callID: call }; const direct = (o: any): string[] | undefined => { if (!o || typeof o !== "object") return undefined; const r: string[] = []; for (const k of ["file_path","filePath","path","absolute_path","AbsolutePath","notebook_path"]) if (k in o) { if (typeof o[k] !== "string") return undefined; r.push(o[k]); } if ("paths" in o) { if (!Array.isArray(o.paths) || o.paths.some((x: unknown) => typeof x !== "string")) return undefined; r.push(...o.paths); } return r.length && r.length <= CAPTURE_MAX_CANDIDATES ? r : undefined; }; let paths = direct(args); if (["multi_edit","multiedit","replace_all"].includes(name)) { const entries = args?.edits ?? args?.replacements; if (!Array.isArray(entries) || !entries.length || entries.length > CAPTURE_MAX_CANDIDATES) paths = undefined; else { paths = paths ?? []; for (const entry of entries) { const more = direct(entry); if (!more || paths.length + more.length > CAPTURE_MAX_CANDIDATES) { paths = undefined; break; } paths.push(...more); } } } if (!paths || paths.some((p) => !p.trim() || [...p].length > CAPTURE_MAX_PATH_CHARS)) return { family: "file", extraction: "missing-or-malformed", callID: call }; return { family: "file", paths, extraction: "extracted", callID: call }; }
function capturePolicy(payload: Record<string, unknown>, cwd: string | undefined): { disposition: CaptureDisposition; protocol?: CaptureProtocol; payload: Record<string, unknown> } { const markerPresent = !!findMarker(cwd); if (CAPTURE_MODE === "allowlist" && !markerPresent) return { disposition: "drop", payload }; const config = captureConfig(cwd); const tool = captureTool(payload); let disposition: CaptureDisposition = "keep"; if (config.state === "invalid" && tool.family === "file") disposition = "metadata-only"; else if (config.state === "active" && tool.family === "search-list") disposition = "drop"; else if (config.state === "active" && tool.family === "file") { if (!tool.paths) disposition = "metadata-only"; else { const candidates = tool.paths.map((p) => captureNormalize(/^(?:\/|\\\\|[A-Za-z]:[\\/])/.test(p) ? p : captureJoin(config.base, p))); if (candidates.some((p) => !p)) disposition = "metadata-only"; else { const budget = { work: 0 }; captureMatch: for (const candidate of candidates as { path: string; windows: boolean }[]) for (const pattern of config.patterns) { if (candidate.windows !== pattern.windows) continue; if (pattern.directory && captureGlob(pattern.directory, candidate.path, pattern.windows, budget)) { disposition = "drop"; break captureMatch; } const match = captureGlob(pattern.path, candidate.path, pattern.windows, budget); if (match === undefined) { disposition = "metadata-only"; break; } if (match) { disposition = "drop"; break captureMatch; } } } } } if (config.state === "inactive") return { disposition, payload }; const protocol: CaptureProtocol = { version: CAPTURE_POLICY_V1, disposition, policy_state: config.state, tool_family: tool.family, path_count: tool.paths?.length ?? 0, extraction_state: tool.extraction }; if (disposition === "metadata-only") { const session = payload.sessionID ?? payload.sessionId ?? payload.session_id; const routing = typeof payload.cwd === "string" ? payload.cwd : cwd; return { disposition, protocol, payload: { ...(typeof session === "string" ? { session_id: session } : {}), ...(typeof routing === "string" ? { cwd: routing } : {}), tool_family: tool.family, tool_name: tool.family, ...(tool.callID ? { tool_call_id: tool.callID } : {}), _ai_memory_capture: protocol } }; } if (disposition === "keep") return { disposition, protocol, payload: { ...payload, _ai_memory_capture: protocol } }; return { disposition, protocol, payload }; }
function captureGlob(pattern: string, candidate: string, insensitive: boolean, budget: { work: number }): boolean | undefined { const p = [...pattern]; const c = [...candidate]; const previous = new Array<boolean>(p.length + 1).fill(false); previous[0] = true; for (let j = 1; j <= p.length; j++) previous[j] = p[j - 1] === "*" && p[j] !== "*" && previous[j - 1]; for (const ch of c) { const current = new Array<boolean>(p.length + 1).fill(false); for (let j = 1; j <= p.length; j++) { if (++budget.work > CAPTURE_MAX_WORK) return undefined; const x = p[j - 1]; current[j] = x === "*" && p[j] === "*" ? false : x === "*" && j >= 2 && p[j - 2] === "*" ? current[j - 2] || previous[j] : x === "*" ? current[j - 1] || (ch !== "/" && previous[j]) : x === "?" ? ch !== "/" && previous[j - 1] : captureCharEq(x, ch, insensitive) && previous[j - 1]; } for (let j = 0; j <= p.length; j++) previous[j] = current[j]; } return previous[p.length]; }
function captureCharEq(a: string, b: string, insensitive: boolean): boolean { return insensitive && a.charCodeAt(0) < 128 && b.charCodeAt(0) < 128 ? a.toLowerCase() === b.toLowerCase() : a === b; }
function captureLiteralPrefix(path: string): string { const glob = path.search(/[*?]/); if (glob < 0) return path; const slash = path.lastIndexOf("/", glob); if (slash < 0) return ""; const head = path.slice(0, slash + 1); const trimmed = head.slice(0, -1); return trimmed && !trimmed.endsWith(":") ? trimmed : head; }
function captureStartsWith(path: string, prefix: string, insensitive: boolean): boolean { const p = path[Symbol.iterator](); for (const x of prefix) { const c = p.next(); if (c.done || !captureCharEq(x, c.value, insensitive)) return false; } return true; }
function captureGlobReaches(glob: string, prefix: string, insensitive: boolean, budget: { work: number }): boolean | undefined { const target = prefix.replace(/\/+$/, ""); const depth = target.split("/").filter(Boolean).length; if (!depth) return true; let seen = 0; let offset = 0; for (const part of glob.split("/")) { offset += part.length; if (part && ++seen === depth) return captureGlob(glob.slice(0, offset), target, insensitive, budget); offset++; } return false; }
// An argv element is one word as given and is also tokenized on its own
// (`bash -lc "<script>"`); joining elements would re-split paths with spaces.
// Elements over 256 chars are scripts, so only their tokens count.
function captureShellCommand(args: Record<string, unknown> | undefined): string | string[] | undefined { if (!args || typeof args !== "object" || Array.isArray(args)) return undefined; const value = "command" in args ? args.command : ("cmd" in args ? args.cmd : args.CommandLine); if (typeof value === "string") return value; if (Array.isArray(value) && value.every((x) => typeof x === "string")) return value as string[]; return undefined; }
// Split only when a policy is active: this runs for every tool event.
function captureShellWordList(command: string | string[]): string[] { if (typeof command === "string") return captureShellWords(command); return command.flatMap((item) => { const tokens = captureShellWords(item); return tokens.length === 1 && tokens[0] === item || [...item].length > 256 ? tokens : [item, ...tokens]; }); }
function captureShellWords(command: string): string[] { const special = (c: string) => /\s/.test(c) || "|&;<>()".includes(c); const chars = [...command]; const words: string[] = []; let word = ""; let inWord = false; let quote = ""; for (let i = 0; i < chars.length; i++) { const c = chars[i]; const next = chars[i + 1]; if (quote) { if (c === quote) quote = ""; else if (quote === '"' && c === "\\" && (next === '"' || next === "\\")) { word += next; i++; } else word += c; } else if (c === "'" || c === '"') { quote = c; inWord = true; } else if (c === "\\" && next !== undefined && (special(next) || next === "'" || next === '"' || next === "\\")) { word += next; i++; inWord = true; } else if (special(c)) { if (inWord) words.push(word); word = ""; inWord = false; } else { word += c; inWord = true; } } if (inWord) words.push(word); return words; }
function captureShellArguments(word: string): string[] { const out = word.startsWith("-") ? [] : [word]; const eq = word.indexOf("="); if (eq >= 0) out.push(word.slice(eq + 1)); return out.filter((argument) => argument.trim() !== ""); }
// Lexical only, like the native hook: nothing is expanded or executed, so
// variables, command substitution, and `cd` state are not followed. A tool's
// own `workdir` replaces the event cwd for relative arguments.
function captureMatchCommand(command: string | string[], config: CaptureConfig, workdir?: string): boolean | undefined { const budget = { work: 0 }; const home = homedir(); const base = workdir === undefined ? config.base : /^(?:\/|\\\\|[A-Za-z]:[\\/])/.test(workdir) ? workdir : config.base && captureJoin(config.base, workdir); for (const word of captureShellWordList(command)) for (const argument of captureShellArguments(word)) { const expanded = argument.startsWith("~/") ? captureJoin(home, argument.slice(2)) : argument; if ([...expanded].length > CAPTURE_MAX_PATH_CHARS) continue; const absolute = /^(?:\/|\\\\|[A-Za-z]:[\\/])/.test(expanded); if (!absolute && !base) continue; const candidate = captureNormalize(absolute ? expanded : captureJoin(base, expanded), config.windowsHost); if (!candidate) continue; const glob = /[*?]/.test(candidate.path); for (const pattern of config.patterns) { if (candidate.windows !== pattern.windows) continue; const under = captureStartsWith(candidate.path, pattern.prefix, pattern.windows); if (!under && !glob) continue; if (under) { const directory = pattern.directory ? captureGlob(pattern.directory, candidate.path, pattern.windows, budget) : false; if (directory !== false) return directory; const match = captureGlob(pattern.path, candidate.path, pattern.windows, budget); if (match !== false) return match; } if (glob) { const reaches = captureGlobReaches(candidate.path, pattern.prefix, pattern.windows, budget); if (reaches !== false) return reaches; } } } return false; }
function captureTool(payload: Record<string, unknown>): { family: CaptureProtocol["tool_family"]; paths?: string[]; extraction: CaptureProtocol["extraction_state"]; callID?: string; command?: string | string[]; shell?: boolean; workdir?: string } { const name = typeof payload.tool === "string" ? payload.tool.toLowerCase() : ""; const args = payload.args as Record<string, unknown> | undefined; const call = ["tool_use_id","toolUseId","tool_call_id","toolCallId","call_id","callId","callID"].map((k) => payload[k]).find((v): v is string => typeof v === "string" && /^[A-Za-z0-9_.-]{1,128}$/.test(v)); if (["search","grep","glob","find","list","ls","list_files","read_dir","list_dir","grep_search","search_files","find_by_name"].includes(name)) return { family: "search-list", extraction: "not-applicable", callID: call }; if (["bash","shell","shell_command","exec","execute","run_command","web_search","search_web","manage_task","manage_subagents","terminal","execute_bash","execute_cmd"].includes(name)) return { family: "non-file", extraction: "extracted", callID: call, command: captureShellCommand(args), shell: name !== "web_search", workdir: typeof args?.workdir === "string" && args.workdir.trim() ? args.workdir : undefined }; if (!["read","write","edit","apply_patch","notebookedit","notebook_edit","create_file","delete_file","rename_file","move_file","multi_edit","multiedit","replace","replace_all"].includes(name)) return { family: "unknown", extraction: "extracted", callID: call }; const direct = (o: any): string[] | undefined => { if (!o || typeof o !== "object") return undefined; const r: string[] = []; for (const k of ["file_path","filePath","path","absolute_path","AbsolutePath","notebook_path","TargetFile"]) if (k in o) { if (typeof o[k] !== "string") return undefined; r.push(o[k]); } if ("paths" in o) { if (!Array.isArray(o.paths) || o.paths.some((x: unknown) => typeof x !== "string")) return undefined; r.push(...o.paths); } return r.length && r.length <= CAPTURE_MAX_CANDIDATES ? r : undefined; }; let paths = direct(args); if (["multi_edit","multiedit","replace_all"].includes(name)) { const entries = args?.edits ?? args?.replacements; if (!Array.isArray(entries) || !entries.length || entries.length > CAPTURE_MAX_CANDIDATES) paths = undefined; else { paths = paths ?? []; for (const entry of entries) { const more = direct(entry); if (!more || paths.length + more.length > CAPTURE_MAX_CANDIDATES) { paths = undefined; break; } paths.push(...more); } } } if (!paths || paths.some((p) => !p.trim() || [...p].length > CAPTURE_MAX_PATH_CHARS)) return { family: "file", extraction: "missing-or-malformed", callID: call }; return { family: "file", paths, extraction: "extracted", callID: call }; }
// An external lifecycle owner (`AI_MEMORY_CAPTURE_OWNER`, any value that is
// non-empty after trimming) takes over capture for this process: the gate runs
// before the capture-policy marker scan, so no disposition work, no marker read
// by this policy, and no queue write, spool write, or capture POST downstream.
// It does not reach the consumer's own routing work (`applyMarkerParams` runs
// earlier in postHook), and context delivery (`fetchHandoff`) is a separate
// path that stays live.
function captureOwnedExternally(): boolean { const owner = typeof process === "undefined" ? undefined : process.env?.AI_MEMORY_CAPTURE_OWNER; return typeof owner === "string" && owner.trim() !== ""; }
// Server profiles (`server = "<name>"`, #992) are routed only by the native
// hook. Any marker up the tree that selects one makes this integration emit
// nothing rather than deliver that repository's capture to the install default.
// Mirrors the native walk: no payload cwd falls back to the process cwd; inside
// home the walk stops at home, outside it continues past the checkout root to an
// organisation-level marker; an unreadable marker counts as a selection.
function captureServerRouted(cwd: string | undefined): boolean { let dir = resolve(cwd ?? process.cwd()); const home = homedir(); let boundary: string | undefined; for (let probe = dir; ; probe = dirname(probe)) { if (probe === home) { boundary = home; break; } if (probe === dirname(probe)) break; } for (;;) { try { if (/^\s*server\s*=/m.test(readFileSync(join(dir, ".ai-memory.toml"), "utf8"))) return true; } catch (e) { if (!["ENOENT", "ENOTDIR", "EISDIR"].includes((e as { code?: string })?.code ?? "")) return true; } if (dir === boundary || dir === dirname(dir)) return false; dir = dirname(dir); } }
function capturePolicy(payload: Record<string, unknown>, cwd: string | undefined): { disposition: CaptureDisposition; protocol?: CaptureProtocol; payload: Record<string, unknown> } { if (captureOwnedExternally()) return { disposition: "drop", payload }; if (captureServerRouted(cwd)) return { disposition: "drop", payload }; const markerPresent = !!findMarker(cwd); if (CAPTURE_MODE === "allowlist" && !markerPresent) return { disposition: "drop", payload }; const config = captureConfig(cwd); const tool = captureTool(payload); let disposition: CaptureDisposition = "keep"; if (config.state === "invalid" && (tool.family === "file" || tool.shell)) disposition = "metadata-only"; else if (config.state === "active" && tool.family === "search-list") disposition = "drop"; else if (config.state === "active" && tool.family === "file") { if (!tool.paths) disposition = "metadata-only"; else { const candidates = tool.paths.map((p) => captureNormalize(/^(?:\/|\\\\|[A-Za-z]:[\\/])/.test(p) ? p : captureJoin(config.base, p), config.windowsHost)); if (candidates.some((p) => !p)) disposition = "metadata-only"; else { const budget = { work: 0 }; captureMatch: for (const candidate of candidates as { path: string; windows: boolean }[]) for (const pattern of config.patterns) { if (candidate.windows !== pattern.windows) continue; if (pattern.directory && captureGlob(pattern.directory, candidate.path, pattern.windows, budget)) { disposition = "drop"; break captureMatch; } const match = captureGlob(pattern.path, candidate.path, pattern.windows, budget); if (match === undefined) { disposition = "metadata-only"; break; } if (match) { disposition = "drop"; break captureMatch; } } } } } else if (config.state === "active" && tool.family === "non-file" && tool.command !== undefined && captureMatchCommand(tool.command, config, tool.workdir) !== false) disposition = "drop"; if (config.state === "inactive") return { disposition, payload }; const protocol: CaptureProtocol = { version: CAPTURE_POLICY_V1, disposition, policy_state: config.state, tool_family: tool.family, path_count: tool.paths?.length ?? 0, extraction_state: tool.extraction }; if (disposition === "metadata-only") { const session = payload.sessionID ?? payload.sessionId ?? payload.session_id; const routing = typeof payload.cwd === "string" ? payload.cwd : cwd; return { disposition, protocol, payload: { ...(typeof session === "string" ? { session_id: session } : {}), ...(typeof routing === "string" ? { cwd: routing } : {}), tool_family: tool.family, tool_name: tool.family, ...(tool.callID ? { tool_call_id: tool.callID } : {}), _ai_memory_capture: protocol } }; } if (disposition === "keep") return { disposition, protocol, payload: { ...payload, _ai_memory_capture: protocol } }; return { disposition, protocol, payload }; }
"##;
TEMPLATE.replace("__AI_MEMORY_CAPTURE_MODE__", capture_mode)
}
@@ -572,6 +632,80 @@ pub(crate) const ZCODE_HOOK_TIMEOUT_MS: u64 = 10_000;
/// (same reasoning as Kiro v2's `max_output_size`).
pub(crate) const ZCODE_HOOK_MAX_OUTPUT_BYTES: usize = 64 * 1024;
/// Hermes Agent lifecycle events ai-memory hooks. Each pair is
/// `(event-name-in-~/.hermes/config.yaml, native `hook --event` value)`.
///
/// Verified against Hermes v0.21.4 (`agent/shell_hooks.py`): a configured
/// `command` is split into argv by `shlex.split` and executed **without a
/// shell**, with the event JSON on stdin — so the generated block invokes the
/// native `hook` subcommand, the same shape Zero and ZCode use, and no
/// `.sh`/`.ps1` bundle is staged. Only the two tool events are wired: their
/// payload carries `tool_name` / `tool_input`, the envelope the router already
/// maps for `agent=hermes`. Session lifecycle stays with the memory-provider
/// plugin (`on_session_end`), so a hook-driven `session-end` cannot
/// double-close a Hermes session.
pub(crate) const HERMES_EVENTS: [(&str, &str); 2] = [
("pre_tool_call", "pre-tool-use"),
("post_tool_call", "post-tool-use"),
];
/// Wall-clock bound written into each Hermes hook entry. Capture POSTs are
/// fire-and-forget, so this only bounds a hung `hook` invocation.
pub(crate) const HERMES_HOOK_TIMEOUT_SECONDS: u64 = 20;
/// The ready-to-paste `hooks:` block for `~/.hermes/config.yaml`.
///
/// `command:` must be a bare argv line — Hermes splits it itself, so the
/// `KEY=value script` prefix every shell-run harness gets would be parsed as
/// extra argv entries. The native platform already emits the argv form
/// (`<exe> [--data-dir …] hook --event … --agent hermes --server-url …`), and
/// the YAML single-quote keeps any POSIX quoting inside it intact. ai-memory
/// deliberately does not write the file: it is YAML the installer would have
/// to splice, and Hermes gates user hooks behind its own acceptance prompt.
#[must_use]
pub(crate) fn build_hermes_hooks_yaml(
server_url: &str,
auth_token: Option<&str>,
data_dir: Option<&Path>,
project_strategy: Option<&str>,
) -> String {
build_hermes_hooks_yaml_for_platform(
server_url,
auth_token,
HookCommandContext::new(
HookCommandPlatform::for_bash_runner(),
"hermes",
data_dir,
project_strategy,
),
)
}
/// Platform-forced variant, so a test can pin POSIX/Windows instead of
/// asserting whatever the machine running the suite happens to be.
fn build_hermes_hooks_yaml_for_platform(
server_url: &str,
auth_token: Option<&str>,
context: HookCommandContext<'_>,
) -> String {
let mut out = String::from("hooks:\n");
for (hermes_event, our_event) in HERMES_EVENTS {
// The native platforms derive the event token from the script stem, so
// the synthetic filename carries the ai-memory event name.
let command = hook_command(
Path::new(&format!("{our_event}.sh")),
server_url,
auth_token,
context,
);
out.push_str(&format!(
" {hermes_event}:\n - command: {}\n timeout: {HERMES_HOOK_TIMEOUT_SECONDS}\n",
yaml_single_quote(&command)
));
}
out
}
/// Devin hook payload for docker/setup-agent script snippets.
/// Devin uses HookShape::Nested (same as Claude Code/Grok) but with
/// DEVIN_EVENTS (PostCompaction instead of PreCompact, no subagent events).
@@ -1728,7 +1862,7 @@ fn to_git_bash_path(path: &str) -> String {
/// command path. Leaves only conservative shell-safe characters unquoted;
/// wraps everything else in single quotes and escapes embedded `'` via
/// `'\''`.
fn shell_quote(s: &str) -> String {
pub(crate) fn shell_quote(s: &str) -> String {
if s.chars().all(|c| {
c.is_ascii_alphanumeric()
|| matches!(c, '-' | '_' | '.' | '/' | ':' | '@' | '%' | '+' | '=' | ',')
@@ -1769,7 +1903,15 @@ fn powershell_quote(s: &str) -> String {
fn powershell_call_operator(agent: &str) -> &'static str {
// Codex evaluates `~/.codex/hooks.json` command strings with PowerShell
// on Windows (#515, reproduced against Codex CLI on a native install).
if agent == "codex" { "& " } else { "" }
// Grok Build CLI does the same for `~/.grok/hooks/*.json` (observed on
// Grok 1.0.41: every ai-memory hook failed with a ParserError in the
// session's `hook_execution` updates, and its hooks guide documents the
// PowerShell `$VAR` rewrite).
if matches!(agent, "codex" | "grok") {
"& "
} else {
""
}
}
fn win_double_quote(s: &str) -> String {
@@ -1795,6 +1937,22 @@ fn codex_windows_command_invokes_rather_than_quotes() {
);
}
/// Grok runs its Windows hooks through PowerShell like Codex, so every
/// ai-memory hook failed to parse and captured nothing.
#[test]
fn grok_windows_command_invokes_rather_than_quotes() {
let cmd = hook_command(
Path::new("stop.sh"),
"http://127.0.0.1:49374",
None,
HookCommandContext::new(HookCommandPlatform::WindowsNative, "grok", None, None),
);
assert!(
cmd.starts_with("& \""),
"grok must use the call operator: {cmd}"
);
}
#[test]
fn claude_code_windows_command_is_unchanged_by_the_codex_fix() {
let cmd = hook_command(
@@ -1860,6 +2018,25 @@ function resolveToken(): string | null {
"#
}
/// `timeoutSignal`, shared by every generated TypeScript integration: a
/// request deadline composed with `hookAbort`. A host that tears its capture
/// state down aborts `hookAbort` once its final deliveries had their drain
/// budget (OpenCode 2 on each location's unload), so whatever is still in
/// flight fails over to the spool at once instead of each waiting out its own
/// timeout. The module-level integrations never abort it.
pub(crate) fn ts_timeout_signal() -> &'static str {
r#"const hookAbort = new AbortController();
function timeoutSignal(ms: number): AbortSignal | undefined {
if (typeof AbortSignal === "undefined") return undefined;
const factory = (AbortSignal as unknown as { timeout?: (ms: number) => AbortSignal }).timeout;
const anyFactory = (AbortSignal as unknown as { any?: (signals: AbortSignal[]) => AbortSignal }).any;
if (!factory) return hookAbort.signal;
return anyFactory ? anyFactory([hookAbort.signal, factory(ms)]) : factory(ms);
}
"#
}
pub(crate) fn ts_spool_runtime() -> &'static str {
r#"
// ---- offline spool (#580): the same on-disk contract as `ai-memory hook` ----
@@ -1879,6 +2056,11 @@ function hookSpoolDir(): string {
return join(base, "hook-spool");
}
// A random prefix per copy of this state keeps two copies in one process
// (OpenCode 2 location instances, OMP and Pi loading each other's extension)
// from renaming onto each other's same-millisecond entry; the counter keeps
// one copy's entries in order.
const spoolSeqPrefix = Math.floor(Math.random() * 0x100000000).toString(16).padStart(8, "0");
let spoolSeq = 0;
function spoolFailedHook(url: URL | string, payload: Record<string, unknown>): void {
@@ -1901,7 +2083,7 @@ function spoolFailedHook(url: URL | string, payload: Record<string, unknown>): v
...(token ? { token } : {}),
attempts: 0,
};
const seq = (spoolSeq++ & 0xffffffff).toString(16).padStart(16, "0");
const seq = spoolSeqPrefix + (spoolSeq++ & 0xffffffff).toString(16).padStart(8, "0");
const name = `${String(createdMs).padStart(13, "0")}-${process.pid}-${seq}.json`;
const tmp = join(dir, `${name}.tmp`);
writeFileSync(tmp, JSON.stringify(entry), { mode: 0o600 });
@@ -1963,6 +2145,60 @@ async function drainHookSpool(): Promise<void> {
"#
}
/// `Some(())` when this machine can execute the emitted TypeScript.
/// Node is not a build dependency of this project, so a box without it
/// (or on a Node too old for type stripping) skips the runtime evidence
/// instead of failing.
#[cfg(test)]
pub(crate) fn node_strip_types_available() -> Option<()> {
let probe = std::process::Command::new("node")
.args(["--experimental-strip-types", "--version"])
.output();
match probe {
Ok(output) if output.status.success() => Some(()),
_ => {
eprintln!(
"skipping Node-required runtime evidence: node lacks --experimental-strip-types"
);
None
}
}
}
/// Two copies of the spool state in one process (OpenCode 2 location
/// instances, OMP and Pi sharing an extensions dir) must not write the
/// same file name in the same millisecond, and every request deadline
/// must compose with the state's abort signal.
#[cfg(test)]
pub(crate) fn assert_shared_ts_delivery_runtime(name: &str, source: &str) {
assert!(
source.contains("const spoolSeqPrefix = Math.floor(Math.random() * 0x100000000)"),
"{name}: spool names need a per-state random prefix"
);
assert!(
source.contains(
"const seq = spoolSeqPrefix + (spoolSeq++ & 0xffffffff).toString(16).padStart(8, \"0\");"
),
"{name}: spool seq must stay 16 hex digits"
);
assert_eq!(
source
.matches("const hookAbort = new AbortController();")
.count(),
1,
"{name}"
);
assert_eq!(
source.matches("function timeoutSignal(").count(),
1,
"{name}"
);
assert!(
source.contains("anyFactory([hookAbort.signal, factory(ms)])"),
"{name}: request deadlines must also honour hookAbort"
);
}
#[cfg(test)]
mod tests {
use super::*;
@@ -2033,27 +2269,17 @@ mod tests {
assert_eq!(h, "Bearer abc123");
}
/// Manual Node-required runtime evidence for the exact TypeScript emitted by
/// Node-required runtime evidence for the exact TypeScript emitted by
/// `ts_capture_policy_v1`. This deliberately executes the emitted source,
/// rather than maintaining a JavaScript copy in the test suite.
/// rather than maintaining a JavaScript copy in the test suite. Ignored
/// because Node is not a build requirement; the Linux CI test job installs
/// Node and runs it explicitly, so a missing Node fails instead of skipping.
#[test]
#[ignore = "manual Node-required generated TypeScript runtime evidence"]
#[ignore = "needs Node >= 22.6; CI runs it with --ignored"]
fn generated_capture_policy_v1_node_runtime_evidence() {
let strip_types = Command::new("node")
.args(["--experimental-strip-types", "--version"])
.output();
let Ok(strip_types) = strip_types else {
eprintln!(
"skipping Node-required runtime evidence: node lacks --experimental-strip-types"
);
let Some(()) = node_strip_types_available() else {
return;
};
if !strip_types.status.success() {
eprintln!(
"skipping Node-required runtime evidence: node lacks --experimental-strip-types"
);
return;
}
let temp = tempfile::tempdir().unwrap();
let module = temp.path().join("capture-policy-runtime-evidence.ts");
@@ -2099,9 +2325,39 @@ for (const vector of fixture.decisions.filter((v: any) => ["open-code", "omp", "
}}
for (const vector of fixture.normalization) {{
const cwd = marker(`[capture]\nignore_paths = [${{JSON.stringify(vector.pattern)}}]\n`);
expectDecision({{ tool: "edit", args: {{ path: vector.candidate }} }}, cwd, vector.match ? "drop" : "keep", "file", "extracted", 1, "fixture-normalization");
// Every normalization vector here uses an absolute pattern+candidate pair
// (so neither ever needs the real marker directory as a join base); the
// fixture's own `cwd` is what decides host flavor. A Windows-shaped `cwd`
// (e.g. "C:/") can't be a real directory on this (POSIX) test runner —
// `resolve("C:/")` would mangle it through the real OS path module — so
// alias the already-written marker file under the fixture's literal `cwd`
// string and pass that string straight through, exactly like the native
// test honors `vector["cwd"]` verbatim.
const hostCwd: string = /^(?:\\\\|\/\/|[A-Za-z]:)/.test(vector.cwd) ? vector.cwd : cwd;
if (hostCwd !== cwd) markerFixtures.set(hostCwd, markerFixtures.get(cwd)!);
expectDecision({{ tool: "edit", args: {{ path: vector.candidate }} }}, hostCwd, vector.match ? "drop" : "keep", "file", "extracted", 1, "fixture-normalization");
}}
expectDecision({{ tool: "edit", args: {{ path: "private/item" }} }}, marker('[capture]\nignore_paths = ["private/**"]\n'), "drop", "file", "extracted", 1, "marker-relative");
// Shell parity with the native hook's lexical command matching (#948): the
// same shared vectors `capture_policy.rs` `shell_fixture_vectors` runs.
const shellRoot = marker(`[capture]\nignore_paths = ${{JSON.stringify(fixture.shell.ignore_paths)}}\n`);
for (const vector of fixture.shell.vectors) {{
const cwd = vector.cwd ? join(shellRoot, vector.cwd) : shellRoot;
markerFixtures.set(cwd, join(shellRoot, ".ai-memory.toml"));
const payload = JSON.parse(JSON.stringify(vector.payload ?? {{ tool: "bash", args: {{ command: vector.command }} }}).replaceAll("{{root}}", shellRoot));
expectDecision(payload, cwd, vector.disposition, "non-file", "extracted", 0, `shell-${{vector.disposition}}: ${{JSON.stringify(vector)}}`);
}}
const bash = (command: string) => ({{ tool: "bash", args: {{ command }} }});
check(capturePolicy(bash("cat docs/adr/x.md"), "/no-marker").disposition === "keep", "shell-inactive-keep");
const shellInvalid = capturePolicy({{ ...bash(`cat ${{privatePath}}`), output: privateBody }}, marker('[capture'));
check(shellInvalid.disposition === "metadata-only" && shellInvalid.protocol?.policy_state === "invalid" && shellInvalid.protocol?.tool_family === "non-file", "shell-invalid-metadata-only");
for (const forbidden of [privatePath, privateBody]) check(!JSON.stringify(shellInvalid.payload).includes(forbidden), "shell-invalid-redaction");
expectDecision({{ tool: "web_search", args: {{ query: "docs/adr" }} }}, marker('[capture'), "keep", "non-file", "extracted", 0, "invalid-commandless-non-file-keep");
expectDecision({{ tool: "bash", args: {{ command: 7 }} }}, marker('[capture'), "metadata-only", "non-file", "extracted", 0, "invalid-unparseable-shell-metadata-only");
const argvScript = "echo x; ".repeat(350) + "ls a?.rs";
const manyPatterns = marker(`[capture]\nignore_paths = [${{Array.from({{ length: 40 }}, (_, i) => `"private${{i}}/**"`).join(",")}}]\n`);
expectDecision({{ tool: "shell", args: {{ command: ["bash", "-lc", argvScript] }} }}, manyPatterns, "keep", "non-file", "extracted", 0, "argv-script-not-budget-dropped");
expectDecision(bash(`cat ${{Array(8).fill("docs/x".repeat(400)).join(" ")}}`), marker(`[capture]\nignore_paths = ["${{"docs/**/?".repeat(100)}}"]\n`), "drop", "non-file", "extracted", 0, "shell-budget-fails-closed");
const nestedRoot = join(markerRoot, "nested-repo");
const nestedCwd = join(nestedRoot, "subdir");
mkdirSync(nestedCwd, {{ recursive: true }});
@@ -2113,7 +2369,8 @@ expectDecision({{ tool: "edit", args: {{ path: "private/item" }} }}, nestedCwd,
expectDecision({{ tool: "edit", args: {{ path: `${{homedir()}}/home-private/item` }} }}, marker('[capture]\nignore_paths = ["~/home-private/**"]\n'), "drop", "file", "extracted", 1, "home-expansion");
expectDecision({{ tool: "edit", args: {{ path: "case/item" }} }}, marker('[capture]\nignore_paths = ["Case/**"]\n'), "keep", "file", "extracted", 1, "posix-case");
expectDecision({{ tool: "edit", args: {{ path: "c:/SECRET/item" }} }}, marker('[capture]\nignore_paths = ["C:/secret/**"]\n'), "drop", "file", "extracted", 1, "windows-drive-case");
expectDecision({{ tool: "edit", args: {{ path: "//SERVER/SHARE/item" }} }}, marker(`[capture]\nignore_paths = ['${{String.raw`\\server\share/**`}}']\n`), "drop", "file", "extracted", 1, "windows-unc-case");
const uncCwd = "C:/"; markerFixtures.set(uncCwd, markerFixtures.get(marker(`[capture]\nignore_paths = ['${{String.raw`\\server\share/**`}}']\n`))!);
expectDecision({{ tool: "edit", args: {{ path: "//SERVER/SHARE/item" }} }}, uncCwd, "drop", "file", "extracted", 1, "windows-unc-case");
expectDecision({{ tool: "edit", args: {{ path: "x" }} }}, marker('[capture'), "metadata-only", "file", "extracted", 1, "malformed-table");
const malformedQuoteCwd = marker('[capture]\nignore_paths = ["private/**');
expectDecision({{ tool: "edit", args: {{ path: privatePath }} }}, malformedQuoteCwd, "metadata-only", "file", "extracted", 1, "malformed-unterminated-quote");
@@ -2130,6 +2387,7 @@ const requests: string[] = []; const queue: Record<string, unknown>[] = [];
function emit(payload: Record<string, unknown>, cwd: string): void {{ const result = capturePolicy(payload, cwd); if (result.disposition === "drop") return; queue.push(result.payload); requests.push(JSON.stringify(result.payload)); }}
const gatedCwd = marker(`[capture]\nignore_paths = [${{JSON.stringify(privatePattern)}}]\n`);
emit({{ tool: "edit", args: {{ path: "/PRIVATE_PATTERN_SENTINEL/item", nested: {{ body: privateBody }} }}, output: privateBody, error: privateBody }}, gatedCwd);
emit({{ tool: "bash", args: {{ command: "cat /PRIVATE_PATTERN_SENTINEL/item" }}, output: privateBody }}, gatedCwd);
check(queue.length === 0 && requests.length === 0, "drop-gates-queue-and-fetch");
emit({{ tool: "edit", args: {{ path: privatePath, nested: {{ body: privateBody }} }}, output: privateBody, error: privateBody }}, malformedQuoteCwd);
const malformedRequest = requests.at(-1) ?? "";
@@ -2151,6 +2409,20 @@ const inactive = capturePolicy(inactivePayload, "/no-marker");
check(inactive.disposition === "keep" && inactive.payload === inactivePayload && !inactive.protocol, "inactive-preserves-object");
const activeKeep = capturePolicy({{ tool: "bash", args: {{ command: privateBody }} }}, gatedCwd);
check(activeKeep.disposition === "keep" && activeKeep.protocol?.version === 1 && activeKeep.protocol.policy_state === "active", "active-keep-adds-protocol");
// #992: only the native hook routes `server` profiles; these integrations
// must drop rather than deliver a routed repository to the install default,
// including under a nested marker that does not repeat the key.
check(capturePolicy(bash("echo hi"), marker('server = "team-b"\n')).disposition === "drop", "server-routed-drops");
check(capturePolicy(bash("echo hi"), marker("server = team-b\n")).disposition === "drop", "server-routed-bare-drops");
check(capturePolicy(bash("echo hi"), marker('servers = "x"\n')).disposition === "keep", "servers-key-is-not-server");
const routedRoot = join(markerRoot, "routed-repo");
const routedChild = join(routedRoot, "child");
mkdirSync(join(routedRoot, ".git"), {{ recursive: true }});
mkdirSync(routedChild, {{ recursive: true }});
writeFileSync(join(routedRoot, ".ai-memory.toml"), 'server = "team-b"\n');
writeFileSync(join(routedChild, ".ai-memory.toml"), 'workspace = "child"\n');
markerFixtures.set(routedChild, join(routedChild, ".ai-memory.toml"));
check(capturePolicy(bash("echo hi"), routedChild).disposition === "drop", "server-routed-inherited-drops");
"#,
fixture = fixture,
policy = ts_capture_policy_v1("denylist"),
@@ -2241,6 +2513,102 @@ check(markedButEmpty.disposition === "keep", "allowlist-marker-present-empty-cap
);
}
/// `AI_MEMORY_CAPTURE_OWNER` hands capture to an external producer: the
/// generated `capturePolicy` must drop before it even looks for a marker.
/// Executed against the real emitted TypeScript rather than a hand-kept
/// copy, and cheap enough (three short Node runs, no fixtures) to stay in
/// the default tier instead of joining the manual evidence test above.
#[test]
fn generated_capture_policy_gates_on_external_capture_owner() {
let Some(()) = node_strip_types_available() else {
return;
};
// `keep()` cancels delete-on-drop: the emitted modules are the
// evidence, so they stay on disk for inspection after the run.
let temp = tempfile::tempdir().unwrap().keep();
eprintln!("capture-owner gate modules retained at {}", temp.display());
const HARNESS: &str = r#"import { closeSync, mkdirSync, openSync, readFileSync as readMarkerText, readSync, writeFileSync } from "node:fs";
import { dirname, join, resolve } from "node:path";
import { homedir } from "node:os";
const markerRoot = process.argv[2]!;
const inheritOnly = process.argv[3] === "inherit";
let markerScans = 0;
const repo = join(markerRoot, "repo");
mkdirSync(repo, { recursive: true });
const markerFile = join(repo, ".ai-memory.toml");
writeFileSync(markerFile, '[capture]\nignore_paths = ["secret/**"]\n');
function findMarker(cwd: string | undefined): string | undefined { markerScans++; return cwd === repo ? markerFile : undefined; }
__AI_MEMORY_POLICY__
function fail(label: string): never { throw new Error(`external-owner gate failed: ${label}`); }
function check(ok: unknown, label: string): asserts ok { if (!ok) fail(label); }
function probe(label: string, owned: boolean): void {
markerScans = 0;
const payload = { tool: "edit", args: { path: "public/item" } };
const result = capturePolicy(payload, repo);
if (owned) {
check(result.disposition === "drop", `${label} disposition`);
// Before the capture-policy marker scan: no disk read by the policy, no
// disposition work, and the payload comes back untouched so nothing
// downstream can queue it.
check(markerScans === 0, `${label} marker scan`);
check(result.protocol === undefined, `${label} protocol`);
check(result.payload === payload, `${label} payload identity`);
} else {
check(result.disposition === "keep", `${label} disposition`);
check(markerScans > 0, `${label} marker scan`);
check(result.protocol?.policy_state === "active", `${label} protocol`);
}
}
if (inheritOnly) {
// The value really arrived through the process environment, not a mutation
// this harness made.
check((process.env.AI_MEMORY_CAPTURE_OWNER ?? "").trim() !== "", "inherited owner missing");
probe("inherited-owner", true);
} else {
for (const [value, owned] of [[undefined, false], ["", false], [" \t\n", false], ["orchestrator-a", true], [" orchestrator-a ", true]] as [string | undefined, boolean][]) {
if (value === undefined) delete process.env.AI_MEMORY_CAPTURE_OWNER;
else process.env.AI_MEMORY_CAPTURE_OWNER = value;
probe(`owner=${JSON.stringify(value)}`, owned);
}
}
"#;
for (mode, inherit) in [("denylist", false), ("denylist", true), ("allowlist", true)] {
let module = temp.join(format!(
"capture-owner-{mode}-{}.ts",
if inherit { "inherit" } else { "matrix" }
));
fs::write(
&module,
HARNESS.replace("__AI_MEMORY_POLICY__", &ts_capture_policy_v1(mode)),
)
.unwrap();
let mut command = Command::new("node");
command.args([
"--experimental-strip-types",
module.to_str().unwrap(),
temp.to_str().unwrap(),
if inherit { "inherit" } else { "matrix" },
]);
if inherit {
command.env("AI_MEMORY_CAPTURE_OWNER", "orchestrator-a");
} else {
command.env_remove("AI_MEMORY_CAPTURE_OWNER");
}
let output = command.output().unwrap();
assert!(
output.status.success(),
"capture-owner gate evidence failed (mode={mode}, inherit={inherit})\nstdout:\n{}\nstderr:\n{}",
String::from_utf8_lossy(&output.stdout),
String::from_utf8_lossy(&output.stderr),
);
}
}
#[test]
fn claude_code_payload_has_all_events() {
let root = PathBuf::from("/host/hooks/claude-code");
@@ -2670,6 +3038,56 @@ check(markedButEmpty.disposition === "keep", "allowlist-marker-present-empty-cap
assert!(args.iter().all(|arg| !arg.contains(r"\\?\")));
}
/// The public entry point must emit a complete `hooks:` map with both tool
/// events, whatever platform it renders for (the suite runs on Windows too).
#[test]
fn hermes_hooks_yaml_covers_both_tool_events() {
let yaml = build_hermes_hooks_yaml("http://127.0.0.1:49374", None, None, None);
assert!(yaml.starts_with("hooks:\n"));
for (event, _) in HERMES_EVENTS {
assert!(yaml.contains(&format!(" {event}:\n")), "missing {event}");
assert!(
yaml.contains(" - command: '"),
"command must be a single-quoted YAML scalar"
);
}
assert_eq!(
yaml.matches(&format!("timeout: {HERMES_HOOK_TIMEOUT_SECONDS}"))
.count(),
HERMES_EVENTS.len()
);
}
/// Hermes splits `command` with `shlex.split` and runs it with no shell, so
/// the generated line must be bare argv: a `KEY=value` prefix would be
/// parsed as extra argv entries and the hook would never reach the server.
#[test]
fn hermes_hooks_yaml_invokes_the_native_command_without_a_shell() {
let yaml = build_hermes_hooks_yaml_for_platform(
"http://127.0.0.1:49374",
None,
HookCommandContext::new(
HookCommandPlatform::PosixNative,
"hermes",
Some(Path::new("/data")),
Some("repo-root"),
),
);
for (_, our_event) in HERMES_EVENTS {
assert!(
yaml.contains(&format!("hook --event {our_event} --agent hermes")),
"missing the native command for {our_event}"
);
}
assert!(
!yaml.contains("AI_MEMORY_HOOK_URL="),
"an env prefix would be parsed as argv, not environment"
);
assert!(yaml.contains("--data-dir /data"));
assert!(yaml.contains("--server-url http://127.0.0.1:49374"));
assert!(yaml.contains("--project-strategy repo-root"));
}
fn exec_args(script: &str, capture_assistant: bool) -> Vec<String> {
let spec = windows_native_exec_spec_with_exe(
Path::new(r"C:\ai-memory\ai-memory.exe"),
crates/ai-memory-cli/src/commands/repair_backfill_timestamps.rs
+393
View File
@@ -0,0 +1,393 @@
//! `ai-memory repair-backfill-timestamps` — thin HTTP client that corrects
//! `sessions.started_at`/`ended_at` for sessions an older `backfill` already
//! imported before it carried the transcript's own event time, flattening
//! every imported session onto the import day.
//!
//! Like every other lifecycle command, this is a thin client: the server
//! (`POST /admin/repair-session-times`) owns validation (including the guard
//! that only a genuinely flattened session is ever rewritten) and the write,
//! one transaction per request through the single `WriterHandle`. This
//! command's only job is the part that must run on the operator's machine —
//! reading the local harness transcripts, which the server never has access
//! to — and it reuses `backfill`'s own discovery (`collect_local_sessions`)
//! and the `ai-memory-workstream` transcript reader (`export_transcript`)
//! rather than re-parsing transcript formats.
//!
//! Dry-run by default: the server validates and reports without writing.
//! `--confirm` applies. See `docs/lifecycle-ops.md` for the full validation
//! list.
use std::collections::BTreeMap;
use std::path::Path;
use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use ai_memory_core::SessionId;
use ai_memory_workstream::{ManagedHarness, build_launch_plan, export_transcript};
use super::backfill::{SessionRef, collect_local_sessions};
use super::run;
use crate::cli::RepairBackfillTimestampsArgs;
use crate::config::Config;
use crate::http_client::{ServerEndpoint, post_json};
/// One request never carries more candidates than the server's
/// `MAX_REPAIR_SESSIONS` (`ai-memory-mcp/src/admin.rs`) accepts; a larger
/// local batch is split into several requests, each its own transaction.
const MAX_SESSIONS_PER_REQUEST: usize = 2_000;
/// One candidate posted to `POST /admin/repair-session-times`.
#[derive(Debug, Clone, Serialize)]
struct SessionTimesItem {
session_id: String,
started_at_us: i64,
#[serde(skip_serializing_if = "Option::is_none")]
ended_at_us: Option<i64>,
}
#[derive(Debug, Serialize)]
struct RepairSessionTimesRequest {
workspace: String,
project: String,
sessions: Vec<SessionTimesItem>,
confirm: bool,
}
/// The server's response, kept in full (not just a summarized subset) so
/// `--json` can pass it straight through and the human report can print
/// concrete old/new values.
#[derive(Debug, Deserialize, Serialize)]
struct RepairSessionTimesResponse {
dry_run: bool,
repaired: Vec<RepairedItem>,
skipped: Vec<SkippedItem>,
}
#[derive(Debug, Deserialize, Serialize)]
struct RepairedItem {
session_id: String,
old_started_at_us: i64,
#[serde(default)]
old_ended_at_us: Option<i64>,
new_started_at_us: i64,
#[serde(default)]
new_ended_at_us: Option<i64>,
#[serde(default)]
end_kept_open: bool,
}
#[derive(Debug, Deserialize, Serialize)]
struct SkippedItem {
session_id: String,
reason: String,
}
/// First/last event timestamp of one local session's transcript, in Unix
/// microseconds. `None` when the transcript could not be read or carried no
/// parseable `occurred_at` on any event.
async fn session_time_span(
home: &Path,
cwd: &Path,
session_dir: Option<&Path>,
session: &SessionRef,
) -> Option<(i64, i64)> {
let transcript = export_transcript(
session.harness,
home,
cwd,
session_dir,
&session.native_session_id,
None,
)
.await
.ok()?;
let mut first_us: Option<i64> = None;
let mut last_us: Option<i64> = None;
for event in &transcript.events {
let Some(us) = event
.occurred_at
.as_deref()
.and_then(|s| s.parse::<jiff::Timestamp>().ok())
.map(jiff::Timestamp::as_microsecond)
else {
continue;
};
first_us = Some(first_us.map_or(us, |f: i64| f.min(us)));
last_us = Some(last_us.map_or(us, |l: i64| l.max(us)));
}
Some((first_us?, last_us?))
}
fn format_us(us: i64) -> String {
jiff::Timestamp::from_microsecond(us)
.map(|t| t.to_string())
.unwrap_or_else(|_| us.to_string())
}
/// Build one candidate per local session with a readable transcript
/// timestamp. `session_dir` is resolved once per harness (`build_launch_plan`
/// only depends on the harness, not the session), not once per session.
async fn build_candidates(
home: &Path,
cwd: &Path,
sessions: &[SessionRef],
) -> (Vec<SessionTimesItem>, usize) {
// A small, fixed set of harnesses: a linear scan avoids requiring `Hash`
// on `ManagedHarness` for a handful of entries.
let mut session_dirs: Vec<(ManagedHarness, Option<std::path::PathBuf>)> = Vec::new();
let mut items = Vec::with_capacity(sessions.len());
let mut uncaptured = 0usize;
for session in sessions {
let session_dir = match session_dirs.iter().find(|(h, _)| *h == session.harness) {
Some((_, dir)) => dir.clone(),
None => {
let dir = build_launch_plan(session.harness, None, Vec::new(), None)
.ok()
.and_then(|plan| plan.session_dir);
session_dirs.push((session.harness, dir.clone()));
dir
}
};
let Some((first_us, last_us)) =
session_time_span(home, cwd, session_dir.as_deref(), session).await
else {
uncaptured += 1;
continue;
};
let session_id = SessionId::from_native(&session.native_session_id);
items.push(SessionTimesItem {
session_id: session_id.to_string(),
started_at_us: first_us,
ended_at_us: Some(last_us),
});
}
(items, uncaptured)
}
/// Run the `repair-backfill-timestamps` subcommand.
///
/// # Errors
/// Returns an error when the scope cannot be resolved, the local harness
/// session stores cannot be located, or the server is unreachable or answers
/// non-2xx.
pub async fn run(config: &Config, args: RepairBackfillTimestampsArgs) -> Result<()> {
let (workspace, project) =
super::resolve_scope(config, args.workspace.as_deref(), args.project.as_deref())?;
let cwd = std::env::current_dir().context("resolving the current working directory")?;
let home = run::native_home(config).context("locating the local harness session stores")?;
// Same discovery `backfill` uses to select what it would import: every
// local native session (across every scanned harness) whose recorded cwd
// matches this checkout. Sharing it means this command looks at exactly
// the sessions `backfill` itself would have found.
let (sessions, limit_hit) = collect_local_sessions(&home, &cwd, None).await;
let (items, uncaptured) = build_candidates(&home, &cwd, &sessions).await;
if items.is_empty() {
println!(
"ai-memory: repair-backfill-timestamps for {workspace}/{project}: no local \
transcript carried a readable timestamp ({uncaptured} session(s) scanned locally, \
none captured); nothing to send."
);
print_scan_limit_note(&limit_hit);
return Ok(());
}
let endpoint = ServerEndpoint::from_config_resolving_auth(config).await;
let mut responses = Vec::new();
for chunk in items.chunks(MAX_SESSIONS_PER_REQUEST) {
let response: RepairSessionTimesResponse = post_json(
&endpoint,
"/admin/repair-session-times",
&RepairSessionTimesRequest {
workspace: workspace.clone(),
project: project.clone(),
sessions: chunk.to_vec(),
confirm: args.confirm,
},
)
.await?;
responses.push(response);
}
if args.json {
println!("{}", serde_json::to_string_pretty(&responses)?);
return Ok(());
}
let dry_run = responses.first().is_none_or(|r| r.dry_run);
let total_repaired: usize = responses.iter().map(|r| r.repaired.len()).sum();
let total_skipped: usize = responses.iter().map(|r| r.skipped.len()).sum();
let mut skipped_by_reason: BTreeMap<&str, usize> = BTreeMap::new();
for response in &responses {
for skip in &response.skipped {
*skipped_by_reason.entry(skip.reason.as_str()).or_default() += 1;
}
}
println!(
"ai-memory: repair-backfill-timestamps for {workspace}/{project}: {total_repaired} \
session(s) {}, {total_skipped} skipped, {uncaptured} local transcript(s) had no \
readable timestamp.",
if dry_run {
"would be repaired"
} else {
"repaired"
},
);
if responses.len() > 1 {
println!(
" sent {} session(s) across {} request(s) of up to {MAX_SESSIONS_PER_REQUEST} each \
(one transaction per request).",
items.len(),
responses.len(),
);
}
for (reason, count) in &skipped_by_reason {
println!(" skipped ({reason}): {count}");
}
print_scan_limit_note(&limit_hit);
for response in &responses {
for repaired in &response.repaired {
println!(
" {}: {} .. {} -> {} .. {}{}",
repaired.session_id,
format_us(repaired.old_started_at_us),
repaired
.old_ended_at_us
.map_or_else(|| "open".to_string(), format_us),
format_us(repaired.new_started_at_us),
repaired
.new_ended_at_us
.map_or_else(|| "open".to_string(), format_us),
if repaired.end_kept_open {
" (end left open)"
} else {
""
},
);
}
}
if dry_run && total_repaired > 0 {
println!(
" dry run: nothing was written. Re-run with --confirm to apply (recorded in \
audit_log on apply, which is the reversibility backing)."
);
}
Ok(())
}
/// A harness whose local scan came back at the cap means older sessions may
/// exist that were not even considered — surface that instead of silently
/// under-reporting.
fn print_scan_limit_note(limit_hit: &[ManagedHarness]) {
for harness in limit_hit {
println!(
" note: {} returned the maximum scanned sessions for this checkout; older \
sessions may exist and were not considered.",
harness.as_str()
);
}
}
#[cfg(test)]
mod tests {
use super::*;
/// `--json` aside, the human report reads real server-supplied old/new
/// values, not values this command invented — this is a compile-time
/// shape check on the (de)serialization, not an HTTP test (the HTTP path
/// is covered by `admin_repair_session_times.rs` on the server side).
#[test]
fn repair_session_times_response_round_trips_through_json() {
let response = RepairSessionTimesResponse {
dry_run: true,
repaired: vec![RepairedItem {
session_id: "11111111-2222-3333-4444-555555555555".into(),
old_started_at_us: 2,
old_ended_at_us: Some(3),
new_started_at_us: 0,
new_ended_at_us: Some(1),
end_kept_open: false,
}],
skipped: vec![SkippedItem {
session_id: "22222222-3333-4444-5555-666666666666".into(),
reason: "not_flattened".into(),
}],
};
let json = serde_json::to_string(&response).unwrap();
let back: RepairSessionTimesResponse = serde_json::from_str(&json).unwrap();
assert_eq!(back.repaired[0].new_started_at_us, 0);
assert_eq!(back.skipped[0].reason, "not_flattened");
}
/// End-to-end proof that this command's own candidate-building reads a
/// planted transcript's real event timestamps (not the file's mtime or
/// any other stand-in), mirroring
/// `backfill::tests::collect_local_sessions_finds_a_planted_claude_session_for_this_cwd`.
/// Unix-gated for the same reason that test is: the fixture uses a
/// POSIX-encoded `~/.claude/projects/<enc-cwd>/` layout.
#[cfg(unix)]
#[tokio::test]
async fn build_candidates_reads_the_planted_transcripts_own_timestamps() {
let home = tempfile::tempdir().unwrap();
let cwd = tempfile::tempdir().unwrap();
let session_dir = home
.path()
.join(".claude")
.join("projects")
.join(cwd.path().to_string_lossy().replace('/', "-"));
std::fs::create_dir_all(&session_dir).unwrap();
let native_id = "11111111-2222-3333-4444-555555555555";
let header =
serde_json::json!({"sessionId": native_id, "cwd": cwd.path().to_string_lossy()});
let first = serde_json::json!({
"type": "user",
"message": {"role": "user", "content": "hello"},
"timestamp": "2020-01-01T00:00:00Z",
});
let last = serde_json::json!({
"type": "assistant",
"message": {"role": "assistant", "content": "hi back"},
"timestamp": "2020-01-01T00:05:00Z",
});
std::fs::write(
session_dir.join("sess.jsonl"),
format!("{header}\n{first}\n{last}\n"),
)
.unwrap();
let (sessions, limit_hit) = crate::commands::backfill::collect_local_sessions_with(
home.path(),
cwd.path(),
None,
|_| None,
)
.await;
assert!(limit_hit.is_empty(), "{limit_hit:?}");
let (items, uncaptured) = build_candidates(home.path(), cwd.path(), &sessions).await;
assert_eq!(uncaptured, 0, "{items:?}");
let claude_item = items
.iter()
.find(|item| item.session_id == native_id)
.unwrap_or_else(|| panic!("no candidate for {native_id} in {items:?}"));
assert_eq!(
claude_item.started_at_us,
"2020-01-01T00:00:00Z"
.parse::<jiff::Timestamp>()
.unwrap()
.as_microsecond()
);
assert_eq!(
claude_item.ended_at_us,
Some(
"2020-01-01T00:05:00Z"
.parse::<jiff::Timestamp>()
.unwrap()
.as_microsecond()
)
);
}
}
crates/ai-memory-cli/src/commands/restore.rs
+469 -36
View File
@@ -1,9 +1,17 @@
//! `ai-memory restore --from <tarball>` — restore a backup tarball.
//!
//! Refuses to overwrite a non-empty data dir unless `--force` is given.
//! Refuses while another `ai-memory` process is alive. After extraction,
//! re-opens the store so any pending migrations run (and a corrupt
//! snapshot fails loudly).
//! Refuses while another `ai-memory` process is alive.
//!
//! The live data is never touched before the archive has proven usable:
//! the tarball is extracted and validated into a staging directory beside
//! the live `wiki/` and `db/`, the restored store is opened there so any
//! pending migrations run (and a corrupt snapshot fails loudly), and only
//! then are the live directories swapped out by rename. A truncated
//! archive, an entry outside the allowed layout, or a snapshot the current
//! binary cannot open therefore leaves the existing data exactly as it was
//! — the moment a restore fails is the moment the operator has no other
//! copy, so the previous state must survive it.
//!
//! # Exception to invariant §16
//!
@@ -18,18 +26,29 @@ use ai_memory_store::Store;
use anyhow::{Context, Result, bail};
use flate2::read::GzDecoder;
use std::path::{Component, Path};
use tracing::info;
use tracing::{info, warn};
use crate::cli::RestoreArgs;
use crate::config::Config;
use crate::process_guard::{busy_message, sibling_processes};
/// Data-dir directories a restore replaces wholesale: whatever the archive
/// holds for each takes the live one's place, and a directory the archive
/// lacks is retired rather than merged with the archive's state. `logs/`,
/// `models/`, `raw/` and anything else beside them are never touched.
const REPLACED_DIRS: &[&str] = &["wiki", "db"];
/// The one file a restore replaces, and only when the archive carries it;
/// otherwise the live copy stays.
const CONFIG_FILE: &str = "config.toml";
/// Run the `restore` subcommand.
///
/// # Errors
/// Returns an error if another `ai-memory` process is running, the
/// data dir is non-empty without `--force`, the tarball cannot be
/// extracted, or the restored store fails to open.
/// extracted, or the restored store fails to open. In every one of those
/// cases the live data dir is left as it was.
pub fn run(config: &Config, args: RestoreArgs) -> Result<()> {
let siblings = sibling_processes();
if !siblings.is_empty() {
@@ -40,37 +59,7 @@ pub fn run(config: &Config, args: RestoreArgs) -> Result<()> {
bail!("source tarball {} not found", args.from.display());
}
let wiki = config.data_dir.join("wiki");
let db = config.data_dir.join("db").join("memory.sqlite");
if (wiki.is_dir() && std::fs::read_dir(&wiki)?.next().is_some()) || db.is_file() {
if !args.force {
bail!(
"refusing to restore: data dir at {} is non-empty (pass --force to overwrite)",
config.data_dir.display(),
);
}
// Force path: drop the existing wiki + db so the tarball can
// populate them cleanly. Keep config.toml, logs/, models/.
for sub in ["wiki", "db"] {
let path = config.data_dir.join(sub);
if path.exists() {
std::fs::remove_dir_all(&path)?;
}
}
}
std::fs::create_dir_all(&config.data_dir)?;
let file = std::fs::File::open(&args.from)
.with_context(|| format!("opening {}", args.from.display()))?;
let decoder = GzDecoder::new(file);
let mut archive = tar::Archive::new(decoder);
unpack_checked_archive(&mut archive, &config.data_dir)
.with_context(|| format!("extracting into {}", config.data_dir.display()))?;
info!(from = %args.from.display(), into = %config.data_dir.display(), "tarball extracted");
// Open + drop the store so refinery applies any pending migrations
// and the SQLite file is validated.
let _store = Store::open(&config.data_dir).context("opening restored store")?;
restore_data_dir(&config.data_dir, &args.from, args.force)?;
info!("restore complete");
println!(
"restored {} -> {}",
@@ -80,6 +69,170 @@ pub fn run(config: &Config, args: RestoreArgs) -> Result<()> {
Ok(())
}
/// Stage, validate, then swap. Nothing under `data_dir` changes until the
/// archive has been fully extracted into a staging directory and the
/// staged store has opened; the live `wiki/`, `db/` and (when the archive
/// carries one) `config.toml` are then exchanged for the staged copies by
/// rename and rolled back if any step of the exchange fails.
fn restore_data_dir(data_dir: &Path, from: &Path, force: bool) -> Result<()> {
let wiki = data_dir.join("wiki");
let db = data_dir.join("db").join("memory.sqlite");
let occupied = (wiki.is_dir() && std::fs::read_dir(&wiki)?.next().is_some()) || db.is_file();
if occupied && !force {
bail!(
"refusing to restore: data dir at {} is non-empty (pass --force to overwrite)",
data_dir.display(),
);
}
std::fs::create_dir_all(data_dir)?;
// Both scratch directories live inside the data dir so every move below
// is a rename on one filesystem, never a copy that could half-complete.
let stamp = format!(
"{}-{}",
jiff::Timestamp::now().strftime("%Y%m%d-%H%M%S"),
std::process::id()
);
let staging = data_dir.join(format!(".restore-staging-{stamp}"));
let previous = data_dir.join(format!(".restore-previous-{stamp}"));
let outcome = stage_then_swap(data_dir, from, &staging, &previous);
// The staging dir is scratch in every outcome: on success its contents
// were moved into place, on failure the live data was never touched.
if staging.exists()
&& let Err(e) = std::fs::remove_dir_all(&staging)
{
warn!(path = %staging.display(), error = %e, "could not remove restore staging dir");
eprintln!(
"warning: could not remove staging dir {} ({e}); delete it by hand",
staging.display()
);
}
outcome
}
fn stage_then_swap(data_dir: &Path, from: &Path, staging: &Path, previous: &Path) -> Result<()> {
std::fs::create_dir(staging)
.with_context(|| format!("creating staging dir {}", staging.display()))?;
// 1. Extract into staging, validating every entry on the way. A
// truncated gzip stream, an unreadable member or a path outside the
// allowed layout fails here, with the live data still untouched.
let file = std::fs::File::open(from).with_context(|| format!("opening {}", from.display()))?;
let decoder = GzDecoder::new(file);
let mut archive = tar::Archive::new(decoder);
unpack_checked_archive(&mut archive, staging)
.with_context(|| format!("extracting {} into {}", from.display(), staging.display()))?;
info!(from = %from.display(), into = %staging.display(), "tarball extracted into staging");
// 2. Open + drop the staged store so refinery applies any pending
// migrations and the SQLite file is validated — still before anything
// live is touched. Dropping the store joins the writer thread and
// closes every connection, so the directory can be renamed afterwards
// on Windows as well.
drop(Store::open(staging).context("opening restored store")?);
// 3. Exchange the live directories for the staged ones.
swap_into_place(data_dir, staging, previous)?;
info!(into = %data_dir.display(), "restored data swapped into place");
Ok(())
}
/// Move the live entries aside into `previous`, move the staged entries
/// into place, then discard `previous`. Every step is a same-filesystem
/// rename; if one fails, the moves already made are reversed so the data
/// dir ends up as it started.
fn swap_into_place(data_dir: &Path, staging: &Path, previous: &Path) -> Result<()> {
std::fs::create_dir(previous).with_context(|| format!("creating {}", previous.display()))?;
// Entries moved from the data dir into `previous`, and staged entries
// already placed live: the two lists a rollback has to undo.
let mut moved_aside: Vec<&str> = Vec::new();
let mut placed: Vec<&str> = Vec::new();
let exchange = (|| -> Result<()> {
for name in REPLACED_DIRS {
let live = data_dir.join(name);
if live.exists() {
std::fs::rename(&live, previous.join(name))
.with_context(|| format!("moving {} aside", live.display()))?;
moved_aside.push(name);
}
}
let live_config = data_dir.join(CONFIG_FILE);
if staging.join(CONFIG_FILE).is_file() && live_config.exists() {
std::fs::rename(&live_config, previous.join(CONFIG_FILE))
.with_context(|| format!("moving {} aside", live_config.display()))?;
moved_aside.push(CONFIG_FILE);
}
for name in REPLACED_DIRS.iter().chain(std::iter::once(&CONFIG_FILE)) {
let staged = staging.join(name);
if staged.exists() {
std::fs::rename(&staged, data_dir.join(name))
.with_context(|| format!("moving {} into place", staged.display()))?;
placed.push(name);
}
}
Ok(())
})();
if let Err(e) = exchange {
if let Err(rollback_err) = roll_back(data_dir, previous, &placed, &moved_aside) {
bail!(
"INCONSISTENT STATE: restore swap failed ({e:#}) and moving the previous data \
back also failed ({rollback_err:#}); the pre-restore wiki/ and db/ are under {} \
— move them back into {} by hand",
previous.display(),
data_dir.display(),
);
}
return Err(e.context("restore swap failed; the previous data was moved back into place"));
}
// The previous data is only discarded once the restored copy is live —
// the documented `--force` semantics, now applied last instead of first.
if let Err(e) = std::fs::remove_dir_all(previous) {
warn!(path = %previous.display(), error = %e, "could not remove pre-restore data");
eprintln!(
"warning: restore succeeded but the pre-restore data under {} could not be \
removed ({e}); delete it by hand",
previous.display()
);
}
Ok(())
}
/// Undo a partial swap: remove whatever staged entries were already placed
/// live, then move the previous entries back. Placed entries are removed
/// before their predecessors return so a rename never finds its target
/// occupied.
fn roll_back(
data_dir: &Path,
previous: &Path,
placed: &[&str],
moved_aside: &[&str],
) -> Result<()> {
for name in placed {
let live = data_dir.join(name);
remove_path(&live).with_context(|| format!("removing half-placed {}", live.display()))?;
}
for name in moved_aside {
let parked = previous.join(name);
std::fs::rename(&parked, data_dir.join(name))
.with_context(|| format!("moving {} back", parked.display()))?;
}
Ok(())
}
fn remove_path(path: &Path) -> std::io::Result<()> {
match std::fs::symlink_metadata(path) {
Ok(meta) if meta.is_dir() => std::fs::remove_dir_all(path),
Ok(_) => std::fs::remove_file(path),
Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()),
Err(e) => Err(e),
}
}
fn unpack_checked_archive<R: std::io::Read>(
archive: &mut tar::Archive<R>,
data_dir: &Path,
@@ -354,4 +507,284 @@ mod tests {
.is_file()
);
}
// ---- stage → validate → swap: a failed restore leaves the live data alone ----
/// A populated data dir as an operator has it: a wiki page, a database
/// file (never opened by these tests, so any bytes do), the config, and
/// the neighbours `logs/` and `raw/` that a restore must never touch.
fn seed_live_data(dir: &Path) {
std::fs::create_dir_all(dir.join("wiki/default/project/notes")).unwrap();
std::fs::write(dir.join("wiki/default/project/notes/old.md"), b"old page").unwrap();
std::fs::create_dir_all(dir.join("db")).unwrap();
std::fs::write(dir.join("db/memory.sqlite"), b"old database bytes").unwrap();
std::fs::write(dir.join("config.toml"), b"# old config\n").unwrap();
std::fs::create_dir_all(dir.join("logs")).unwrap();
std::fs::write(dir.join("logs/app.log"), b"log").unwrap();
std::fs::create_dir_all(dir.join("raw")).unwrap();
std::fs::write(dir.join("raw/segment.jsonl"), b"{}").unwrap();
}
/// Scratch directories a restore leaves behind only when its cleanup
/// failed: none may survive a run, successful or not.
fn restore_scratch_dirs(data_dir: &Path) -> Vec<std::path::PathBuf> {
std::fs::read_dir(data_dir)
.unwrap()
.filter_map(Result::ok)
.map(|e| e.path())
.filter(|p| {
p.file_name()
.and_then(|n| n.to_str())
.is_some_and(|n| n.starts_with(".restore-"))
})
.collect()
}
fn assert_live_data_untouched(dir: &Path) {
assert_eq!(
std::fs::read(dir.join("wiki/default/project/notes/old.md")).unwrap(),
b"old page"
);
assert_eq!(
std::fs::read(dir.join("db/memory.sqlite")).unwrap(),
b"old database bytes"
);
assert_eq!(
std::fs::read(dir.join("config.toml")).unwrap(),
b"# old config\n"
);
assert_eq!(std::fs::read(dir.join("logs/app.log")).unwrap(), b"log");
assert_eq!(std::fs::read(dir.join("raw/segment.jsonl")).unwrap(), b"{}");
let scratch = restore_scratch_dirs(dir);
assert!(scratch.is_empty(), "scratch dirs left behind: {scratch:?}");
}
/// Bytes of a real, migrated `memory.sqlite` — what a `backup` tarball
/// carries.
fn migrated_sqlite_bytes() -> Vec<u8> {
let tmp = tempfile::TempDir::new().unwrap();
drop(Store::open(tmp.path()).unwrap());
std::fs::read(tmp.path().join("db/memory.sqlite")).unwrap()
}
/// A gzipped tarball holding the given regular files.
fn gz_tarball(entries: &[(&str, &[u8])]) -> Vec<u8> {
use std::io::Write as _;
let mut tar_bytes = Vec::new();
{
let mut builder = tar::Builder::new(&mut tar_bytes);
for (path, body) in entries {
let mut header = tar::Header::new_gnu();
header.set_path(path).unwrap();
header.set_size(body.len() as u64);
header.set_mode(0o644);
header.set_cksum();
builder.append(&header, *body).unwrap();
}
builder.finish().unwrap();
}
let mut gz = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
gz.write_all(&tar_bytes).unwrap();
gz.finish().unwrap()
}
fn good_archive() -> Vec<u8> {
let db = migrated_sqlite_bytes();
gz_tarball(&[
("wiki/default/project/notes/new.md", b"new page"),
("db/memory.sqlite", db.as_slice()),
("config.toml", b"# restored config\n"),
])
}
fn write_tarball(dir: &Path, name: &str, bytes: &[u8]) -> std::path::PathBuf {
let path = dir.join(name);
std::fs::write(&path, bytes).unwrap();
path
}
#[test]
fn restore_force_keeps_live_data_when_the_tarball_is_not_a_gzip_stream() {
let data = tempfile::TempDir::new().unwrap();
seed_live_data(data.path());
let scratch = tempfile::TempDir::new().unwrap();
let from = write_tarball(scratch.path(), "bad.tar.gz", b"this is not a gzip stream");
let err = restore_data_dir(data.path(), &from, true).unwrap_err();
assert!(
format!("{err:#}").contains("extracting"),
"unexpected error: {err:#}"
);
assert_live_data_untouched(data.path());
}
#[test]
fn restore_force_keeps_live_data_when_the_tarball_is_truncated() {
let data = tempfile::TempDir::new().unwrap();
seed_live_data(data.path());
let scratch = tempfile::TempDir::new().unwrap();
let whole = good_archive();
let from = write_tarball(scratch.path(), "cut.tar.gz", &whole[..whole.len() / 2]);
let err = restore_data_dir(data.path(), &from, true).unwrap_err();
assert!(
format!("{err:#}").contains("extracting"),
"unexpected error: {err:#}"
);
assert_live_data_untouched(data.path());
}
#[test]
fn restore_force_keeps_live_data_when_an_entry_is_outside_the_allowed_layout() {
let data = tempfile::TempDir::new().unwrap();
seed_live_data(data.path());
let scratch = tempfile::TempDir::new().unwrap();
// Valid entries first, so extraction is already under way when the
// stray one is met — the failure must still leave nothing behind.
let db = migrated_sqlite_bytes();
let from = write_tarball(
scratch.path(),
"stray.tar.gz",
&gz_tarball(&[
("wiki/default/project/notes/new.md", b"new page"),
("db/memory.sqlite", db.as_slice()),
("db/extra.sqlite", b"stray"),
]),
);
let err = restore_data_dir(data.path(), &from, true).unwrap_err();
assert!(
format!("{err:#}").contains("unexpected path"),
"unexpected error: {err:#}"
);
assert_live_data_untouched(data.path());
}
#[test]
fn restore_force_keeps_live_data_when_the_restored_store_cannot_open() {
let data = tempfile::TempDir::new().unwrap();
seed_live_data(data.path());
let scratch = tempfile::TempDir::new().unwrap();
let from = write_tarball(
scratch.path(),
"notadb.tar.gz",
&gz_tarball(&[
("wiki/default/project/notes/new.md", b"new page"),
("db/memory.sqlite", &[0xFF; 4096]),
]),
);
let err = restore_data_dir(data.path(), &from, true).unwrap_err();
assert!(
format!("{err:#}").contains("opening restored store"),
"unexpected error: {err:#}"
);
assert_live_data_untouched(data.path());
}
#[test]
fn restore_refuses_a_populated_data_dir_without_force() {
let data = tempfile::TempDir::new().unwrap();
seed_live_data(data.path());
let scratch = tempfile::TempDir::new().unwrap();
let from = write_tarball(scratch.path(), "good.tar.gz", &good_archive());
let err = restore_data_dir(data.path(), &from, false).unwrap_err();
assert!(
err.to_string().contains("pass --force"),
"unexpected error: {err:#}"
);
assert_live_data_untouched(data.path());
}
#[test]
fn restore_force_replaces_the_live_data_with_the_archive() {
let data = tempfile::TempDir::new().unwrap();
seed_live_data(data.path());
let scratch = tempfile::TempDir::new().unwrap();
let from = write_tarball(scratch.path(), "good.tar.gz", &good_archive());
restore_data_dir(data.path(), &from, true).unwrap();
assert_eq!(
std::fs::read(data.path().join("wiki/default/project/notes/new.md")).unwrap(),
b"new page"
);
assert!(
!data
.path()
.join("wiki/default/project/notes/old.md")
.exists(),
"the previous wiki must not be merged into the restored one"
);
assert_eq!(
std::fs::read(data.path().join("config.toml")).unwrap(),
b"# restored config\n"
);
assert_eq!(
std::fs::read(data.path().join("logs/app.log")).unwrap(),
b"log"
);
assert_eq!(
std::fs::read(data.path().join("raw/segment.jsonl")).unwrap(),
b"{}"
);
let scratch_dirs = restore_scratch_dirs(data.path());
assert!(
scratch_dirs.is_empty(),
"scratch dirs left behind: {scratch_dirs:?}"
);
// The swapped-in store is the migrated one and opens cleanly in place.
drop(Store::open(data.path()).unwrap());
}
#[test]
fn restore_keeps_the_live_config_when_the_archive_carries_none() {
let data = tempfile::TempDir::new().unwrap();
seed_live_data(data.path());
let scratch = tempfile::TempDir::new().unwrap();
let db = migrated_sqlite_bytes();
let from = write_tarball(
scratch.path(),
"noconfig.tar.gz",
&gz_tarball(&[
("wiki/default/project/notes/new.md", b"new page"),
("db/memory.sqlite", db.as_slice()),
]),
);
restore_data_dir(data.path(), &from, true).unwrap();
assert_eq!(
std::fs::read(data.path().join("config.toml")).unwrap(),
b"# old config\n"
);
assert_eq!(
std::fs::read(data.path().join("wiki/default/project/notes/new.md")).unwrap(),
b"new page"
);
assert!(restore_scratch_dirs(data.path()).is_empty());
}
#[test]
fn restore_into_an_empty_data_dir_needs_no_force() {
let data = tempfile::TempDir::new().unwrap();
let scratch = tempfile::TempDir::new().unwrap();
let from = write_tarball(scratch.path(), "good.tar.gz", &good_archive());
restore_data_dir(data.path(), &from, false).unwrap();
assert_eq!(
std::fs::read(data.path().join("wiki/default/project/notes/new.md")).unwrap(),
b"new page"
);
assert!(data.path().join("db/memory.sqlite").is_file());
assert!(restore_scratch_dirs(data.path()).is_empty());
}
}
crates/ai-memory-cli/src/commands/resume.rs
+3
View File
@@ -161,8 +161,11 @@ pub async fn run(config: &Config, args: ResumeArgs) -> Result<i32> {
new_workstream: None,
executable: None,
yolo: args.yolo,
true_yolo: args.true_yolo,
fresh: args.fresh,
no_autowire: false,
env: Vec::new(),
env_file: None,
harness,
native_args: Vec::new(),
},
crates/ai-memory-cli/src/commands/run.rs
+2542 -93
View File
File diff suppressed because it is too large Load Diff
crates/ai-memory-cli/src/commands/run_autowire.rs
+1147 -84
View File
File diff suppressed because it is too large Load Diff
crates/ai-memory-cli/src/commands/serve.rs
+809 -39
View File
File diff suppressed because it is too large Load Diff
crates/ai-memory-cli/src/commands/server.rs
+268
View File
@@ -0,0 +1,268 @@
//! `ai-memory server add|list|remove` — manage local server profiles (#992).
//!
//! Local-only: writes `<data_dir>/servers.toml` and `<data_dir>/auth-tokens/`,
//! never contacts a server. See `server_profiles` for the routing rules.
use std::io::Write;
use std::path::Path;
use anyhow::{Context, Result, anyhow, bail};
use crate::cli::{ServerAddArgs, ServerArgs, ServerCommand, ServerListArgs, ServerRemoveArgs};
use crate::config::Config;
use crate::server_profiles::{self, ProfileName};
/// Run a `server` subcommand.
///
/// # Errors
/// Returns an error for an invalid name, URL, or root, an invalid existing
/// `servers.toml`, or a failed write.
pub fn run(config: &Config, args: ServerArgs) -> Result<()> {
let mut stdout = std::io::stdout();
match args.command {
ServerCommand::Add(args) => {
let stdin_token = if args.auth_token_stdin {
Some(read_token_line(&mut std::io::stdin().lock())?)
} else {
None
};
add(&config.data_dir, args, stdin_token, &mut stdout)
}
ServerCommand::List(args) => list(&config.data_dir, &args, &mut stdout),
ServerCommand::Remove(args) => remove(&config.data_dir, &args, &mut stdout),
}
}
fn parse_name(raw: &str) -> Result<ProfileName> {
ProfileName::parse(raw).ok_or_else(|| {
anyhow!(
"`{raw}` is not a valid profile name: use 1-64 lowercase letters, digits, `-` or `_`, \
starting with a letter or digit"
)
})
}
fn read_token_line(input: &mut impl std::io::BufRead) -> Result<String> {
let mut line = String::new();
input
.read_line(&mut line)
.context("reading the token from stdin")?;
let token = line.trim();
if token.is_empty() {
bail!("--auth-token-stdin was given but stdin had no token");
}
Ok(token.to_owned())
}
fn add(
data_dir: &Path,
args: ServerAddArgs,
stdin_token: Option<String>,
out: &mut impl Write,
) -> Result<()> {
let name = parse_name(&args.name)?;
let token = stdin_token.or(args.auth_token);
let outcome = server_profiles::add(data_dir, &name, &args.url, &args.roots, token.as_deref())?;
writeln!(out, "Registered server profile `{name}`.")?;
if outcome.roots_kept {
writeln!(
out,
"Kept the roots already registered for `{name}`; pass --root to replace them."
)?;
}
if outcome.token_discarded {
writeln!(
out,
"The URL changed, so the token stored for the previous URL was removed."
)?;
}
if server_profiles::read_token(data_dir, &name).is_none() {
writeln!(
out,
"No token is stored for `{name}`: repositories selecting it drop their capture \
until one is added with `ai-memory server add {name} --url … --auth-token-stdin`."
)?;
}
let registry = server_profiles::load(data_dir)?;
let unrooted: Vec<&str> = registry
.profiles
.iter()
.filter(|(_, profile)| profile.roots.is_empty())
.map(|(name, _)| name.as_str())
.collect();
if registry.profiles.len() > 1 && !unrooted.is_empty() {
writeln!(
out,
"Several profiles are registered, so a profile without --root is refused. \
Add roots to: {}",
unrooted.join(", ")
)?;
}
Ok(())
}
fn list(data_dir: &Path, args: &ServerListArgs, out: &mut impl Write) -> Result<()> {
let registry = server_profiles::load(data_dir)?;
if args.json {
let rows: Vec<_> = registry
.profiles
.iter()
.map(|(name, profile)| {
serde_json::json!({
"name": name.as_str(),
"url": profile.url,
"roots": profile.roots,
"token": token_state(data_dir, name),
})
})
.collect();
writeln!(out, "{}", serde_json::to_string_pretty(&rows)?)?;
return Ok(());
}
if registry.profiles.is_empty() {
writeln!(out, "No server profiles registered.")?;
return Ok(());
}
for (name, profile) in &registry.profiles {
let roots = if profile.roots.is_empty() {
"(any)".to_owned()
} else {
profile.roots.join(", ")
};
writeln!(
out,
"{name}\n url: {}\n roots: {roots}\n token: {}",
profile.url,
token_state(data_dir, name)
)?;
}
Ok(())
}
fn token_state(data_dir: &Path, name: &ProfileName) -> &'static str {
if server_profiles::read_token(data_dir, name).is_some() {
"stored"
} else {
"missing"
}
}
fn remove(data_dir: &Path, args: &ServerRemoveArgs, out: &mut impl Write) -> Result<()> {
let name = parse_name(&args.name)?;
if server_profiles::remove(data_dir, &name)? {
writeln!(out, "Removed server profile `{name}`.")?;
} else {
writeln!(out, "No server profile named `{name}`.")?;
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
fn add_args(name: &str, url: &str, roots: &[&str], token: Option<&str>) -> ServerAddArgs {
ServerAddArgs {
name: name.to_owned(),
url: url.to_owned(),
roots: roots.iter().map(|r| (*r).to_owned()).collect(),
auth_token: token.map(str::to_owned),
auth_token_stdin: false,
}
}
fn output(run: impl FnOnce(&mut Vec<u8>) -> Result<()>) -> String {
let mut out = Vec::new();
run(&mut out).unwrap();
String::from_utf8(out).unwrap()
}
/// `server add` that must succeed; returns what it printed.
fn add_ok(data_dir: &Path, args: ServerAddArgs, stdin_token: Option<String>) -> String {
output(|out| add(data_dir, args, stdin_token, out))
}
/// The listing reports whether a token is stored and never the token.
#[test]
fn list_never_prints_a_token() {
let dd = tempfile::tempdir().unwrap();
let root = dd.path().join("b").to_string_lossy().into_owned();
add_ok(
dd.path(),
add_args("team-b", "https://b.example", &[&root], Some("SECRET-B")),
None,
);
add_ok(
dd.path(),
add_args("team-c", "https://c.example", &[&root], None),
None,
);
for json in [false, true] {
let listed = output(|out| list(dd.path(), &ServerListArgs { json }, out));
assert!(!listed.contains("SECRET-B"), "{listed}");
assert!(listed.contains("https://b.example"), "{listed}");
assert!(listed.contains("stored"), "{listed}");
assert!(listed.contains("missing"), "{listed}");
}
}
#[test]
fn a_token_from_stdin_wins_and_is_stored() {
let dd = tempfile::tempdir().unwrap();
let token = read_token_line(&mut std::io::Cursor::new("from-stdin\n")).unwrap();
add_ok(
dd.path(),
add_args("a", "https://a.example", &[], None),
Some(token),
);
let name = ProfileName::parse("a").unwrap();
assert_eq!(
server_profiles::read_token(dd.path(), &name).as_deref(),
Some("from-stdin")
);
assert!(read_token_line(&mut std::io::Cursor::new("\n")).is_err());
}
#[test]
fn adding_a_second_unrooted_profile_warns_that_it_will_be_refused() {
let dd = tempfile::tempdir().unwrap();
add_ok(
dd.path(),
add_args("a", "https://a.example", &[], Some("t")),
None,
);
let text = add_ok(
dd.path(),
add_args("b", "https://b.example", &[], Some("t")),
None,
);
assert!(text.contains("Add roots to: a, b"), "{text}");
}
#[test]
fn invalid_names_are_refused_before_anything_is_written() {
let dd = tempfile::tempdir().unwrap();
let mut out = Vec::new();
let bad = add_args("../x", "https://a.example", &[], Some("t"));
assert!(add(dd.path(), bad, None, &mut out).is_err());
assert!(!dd.path().join("servers.toml").exists());
let remove_bad = ServerRemoveArgs {
name: "../x".into(),
};
assert!(remove(dd.path(), &remove_bad, &mut out).is_err());
}
#[test]
fn remove_reports_whether_the_profile_existed() {
let dd = tempfile::tempdir().unwrap();
add_ok(
dd.path(),
add_args("a", "https://a.example", &[], Some("t")),
None,
);
let args = ServerRemoveArgs { name: "a".into() };
assert!(output(|out| remove(dd.path(), &args, out)).contains("Removed"));
assert!(output(|out| remove(dd.path(), &args, out)).contains("No server profile"));
}
}
crates/ai-memory-cli/src/commands/setup_agent.rs
+49 -11
View File
@@ -9,14 +9,16 @@
//!
//! `setup-agent` bundles the extract + render into one command:
//!
//! docker run --rm \
//! -v "$HOME/.ai-memory:/host" \
//! akitaonrails/ai-memory:latest \
//! setup-agent \
//! --agent claude-code \
//! --to /host/hooks \
//! --host-prefix "$HOME/.ai-memory/hooks" \
//! --auth-token "$TOKEN"
//! ```text
//! docker run --rm \
//! -v "$HOME/.ai-memory:/host" \
//! akitaonrails/ai-memory:latest \
//! setup-agent \
//! --agent claude-code \
//! --to /host/hooks \
//! --host-prefix "$HOME/.ai-memory/hooks" \
//! --auth-token "$TOKEN"
//! ```
//!
//! 1. Copies `/usr/local/share/ai-memory/hooks/claude-code/*.{sh,ps1}` into
//! `/host/hooks/claude-code/` (which on the host is
@@ -82,6 +84,12 @@ pub fn run(config: &Config, args: SetupAgentArgs) -> Result<()> {
emit_zcode(&args)?;
return Ok(());
}
// Hermes also runs the native `hook` command directly (no shell), so there
// are no scripts to stage — setup-agent prints the YAML block instead.
if matches!(args.agent, AgentChoice::Hermes) {
emit_hermes(&args)?;
return Ok(());
}
let Some(agent_sub) = args.agent.script_hook_subdir() else {
bail!("internal: generated integration should have returned before staging hooks")
};
@@ -174,7 +182,8 @@ pub fn run(config: &Config, args: SetupAgentArgs) -> Result<()> {
| AgentChoice::Omp
| AgentChoice::Openclaw
| AgentChoice::Zero
| AgentChoice::Zcode => {
| AgentChoice::Zcode
| AgentChoice::Hermes => {
bail!(
"internal: generated integration should have returned before emitting staged hooks"
)
@@ -243,6 +252,35 @@ fn emit_zcode(args: &SetupAgentArgs) -> Result<()> {
Ok(())
}
/// Print the `hooks:` block for `~/.hermes/config.yaml` (the follow-up to
/// #623, where Hermes was accepted for capture/storage but had no installer).
/// No scripts are staged: Hermes splits the configured `command` itself and
/// runs the ai-memory binary directly with the JSON payload on stdin. The
/// binary must be reachable on the host that runs Hermes — for docker-wrapper
/// setups install the native binary or point `command` at the wrapper.
fn emit_hermes(args: &SetupAgentArgs) -> Result<()> {
let block = crate::commands::render_shared::build_hermes_hooks_yaml(
&args.server_url,
args.auth_token.as_deref(),
None,
None,
);
println!("# Hermes Agent — merge the `hooks:` block into ~/.hermes/config.yaml");
println!("# The `command` must be an ai-memory binary reachable on the host that");
println!("# runs Hermes; prefer `ai-memory install-hooks --agent hermes` from");
println!("# that host so the path is resolved for you.");
if args.auth_token.is_some() {
println!("# Treat the config as sensitive (chmod 600).");
}
println!("# NOTE: Hermes splits each `command` with shlex.split and runs it with");
println!("# no shell, so the generated line is argv, not a shell command.");
println!("# NOTE: tool observations only — session lifecycle stays with the");
println!("# ai-memory memory-provider plugin.");
println!();
println!("{block}");
Ok(())
}
fn normalise_hook_server_url(url: &str) -> String {
url.trim().trim_end_matches('/').to_string()
}
@@ -545,12 +583,12 @@ fn source_candidates(explicit: Option<&Path>, sub: &str, exe: Option<PathBuf>) -
if let Some(exe) = exe {
// Release tarball: `hooks/` sits in the same dir as the binary.
if let Some(dir) = exe.parent() {
v.push(dir.join("hooks").join(sub));
v.push(dir.join(crate::install_layout::HOOKS_DIR_NAME).join(sub));
}
// Repo-local fallback for `cargo run setup-agent` during dev:
// target/<profile>/<bin> → repo root.
if let Some(root) = exe.parent().and_then(Path::parent).and_then(Path::parent) {
v.push(root.join("hooks").join(sub));
v.push(root.join(crate::install_layout::HOOKS_DIR_NAME).join(sub));
}
}
v
crates/ai-memory-cli/src/commands/show.rs
+4 -1
View File
@@ -127,7 +127,7 @@ struct JsonOutput {
/// Pick a local checkout and harness, then delegate to managed `run`.
pub async fn run(config: &Config, args: ShowArgs) -> Result<i32> {
if args.json && (args.yolo || args.fresh || !args.native_args.is_empty()) {
if args.json && (args.yolo || args.true_yolo || args.fresh || !args.native_args.is_empty()) {
bail!("--json only lists launch options; do not combine it with launch arguments");
}
let root = std::env::current_dir()
@@ -225,8 +225,11 @@ pub async fn run(config: &Config, args: ShowArgs) -> Result<i32> {
new_workstream: None,
executable: None,
yolo: args.yolo,
true_yolo: args.true_yolo,
fresh: args.fresh,
no_autowire: false,
env: Vec::new(),
env_file: None,
harness: Some(harness),
native_args: args.native_args,
},
crates/ai-memory-cli/src/commands/uninstall.rs
+96 -12
View File
@@ -11,7 +11,9 @@ use crate::cli::UninstallArgs;
use crate::commands::apply_shared::apply_atomic;
use crate::commands::apply_shared::mutate_json;
use crate::commands::apply_shared::mutate_toml;
use crate::commands::path_util::{claude_config_dir, claude_config_paths, home_dir};
use crate::commands::path_util::{
agent_config_home, claude_config_dir, claude_config_paths, home_dir,
};
use crate::commands::{data_purge, install_hooks, install_mcp, openclaw_plugin};
use crate::config::Config;
use ai_memory_core::routing_skills::{
@@ -72,6 +74,7 @@ enum DeleteKind {
OpenClawEntrypoint,
KiroCliV3Hooks,
ManagedSkill,
AutowireSentinel,
}
impl DeleteKind {
@@ -86,6 +89,7 @@ impl DeleteKind {
Self::OpenClawEntrypoint => "OpenClaw plugin entrypoint",
Self::KiroCliV3Hooks => "Kiro CLI v3 hook file",
Self::ManagedSkill => "managed Agent Skill",
Self::AutowireSentinel => "auto-wire sentinel",
}
}
}
@@ -130,6 +134,38 @@ fn push_rewrite(plan: &mut Vec<PlannedChange>, path: PathBuf, removed: Vec<Strin
});
}
/// The OMP agent dirs an ai-memory install may have written to, active first:
/// the one OMP loads now (`--profile`, `OMP_PROFILE`, `PI_PROFILE`, else
/// `PI_CODING_AGENT_DIR`), the default profile's, where earlier releases put
/// the extension whenever `PI_CODING_AGENT_DIR` was set, the active one as
/// earlier releases resolved it (they ignored `PI_CONFIG_DIR` and wrote under
/// `~/.omp`), and `~/.omp/agent`, where `install-mcp` always wrote the MCP
/// entry. A profile name OMP refuses is reported and skipped instead of
/// aborting every other agent's cleanup.
fn omp_agent_dirs(home: Option<&Path>, profile: Option<&str>) -> Vec<PathBuf> {
let Some(home) = home else {
return Vec::new();
};
let env = |name: &str| std::env::var_os(name);
let without_config_dir = |name: &str| (name != "PI_CONFIG_DIR").then(|| env(name)).flatten();
let mut dirs = Vec::with_capacity(4);
match ai_memory_workstream::omp_agent_dir(home, profile, env) {
Ok(dir) => dirs.push(dir),
Err(error) => eprintln!("warning: skipping the active OMP profile: {error:#}"),
}
let fallbacks = [
ai_memory_workstream::omp_agent_dir(home, Some("default"), env).ok(),
ai_memory_workstream::omp_agent_dir(home, profile, without_config_dir).ok(),
Some(home.join(".omp").join("agent")),
];
for dir in fallbacks.into_iter().flatten() {
if !dirs.contains(&dir) {
dirs.push(dir);
}
}
dirs
}
fn push_generated_delete(plan: &mut Vec<PlannedChange>, path: PathBuf, kind: DeleteKind) {
if generated_file_is_ours(&path, kind) {
plan.push(PlannedChange::DeleteFile { path, kind });
@@ -139,13 +175,19 @@ fn push_generated_delete(plan: &mut Vec<PlannedChange>, path: PathBuf, kind: Del
/// Build the full removal plan by reading each existing config file and
/// running the matching pure stripper. Missing files / no-matches
/// produce no entry. `name`/`url` identify the MCP server.
fn build_plan(args: &UninstallArgs) -> anyhow::Result<Vec<PlannedChange>> {
fn build_plan(args: &UninstallArgs, data_dir: &Path) -> anyhow::Result<Vec<PlannedChange>> {
let mut plan = Vec::new();
let want = |k: crate::cli::UninstallOnly| args.only.is_none() || args.only == Some(k);
let name = args.mcp_name.as_deref();
let url = args.mcp_url.as_str();
let home = home_dir();
let claude_config_dir = claude_config_dir(std::env::var_os("CLAUDE_CONFIG_DIR"));
let omp_dirs = if want(crate::cli::UninstallOnly::Hooks) || want(crate::cli::UninstallOnly::Mcp)
{
omp_agent_dirs(home.as_deref(), args.profile.as_deref())
} else {
Vec::new()
};
// ---- Hooks (JSON configs) ----
if want(crate::cli::UninstallOnly::Hooks) {
@@ -299,12 +341,10 @@ fn build_plan(args: &UninstallArgs) -> anyhow::Result<Vec<PlannedChange>> {
let plugin2 = install_hooks::opencode2_plugin_path()?;
push_generated_delete(&mut plan, plugin2, DeleteKind::OpenCode2Plugin);
let omp_profile = args.profile.as_deref();
let omp = install_hooks::omp_extension_path(omp_profile)?;
push_generated_delete(&mut plan, omp.clone(), DeleteKind::OmpExtension);
let legacy_omp = omp.with_file_name("ai-memory.ts");
if legacy_omp != omp {
for dir in &omp_dirs {
let omp = dir.join("extensions").join("ai-memory-omp.ts");
let legacy_omp = omp.with_file_name("ai-memory.ts");
push_generated_delete(&mut plan, omp, DeleteKind::OmpExtension);
push_generated_delete(&mut plan, legacy_omp, DeleteKind::OmpExtension);
}
@@ -367,6 +407,17 @@ fn build_plan(args: &UninstallArgs) -> anyhow::Result<Vec<PlannedChange>> {
Path::new(".claude.json"),
Path::new(".claude.json"),
)
} else if matches!(client, Codex) {
// Older installs wrote the Codex MCP entry to ~/.codex/config.toml
// even with CODEX_HOME set, so sweep that file too.
claude_config_paths(
home.as_deref(),
agent_config_home(std::env::var_os("CODEX_HOME")).as_deref(),
Path::new(".codex/config.toml"),
Path::new("config.toml"),
)
} else if matches!(client, Omp) {
omp_dirs.iter().map(|dir| dir.join("mcp.json")).collect()
} else {
let Ok(path) = install_mcp::mcp_config_path(client) else {
continue;
@@ -438,6 +489,24 @@ fn build_plan(args: &UninstallArgs) -> anyhow::Result<Vec<PlannedChange>> {
}
}
// ---- Auto-wire sentinels (ai-memory's own state) ----
// `ai-memory run` skips wiring while a sentinel exists, so one left behind
// would stop the next managed launch from reinstalling what this removes.
// All of them go: one keyed on a config home this environment cannot see
// costs only an idempotent re-apply.
if want(crate::cli::UninstallOnly::Hooks) || want(crate::cli::UninstallOnly::Mcp) {
let dir = crate::commands::run_autowire::autowire_state_dir(data_dir);
let mut sentinels: Vec<PathBuf> = std::fs::read_dir(&dir)
.into_iter()
.flatten()
.filter_map(|entry| entry.ok().map(|entry| entry.path()))
.collect();
sentinels.sort();
for path in sentinels {
push_generated_delete(&mut plan, path, DeleteKind::AutowireSentinel);
}
}
Ok(plan)
}
@@ -599,7 +668,7 @@ pub fn run(config: &Config, args: UninstallArgs) -> anyhow::Result<()> {
let name = args.mcp_name.clone();
let url = args.mcp_url.clone();
let plan = build_plan(&args)?;
let plan = build_plan(&args, &config.data_dir)?;
print_plan(&plan);
if args.purge_data {
for path in data_purge::purge_preview(&config.data_dir) {
@@ -642,14 +711,26 @@ pub fn run(config: &Config, args: UninstallArgs) -> anyhow::Result<()> {
}
// Removing the hooks removes the only readers of the stored bearer, so
// leaving it on disk would strand a live credential (#552). Best-effort:
// an unremovable file must not fail a teardown that otherwise succeeded.
if let Err(error) = crate::config::clear_hook_auth_token(&config.data_dir) {
// leaving it on disk would strand a live credential (#552). An uninstall
// that keeps the hooks (`--only mcp|instructions|skills`) keeps it: the
// native hook, the shell hooks and the generated TypeScript integrations
// (the Pi one also bridges MCP) all still read it. Best-effort: an
// unremovable file must not fail a teardown that otherwise succeeded.
let hooks_removed = args.only.is_none() || args.only == Some(crate::cli::UninstallOnly::Hooks);
if hooks_removed && let Err(error) = crate::config::clear_hook_auth_token(&config.data_dir) {
eprintln!(
"ai-memory uninstall warning: could not remove the stored auth token under {}: {error}",
config.data_dir.display()
);
}
// Server-profile tokens (#992) are read by the same hooks, so the same
// reasoning applies; the registry of URLs stays.
if let Err(error) = crate::server_profiles::clear_tokens(&config.data_dir) {
eprintln!(
"ai-memory uninstall warning: could not remove the stored server-profile tokens under {}: {error}",
config.data_dir.display()
);
}
if args.purge_data {
for path in data_purge::purge_data_dirs(&config.data_dir)? {
@@ -1051,6 +1132,9 @@ fn generated_file_is_ours(path: &Path, kind: DeleteKind) -> bool {
content.contains("Auto-generated by `ai-memory install-hooks --agent omp --apply`")
&& content.contains("const AGENT = \"omp\";")
}
// Auto-wire writes its sentinels empty; anything else in that
// directory is not one of them.
DeleteKind::AutowireSentinel => content.is_empty(),
DeleteKind::PiExtension => {
content.contains("Auto-generated by `ai-memory install-hooks --agent pi --apply`")
&& content.contains("const AGENT = \"pi\";")
crates/ai-memory-cli/src/commands/upgrade.rs
+1436
View File
File diff suppressed because it is too large Load Diff
crates/ai-memory-cli/src/commands/user.rs
+4 -1
View File
@@ -95,6 +95,9 @@ pub async fn run(config: &Config, args: UserArgs) -> Result<()> {
UserCommand::Disable(args) => disable(&ep, args).await,
UserCommand::Enable(args) => enable(&ep, args).await,
UserCommand::Patch(args) => patch(&ep, args).await,
UserCommand::Grant(args) => crate::commands::grant::grant(&ep, &args).await,
UserCommand::Revoke(args) => crate::commands::grant::revoke(&ep, &args).await,
UserCommand::Grants(args) => crate::commands::grant::list_for_user(&ep, &args).await,
}
}
@@ -409,7 +412,7 @@ fn confirm(prompt: &str) -> Result<()> {
Ok(())
}
fn url_encode(s: &str) -> String {
pub(crate) fn url_encode(s: &str) -> String {
let mut out = String::with_capacity(s.len());
for &b in s.as_bytes() {
if b.is_ascii_alphanumeric() || matches!(b, b'-' | b'_' | b'.') {
crates/ai-memory-cli/src/config.rs
+1516 -36
View File
File diff suppressed because it is too large Load Diff
crates/ai-memory-cli/src/http_client.rs
+8
View File
@@ -138,6 +138,14 @@ impl ServerEndpoint {
)
}
/// Build for a target the spawning hook already resolved: its URL
/// (including any mount path, as `install-hooks` rendered it) and the
/// bearer it authenticates with. No config or environment is consulted.
#[must_use]
pub(crate) fn for_hook_target(url: String, token: Option<String>) -> Self {
Self::build(Some(url), token, true, None)
}
/// Build from an explicit URL + token pair (useful for tests that
/// cannot safely mutate the process environment).
///
crates/ai-memory-cli/src/install_layout.rs
+29
View File
@@ -0,0 +1,29 @@
//! Release and install path basenames shared by packaging-aware CLI commands.
//!
//! These names match the shipped binary (`[[bin]] name`), release archive
//! entries, and the binary-sibling hooks bundle. They are not domain
//! constants — keep them in the CLI crate, not `ai-memory-core`.
/// Shipped binary basename (`[[bin]] name`, `/proc/*/comm`).
///
/// Release archives use [`shipped_binary_name`] instead — Windows zips ship
/// `ai-memory.exe`.
pub const BINARY_NAME: &str = "ai-memory";
/// Sibling hooks bundle dir in release archives and install prefixes.
pub const HOOKS_DIR_NAME: &str = "hooks";
/// On-disk / archive basename for the release binary on this host.
///
/// Matches `release.yml`: Unix tarballs ship `ai-memory`; the Windows zip
/// ships `ai-memory.exe`.
pub fn shipped_binary_name() -> &'static str {
#[cfg(windows)]
{
"ai-memory.exe"
}
#[cfg(not(windows))]
{
BINARY_NAME
}
}
crates/ai-memory-cli/src/lib.rs
+11
View File
@@ -24,9 +24,11 @@ mod cli;
mod commands;
mod config;
mod http_client;
mod install_layout;
mod logging;
mod marker;
mod process_guard;
mod server_profiles;
use cli::{Cli, Command};
use config::Config;
@@ -86,6 +88,9 @@ pub async fn run() -> Result<()> {
Command::Status(args) => commands::status::run(&config, args).await,
Command::Doctor(args) => commands::doctor::run(&config, args).await,
Command::Backfill(args) => commands::backfill::run(&config, args).await,
Command::RepairBackfillTimestamps(args) => {
commands::repair_backfill_timestamps::run(&config, args).await
}
Command::Run(args) => {
let exit_code = commands::run::run(&config, args).await?;
if exit_code != 0 {
@@ -129,6 +134,9 @@ pub async fn run() -> Result<()> {
Command::Serve(args) => commands::serve::run(&config, args).await,
Command::Reset(args) => commands::reset::run(&config, args),
Command::Compact(args) => commands::compact::run(&config, args).await,
Command::ReclaimLedgerVersions(args) => {
commands::reclaim_ledger_versions::run(&config, args).await
}
Command::Backup(args) => commands::backup::run(&config, args).await,
Command::ExportOkf(args) => commands::export_okf::run(&config, args).await,
Command::Restore(args) => commands::restore::run(&config, args),
@@ -164,9 +172,12 @@ pub async fn run() -> Result<()> {
Command::MoveProject(args) => commands::move_project::run(&config, args).await,
Command::MoveSession(args) => commands::move_session::run(&config, args).await,
Command::Uninstall(args) => commands::uninstall::run(&config, args),
Command::Upgrade(args) => commands::upgrade::run(&config, args).await,
Command::Auth(args) => commands::auth::run(&config, args).await,
Command::User(args) => commands::user::run(&config, args).await,
Command::ApiKey(args) => commands::api_key::run(&config, args).await,
Command::Project(args) => commands::project::run(&config, args).await,
Command::Server(args) => commands::server::run(&config, args),
// `Completions` is handled in the fast-path above (before config/tracing).
Command::Completions(args) => commands::completions::run(args),
}
crates/ai-memory-cli/src/logging.rs
+104 -3
View File
@@ -21,6 +21,7 @@ use anyhow::Result;
use tracing_appender::non_blocking::WorkerGuard;
use tracing_appender::rolling::{RollingFileAppender, Rotation};
use tracing_subscriber::EnvFilter;
use tracing_subscriber::filter::LevelFilter;
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::util::SubscriberInitExt;
@@ -87,6 +88,37 @@ fn resolve_file_appender(
(None, notices)
}
/// The `EnvFilter` directive used when `RUST_LOG` is unset.
///
/// Two overrides bracket the operator's `log_level`, and order is
/// load-bearing because a later directive wins in an `EnvFilter`:
///
/// - `rmcp=warn` is **prepended**, so it is the weakest directive and an
/// operator can restore the external MCP SDK's per-request info logs
/// through `log_level` (e.g. `info,rmcp=info`) without setting `RUST_LOG`.
/// Left at info, `rmcp` alone is ~half the default server log (#894).
/// A target directive also beats a *quieter* global level, so the cap is
/// left out when `log_level` is already `warn` or quieter: it may only
/// lower rmcp, never re-enable warnings an `error`/`off` level silenced.
/// - `tracing_appender=warn` stays **appended**, so it is the strongest and
/// cannot be lowered through `log_level`. That guard is invariant #15: the
/// appender must never log at its own level or it feeds itself (the loop
/// that filled 137 GB for agentmemory #519).
///
/// `RUST_LOG` (`EnvFilter::try_from_default_env`) still overrides all of this.
fn default_filter(log_level: &str) -> String {
// The last bare level in the list is the global one EnvFilter applies.
let global = log_level
.split(',')
.filter_map(|directive| directive.trim().parse::<LevelFilter>().ok())
.next_back();
if global.is_some_and(|level| level <= LevelFilter::WARN) {
format!("{log_level},tracing_appender=warn")
} else {
format!("rmcp=warn,{log_level},tracing_appender=warn")
}
}
/// Initialise the global tracing subscriber.
///
/// Returns a guard whose drop flushes any pending log lines; `None` when no
@@ -106,9 +138,8 @@ pub fn init(config: &Config, warnings: DegradeWarnings) -> Result<Option<WorkerG
}
}
let default_filter = format!("{},tracing_appender=warn", config.log_level);
let env_filter =
EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new(default_filter));
let env_filter = EnvFilter::try_from_default_env()
.unwrap_or_else(|_| EnvFilter::new(default_filter(&config.log_level)));
let stderr_layer = tracing_subscriber::fmt::layer()
.with_target(true)
@@ -139,6 +170,76 @@ pub fn init(config: &Config, warnings: DegradeWarnings) -> Result<Option<WorkerG
mod tests {
use super::*;
/// The effective per-target level `EnvFilter` resolves the default filter
/// to. `EnvFilter`'s `Display` reprints its live directives (last wins per
/// target), so it reflects the real conflict resolution, not the raw
/// string. Returns `None` for a target the filter carries no directive for.
fn effective_level(log_level: &str, target: &str) -> Option<String> {
let printed = EnvFilter::new(default_filter(log_level)).to_string();
printed
.split(',')
.find_map(|d| d.strip_prefix(&format!("{target}=")).map(str::to_owned))
}
#[test]
fn default_filter_suppresses_rmcp_and_the_appender() {
// (a) With a plain `info` level, both the noisy external MCP SDK and
// the appender are pinned to warn.
assert_eq!(effective_level("info", "rmcp").as_deref(), Some("warn"));
assert_eq!(
effective_level("info", "tracing_appender").as_deref(),
Some("warn")
);
}
#[test]
fn operator_can_restore_rmcp_through_log_level() {
// (b) `rmcp=warn` is prepended (weakest), so a log_level directive for
// the same target wins and brings the SDK's info logs back — while the
// appender stays warn.
assert_eq!(
effective_level("info,rmcp=info", "rmcp").as_deref(),
Some("info"),
"an operator must be able to restore rmcp via log_level"
);
assert_eq!(
effective_level("info,rmcp=info", "tracing_appender").as_deref(),
Some("warn"),
"restoring rmcp must not disturb the appender guard"
);
}
#[test]
fn a_quieter_log_level_is_not_overridden_for_rmcp() {
// (d) A target directive beats the global level whichever is louder,
// so a prepended `rmcp=warn` under `log_level = "error"` or `"off"`
// would re-enable the SDK's warnings the operator had silenced. The
// cap may only lower rmcp: here it must carry no directive at all.
for quieter in ["warn", "error", "off", "debug,error"] {
assert_eq!(
effective_level(quieter, "rmcp"),
None,
"log_level {quieter:?} must govern rmcp itself"
);
}
assert_eq!(
effective_level("off,rmcp=info", "rmcp").as_deref(),
Some("info"),
"an explicit rmcp directive still wins"
);
}
#[test]
fn log_level_cannot_lower_the_appender_below_warn() {
// (c) `tracing_appender=warn` is appended (strongest), so no log_level
// directive can lower it — the invariant #15 feedback-loop guard.
assert_eq!(
effective_level("info,tracing_appender=trace", "tracing_appender").as_deref(),
Some("warn"),
"the appender guard must be non-overridable through log_level"
);
}
// Issue #158: the log directory EXISTS but the filesystem is read-only —
// dir creation "succeeds", file creation fails. The old code panicked
// here (RollingFileAppender::new); the chain must fall through to the
crates/ai-memory-cli/src/marker.rs
+233 -18
View File
@@ -113,7 +113,9 @@ pub(crate) fn find_marker(cwd: &str) -> Option<PathBuf> {
}
fn find_marker_with_home(cwd: &str, home: Option<&Path>) -> Option<PathBuf> {
find_marker_matching(cwd, home, |_| true)
find_marker_matching(cwd, home, OutsideHome::StopAtCheckoutRoot, |path| {
Some(path.to_path_buf())
})
}
/// Like [`find_marker`], but skips a marker that declares nothing beyond a
@@ -129,34 +131,49 @@ pub(crate) fn find_settings_marker(cwd: &str) -> Option<PathBuf> {
}
fn find_settings_marker_with_home(cwd: &str, home: Option<&Path>) -> Option<PathBuf> {
find_marker_matching(cwd, home, |path| {
std::fs::read_to_string(path).is_ok_and(|text| declares_more_than_capture(&text))
find_marker_matching(cwd, home, OutsideHome::StopAtCheckoutRoot, |path| {
std::fs::read_to_string(path)
.is_ok_and(|text| declares_more_than_capture(&text))
.then(|| path.to_path_buf())
})
}
/// Shared walk-up-from-`cwd`-toward-`$HOME` used by [`find_marker_with_home`]
/// and [`find_settings_marker_with_home`]; `matches` decides whether a marker
/// file found along the way stops the walk (returned) or is skipped in favor
/// of the next ancestor. The HOME/checkout-root boundary is identical either
/// way — only which markers count as a stopping point differs.
fn find_marker_matching(
/// Where a walk that starts outside `$HOME` stops.
#[derive(Clone, Copy)]
enum OutsideHome {
/// At the nearest `.git` root, or `cwd` itself outside any checkout.
StopAtCheckoutRoot,
/// At the filesystem root.
WalkToRoot,
}
/// Shared walk-up-from-`cwd`-toward-`$HOME` used by every marker lookup;
/// `matches` maps a marker file found along the way to a result that stops
/// the walk, or `None` to continue to the next ancestor. Inside `$HOME` the
/// walk stops at `$HOME`; `outside_home` picks the stop for a start outside it.
fn find_marker_matching<T>(
cwd: &str,
home: Option<&Path>,
matches: impl Fn(&Path) -> bool,
) -> Option<PathBuf> {
outside_home: OutsideHome,
mut matches: impl FnMut(&Path) -> Option<T>,
) -> Option<T> {
let start = absolute_normalized(Path::new(cwd));
let home = home.map(absolute_normalized);
let boundary = match home.as_deref() {
Some(home) if start.starts_with(home) => Some(home.to_path_buf()),
Some(_) => Some(checkout_root(&start).unwrap_or_else(|| start.clone())),
None => None,
let boundary = match (home.as_deref(), outside_home) {
(Some(home), _) if start.starts_with(home) => Some(home.to_path_buf()),
(Some(_), OutsideHome::StopAtCheckoutRoot) => {
Some(checkout_root(&start).unwrap_or_else(|| start.clone()))
}
(Some(_), OutsideHome::WalkToRoot) | (None, _) => None,
};
let mut dir = start.as_path();
loop {
let candidate = dir.join(".ai-memory.toml");
if candidate.is_file() && matches(&candidate) {
return Some(candidate);
if candidate.is_file()
&& let Some(found) = matches(&candidate)
{
return Some(found);
}
if boundary.as_deref() == Some(dir) {
return None;
@@ -181,11 +198,12 @@ fn find_marker_matching(
/// the file. That is conservative on purpose: it can only turn a marker INTO
/// a boundary, never wrongly make one transparent.
fn declares_more_than_capture(text: &str) -> bool {
const QUOTED_KEYS: [&str; 4] = [
const QUOTED_KEYS: [&str; 5] = [
"workspace",
"project",
"project_strategy",
"drop_subagent_captures",
"identity",
];
const FLAG_KEYS: [&str; 3] = ["default_global", "inject_on_session_start", "max_chars"];
QUOTED_KEYS
@@ -194,6 +212,78 @@ fn declares_more_than_capture(text: &str) -> bool {
|| FLAG_KEYS
.iter()
.any(|key| parse_flag_in(text, key).is_some())
|| server_selection_in(text).is_some()
}
/// A marker's `server = "<profile>"` selection (#992), and the directory of
/// the marker that made it.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct ServerSelection {
/// The raw value, validated later by `server_profiles::ProfileName`;
/// `None` when a marker on the walk exists but could not be read, which
/// the caller must treat as a refused selection, never as "no selection".
pub(crate) name: Option<String>,
/// Directory holding the declaring marker, lexically normalized.
pub(crate) marker_dir: PathBuf,
}
/// The nearest marker on the walk from `cwd` that declares `server`.
///
/// Deliberately *not* [`find_settings_marker`]'s nearest-marker rule, in two
/// ways, both of which would otherwise route a profile's capture to the
/// install-default server:
///
/// - Routing is inherited down the tree: a nested marker that sets only
/// `workspace` must not reset a subdirectory of a profile-routed tree. A
/// nested marker can only select a different profile, which that
/// profile's `roots` then have to admit.
/// - Outside `$HOME` the walk does not stop at the checkout root, so an
/// organisation-level marker above a repository (`/srv/work/team-b/`,
/// `/Volumes/…`) still routes it. A marker planted higher up can only name
/// a profile this operator registered, which `roots` gate once there are
/// several.
///
/// It also fails closed on content: a marker it cannot read is a refused
/// selection, and a UTF-8 BOM or stray non-UTF-8 byte cannot hide the key.
///
/// `home` is the walk boundary inside `$HOME`; the caller passes the one it
/// also expands `~/` roots against.
pub(crate) fn find_server_selection(cwd: &str, home: Option<&Path>) -> Option<ServerSelection> {
find_marker_matching(cwd, home, OutsideHome::WalkToRoot, |path| {
let name = match std::fs::read(path) {
Ok(bytes) => Some(server_selection_in(&String::from_utf8_lossy(&bytes))?),
Err(_) => None,
};
Some(ServerSelection {
name,
marker_dir: path.parent().map(Path::to_path_buf).unwrap_or_default(),
})
})
}
/// Line-based like [`parse_key_in`], but it fails closed on shape: a
/// `server = team-b` without quotes, or an empty `server = ""`, still counts
/// as a selection (and is then rejected by name validation) instead of being
/// ignored and silently delivered to the install default. Section headers are
/// not tracked, so a `server` key under any table is treated the same way.
fn server_selection_in(text: &str) -> Option<String> {
for line in text.lines() {
// A BOM is not whitespace to `trim_start`, and would hide a first-line key.
let line = line.trim_start_matches('\u{feff}').trim_start();
let Some(rest) = line.strip_prefix("server") else {
continue;
};
let Some(value) = rest.trim_start().strip_prefix('=') else {
continue;
};
let value = value.trim();
let value = match value.strip_prefix('"') {
Some(quoted) => quoted.split_once('"').map_or(quoted, |(inner, _)| inner),
None => value.split('#').next().unwrap_or("").trim(),
};
return Some(value.to_owned());
}
None
}
/// Make `path` absolute and resolve its `.`/`..` components, WITHOUT
@@ -662,4 +752,129 @@ project = "infra" # this is fine
);
assert_eq!(normalized, link.join("file.txt"));
}
// ── #992: `server` profile selection ─────────────────────────────────
#[test]
fn server_selection_parses_quoted_bare_and_empty_values() {
assert_eq!(
server_selection_in("server = \"team-b\" # comment\n").as_deref(),
Some("team-b")
);
assert_eq!(
server_selection_in(" server=team-b # comment\n").as_deref(),
Some("team-b")
);
assert_eq!(server_selection_in("server = \"\"\n").as_deref(), Some(""));
assert_eq!(
server_selection_in("servers = \"x\"\nserver_url = \"y\"\n"),
None
);
assert_eq!(server_selection_in("workspace = \"a\"\n"), None);
}
/// A marker declaring only `server` is a settings boundary like any other
/// root-level key; a `[capture]`-only marker stays transparent.
#[test]
fn a_server_only_marker_is_a_settings_boundary() {
assert!(declares_more_than_capture("server = \"team-b\"\n"));
assert!(declares_more_than_capture("server = team-b\n"));
assert!(!declares_more_than_capture(
"[capture]\nignore_paths = [\"x/**\"]\n"
));
}
/// Routing is inherited: a nested marker that sets only `workspace`, or
/// only `[capture]`, keeps the ancestor's profile. Without this, a
/// sub-project marker would silently send a profile-routed tree to the
/// install-default server.
#[test]
fn nested_markers_without_server_inherit_the_ancestor_selection() {
let tmp = TempDir::new().unwrap();
write_marker(tmp.path(), "workspace = \"team-b\"\nserver = \"team-b\"\n");
let scoped = tmp.path().join("scoped");
let capture_only = scoped.join("capture-only");
fs::create_dir_all(&capture_only).unwrap();
write_marker(&scoped, "workspace = \"other\"\n");
write_marker(&capture_only, "[capture]\nignore_paths = [\"x/**\"]\n");
let selection = find_server_selection(capture_only.to_str().unwrap(), Some(tmp.path()))
.expect("the ancestor's selection applies");
assert_eq!(selection.name.as_deref(), Some("team-b"));
assert_eq!(selection.marker_dir, absolute_normalized(tmp.path()));
}
#[test]
fn the_nearest_server_declaration_wins() {
let tmp = TempDir::new().unwrap();
write_marker(tmp.path(), "server = \"team-a\"\n");
let inner = tmp.path().join("inner");
fs::create_dir_all(&inner).unwrap();
write_marker(&inner, "server = \"team-b\"\n");
let selection = find_server_selection(inner.to_str().unwrap(), Some(tmp.path())).unwrap();
assert_eq!(selection.name.as_deref(), Some("team-b"));
assert_eq!(selection.marker_dir, absolute_normalized(&inner));
}
#[test]
fn no_server_key_anywhere_is_no_selection() {
let tmp = TempDir::new().unwrap();
write_marker(tmp.path(), "workspace = \"a\"\n");
assert_eq!(
find_server_selection(tmp.path().to_str().unwrap(), Some(tmp.path())),
None
);
}
/// Outside `$HOME`, an organisation-level marker above the repository's
/// own `.git` still routes it; stopping at the checkout root would send
/// the repository to the install default.
#[test]
fn outside_home_the_walk_reaches_a_marker_above_the_checkout_root() {
let tmp = TempDir::new().unwrap();
let org = tmp.path().join("org");
let repo = org.join("repo");
fs::create_dir_all(repo.join(".git")).unwrap();
write_marker(&org, "server = \"team-b\"\n");
write_marker(&repo, "workspace = \"api\"\n");
let elsewhere = tmp.path().join("home");
let selection = find_server_selection(repo.to_str().unwrap(), Some(&elsewhere)).unwrap();
assert_eq!(selection.name.as_deref(), Some("team-b"));
assert_eq!(selection.marker_dir, absolute_normalized(&org));
}
/// A BOM or a non-UTF-8 byte must not hide the key.
#[test]
fn encoding_noise_cannot_hide_a_server_key() {
let tmp = TempDir::new().unwrap();
for bytes in [
b"\xEF\xBB\xBFserver = \"team-b\"\n".as_slice(),
b"# caf\xE9\nserver = \"team-b\"\n".as_slice(),
] {
fs::write(tmp.path().join(".ai-memory.toml"), bytes).unwrap();
let selection =
find_server_selection(tmp.path().to_str().unwrap(), Some(tmp.path())).unwrap();
assert_eq!(selection.name.as_deref(), Some("team-b"), "{bytes:?}");
}
}
/// A marker that exists but cannot be read is a refused selection, not
/// "no selection" — it may well declare a profile.
#[cfg(unix)]
#[test]
fn an_unreadable_marker_is_a_refused_selection() {
use std::os::unix::fs::PermissionsExt as _;
let tmp = TempDir::new().unwrap();
let marker = write_marker(tmp.path(), "server = \"team-b\"\n");
fs::set_permissions(&marker, fs::Permissions::from_mode(0o000)).unwrap();
if fs::read(&marker).is_ok() {
// Running as root: permissions cannot make the file unreadable.
return;
}
let selection =
find_server_selection(tmp.path().to_str().unwrap(), Some(tmp.path())).unwrap();
assert_eq!(selection.name, None);
}
}
crates/ai-memory-cli/src/process_guard.rs
+19 -1
View File
@@ -11,12 +11,30 @@ use std::ffi::OsStr;
use sysinfo::System;
/// Binary name to match against `/proc/*/comm` (or platform equivalent).
pub const BIN_NAME: &str = "ai-memory";
pub const BIN_NAME: &str = crate::install_layout::BINARY_NAME;
/// Return PIDs of *other* `ai-memory` processes (excluding the current
/// process and any threads of it).
#[must_use]
pub fn sibling_processes() -> Vec<sysinfo::Pid> {
// Test injection: a comma-separated list of fake PIDs to report as alive
// siblings, bypassing both the real scan AND the `cfg!(test)` opt-out
// below. This is what lets the guard's REFUSAL path be exercised by an
// in-process test (reset / reindex / restore / uninstall --purge-data all
// call `sibling_processes()` directly, and `cfg!(test)` alone would
// otherwise force every in-process test onto the "no siblings" branch).
// Checked first, and not itself gated by `cfg!(test)`, matching the
// existing `AI_MEMORY_TEST_NO_PROCESS_GUARD` opt-out below: neither is
// reachable in a normal shipped run because neither is ever set outside
// a test harness's own env.
if let Ok(raw) = std::env::var("AI_MEMORY_TEST_FORCE_SIBLING_PIDS") {
return raw
.split(',')
.filter(|s| !s.trim().is_empty())
.filter_map(|s| s.trim().parse::<u32>().ok())
.map(sysinfo::Pid::from_u32)
.collect();
}
// Test opt-out. The destructive-command tests would otherwise flake
// non-deterministically: a dev box (and a parallel test run) almost always
// has some *other* `ai-memory` process alive, which the real scan rightly
crates/ai-memory-cli/src/server_profiles.rs
+802
View File
@@ -0,0 +1,802 @@
//! Named server profiles: per-repository routing of hook capture (#992).
//!
//! A machine can register several ai-memory servers under local names in
//! `<data_dir>/servers.toml`; a repository's `.ai-memory.toml` then selects one
//! with `server = "<name>"`. The marker is committed repository content — so
//! untrusted input — which is why it carries only a *name*: a cloned repository
//! can choose among servers this operator registered, never introduce a new
//! destination, and never see a credential.
//!
//! Every failure resolves to a [`Rejection`] and the caller drops the event.
//! Falling back to the install-time `--server-url` would deliver one team's
//! capture to another team's server, which is the exact outcome profiles exist
//! to prevent.
//!
//! Tokens live one per file under `<data_dir>/auth-tokens/<name>`, `0600`, and
//! never in `servers.toml`, a rendered hook config, or a process argv.
//!
//! Read on the hook hot path, so parsing uses `toml_edit` over this one file
//! and nothing else — no figment merge, no environment — for the same reason
//! `hook_spool::configured_server_url` does.
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use anyhow::{Context as _, Result, bail};
/// `<data_dir>/servers.toml`.
const REGISTRY_FILE: &str = "servers.toml";
/// `<data_dir>/auth-tokens/` — one owner-only file per profile.
const TOKEN_DIR: &str = "auth-tokens";
/// A registry is a handful of short tables; anything larger is not one.
const MAX_REGISTRY_BYTES: u64 = 64 * 1024;
const MAX_NAME_CHARS: usize = 64;
/// A validated profile name: `[a-z0-9][a-z0-9_-]{0,63}`, and not a Windows
/// reserved device name.
///
/// Parsed once, because the name becomes a file name under `auth-tokens/`: a
/// marker naming `../auth-token` must never reach the filesystem as a path,
/// and `nul` or `con` would open a device instead of a file on Windows. The
/// device names are refused on every platform so a registry stays portable.
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub(crate) struct ProfileName(String);
impl ProfileName {
pub(crate) fn parse(raw: &str) -> Option<Self> {
let mut chars = raw.chars();
let first = chars.next()?;
let valid = raw.len() <= MAX_NAME_CHARS
&& (first.is_ascii_lowercase() || first.is_ascii_digit())
&& chars.all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '-' || c == '_')
&& !ai_memory_core::is_dos_device_name(raw);
valid.then(|| Self(raw.to_owned()))
}
pub(crate) fn as_str(&self) -> &str {
&self.0
}
}
impl std::fmt::Display for ProfileName {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(&self.0)
}
}
/// One registered server.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct Profile {
/// Server URL, trailing slashes trimmed.
pub(crate) url: String,
/// Directories allowed to select this profile, as written (`~/` allowed).
pub(crate) roots: Vec<String>,
}
/// The parsed `servers.toml`. An absent file is an empty registry.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub(crate) struct Registry {
pub(crate) profiles: BTreeMap<ProfileName, Profile>,
}
/// Why a marker's `server` selection was refused. Each one drops the event.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum Rejection {
/// The marker's value is not a valid profile name.
InvalidName,
/// `servers.toml` is unreadable, oversized, or malformed.
InvalidRegistry,
/// No profile by that name is registered.
UnknownProfile,
/// The profile has no stored token.
NoToken,
/// Several profiles are registered and this one declares no `roots`, so
/// any repository could select it.
RootsRequired,
/// The marker sits outside every one of the profile's `roots`.
OutsideRoots,
/// A marker on the walk exists but could not be read, so whether it
/// selects a profile is unknown.
UnreadableMarker,
}
impl Rejection {
/// Stable, secret-free label for `hook --check-capture` and warnings.
pub(crate) fn as_str(self) -> &'static str {
match self {
Self::InvalidName => "rejected-invalid-profile-name",
Self::InvalidRegistry => "rejected-invalid-registry",
Self::UnknownProfile => "rejected-unknown-profile",
Self::NoToken => "rejected-no-token",
Self::RootsRequired => "rejected-roots-required",
Self::OutsideRoots => "rejected-outside-roots",
Self::UnreadableMarker => "rejected-unreadable-marker",
}
}
}
/// A selection that passed every check: where to deliver, and with what.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct ResolvedServer {
pub(crate) name: ProfileName,
pub(crate) url: String,
pub(crate) token: String,
}
fn registry_path(data_dir: &Path) -> PathBuf {
data_dir.join(REGISTRY_FILE)
}
fn token_path(data_dir: &Path, name: &ProfileName) -> PathBuf {
data_dir.join(TOKEN_DIR).join(name.as_str())
}
/// Load `servers.toml`. A missing file is an empty registry; anything else
/// that is not exactly the documented shape is an error, so a typo can only
/// ever make routing refuse, never make it guess.
pub(crate) fn load(data_dir: &Path) -> Result<Registry> {
read_bounded(&registry_path(data_dir))?.map_or_else(
|| Ok(Registry::default()),
|text| validate(&parse_doc(&text)?),
)
}
fn read_bounded(path: &Path) -> Result<Option<String>> {
let meta = match std::fs::metadata(path) {
Ok(meta) => meta,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(None),
Err(e) => return Err(e).with_context(|| format!("cannot read {REGISTRY_FILE}")),
};
if meta.len() > MAX_REGISTRY_BYTES {
bail!("{REGISTRY_FILE} exceeds {MAX_REGISTRY_BYTES} bytes");
}
std::fs::read_to_string(path)
.map(Some)
.with_context(|| format!("cannot read {REGISTRY_FILE}"))
}
fn parse_doc(text: &str) -> Result<toml_edit::DocumentMut> {
text.parse()
.with_context(|| format!("{REGISTRY_FILE} is not valid TOML"))
}
fn validate(doc: &toml_edit::DocumentMut) -> Result<Registry> {
let mut registry = Registry::default();
for (key, item) in doc.iter() {
if key != "servers" {
bail!("unknown top-level key `{key}` in {REGISTRY_FILE}");
}
let servers = item
.as_table_like()
.with_context(|| format!("`servers` in {REGISTRY_FILE} must be a table"))?;
for (raw_name, entry) in servers.iter() {
let name = ProfileName::parse(raw_name)
.with_context(|| format!("invalid profile name `{raw_name}`"))?;
let table = entry
.as_table_like()
.with_context(|| format!("profile `{name}` must be a table"))?;
let profile = parse_profile(table).with_context(|| format!("profile `{name}`"))?;
registry.profiles.insert(name, profile);
}
}
Ok(registry)
}
fn parse_profile(table: &dyn toml_edit::TableLike) -> Result<Profile> {
let mut url = None;
let mut roots = Vec::new();
for (key, value) in table.iter() {
match key {
"url" => {
url = Some(validate_url(
value.as_str().context("`url` must be a string")?,
)?);
}
"roots" => {
for root in value.as_array().context("`roots` must be an array")? {
let root = root
.as_str()
.context("every `roots` entry must be a string")?;
validate_root(root)?;
roots.push(root.to_owned());
}
}
other => bail!("unknown key `{other}`"),
}
}
Ok(Profile {
url: url.context("no `url`")?,
roots,
})
}
/// Accept only `http(s)://` URLs with a host, normalised without trailing
/// slashes — rejected here rather than on every hook request.
fn validate_url(raw: &str) -> Result<String> {
let trimmed = raw.trim().trim_end_matches('/');
let parsed = reqwest::Url::parse(trimmed).with_context(|| format!("`{raw}` is not a URL"))?;
if !matches!(parsed.scheme(), "http" | "https") || parsed.host_str().is_none_or(str::is_empty) {
bail!("`{raw}` is not an http:// or https:// URL with a host");
}
Ok(trimmed.to_owned())
}
/// A root must be absolute or `~/`-relative: a relative root would resolve
/// against whatever directory the hook happened to run in.
fn validate_root(raw: &str) -> Result<()> {
if raw == "~" || raw.starts_with("~/") || Path::new(raw).is_absolute() {
Ok(())
} else {
bail!("root `{raw}` must be absolute or start with `~/`")
}
}
fn expand_root(raw: &str, home: Option<&Path>) -> Option<PathBuf> {
let path = if raw == "~" {
home?.to_path_buf()
} else if let Some(rest) = raw.strip_prefix("~/") {
home?.join(rest)
} else {
PathBuf::from(raw)
};
Some(crate::marker::absolute_normalized(&path))
}
/// Read a profile's stored token. Trailing newline trimmed; blank is `None`.
pub(crate) fn read_token(data_dir: &Path, name: &ProfileName) -> Option<String> {
crate::config::read_trimmed_secret(&token_path(data_dir, name))
}
/// Resolve a marker's `server` selection.
///
/// `marker_dir` is the directory of the marker that declared the key, in the
/// same lexical namespace as the hook's cwd (see `marker::absolute_normalized`);
/// `roots` are compared against it component-wise, so `~/work/team-b` admits
/// `~/work/team-b/repo` but not `~/work/team-bb`.
pub(crate) fn resolve(
data_dir: &Path,
raw_name: &str,
marker_dir: &Path,
home: Option<&Path>,
) -> Result<ResolvedServer, Rejection> {
let name = ProfileName::parse(raw_name).ok_or(Rejection::InvalidName)?;
let registry = load(data_dir).map_err(|_| Rejection::InvalidRegistry)?;
let profile = registry
.profiles
.get(&name)
.ok_or(Rejection::UnknownProfile)?;
if profile.roots.is_empty() {
if registry.profiles.len() > 1 {
return Err(Rejection::RootsRequired);
}
} else {
let marker_dir = crate::marker::absolute_normalized(marker_dir);
let admitted = profile
.roots
.iter()
.filter_map(|root| expand_root(root, home))
.any(|root| marker_dir.starts_with(&root));
if !admitted {
return Err(Rejection::OutsideRoots);
}
}
with_token(data_dir, name, profile)
}
/// A registered profile's URL and stored token, without the `roots` check.
///
/// For a caller acting on a selection the hook already vetted against the
/// marker that made it — the backfill the hook spawned.
pub(crate) fn lookup(data_dir: &Path, name: &ProfileName) -> Result<ResolvedServer, Rejection> {
let registry = load(data_dir).map_err(|_| Rejection::InvalidRegistry)?;
let profile = registry
.profiles
.get(name)
.ok_or(Rejection::UnknownProfile)?;
with_token(data_dir, name.clone(), profile)
}
fn with_token(
data_dir: &Path,
name: ProfileName,
profile: &Profile,
) -> Result<ResolvedServer, Rejection> {
let token = read_token(data_dir, &name).ok_or(Rejection::NoToken)?;
Ok(ResolvedServer {
name,
url: profile.url.clone(),
token,
})
}
/// Register or replace a profile, keeping any hand-written comments and other
/// profiles in `servers.toml` intact, then store its token when one is given.
///
/// The file is validated before it is rewritten, so `add` refuses to build on
/// a registry the hook would already reject as a whole.
pub(crate) fn add(
data_dir: &Path,
name: &ProfileName,
url: &str,
roots: &[String],
token: Option<&str>,
) -> Result<AddOutcome> {
let url = validate_url(url)?;
let roots_were_omitted = roots.is_empty();
for root in roots {
validate_root(root)?;
}
let path = registry_path(data_dir);
let mut doc = parse_doc(&read_bounded(&path)?.unwrap_or_default())?;
let existing = validate(&doc)?.profiles.remove(name);
// Re-running `add` to rotate a token or fix a URL must not silently lift
// the roots restriction: omitted roots keep the ones already registered.
let roots = match &existing {
Some(profile) if roots_were_omitted => profile.roots.clone(),
_ => roots.to_vec(),
};
// A stored token belongs to the URL it was registered with. Once the URL
// changes it is discarded *before* the new URL is written, so there is no
// moment at which one server's credential is paired with another server.
let token = token.map(str::trim).filter(|t| !t.is_empty());
let url_changed = existing.as_ref().is_some_and(|profile| profile.url != url);
let token_discarded = url_changed && discard_token(data_dir, name)? && token.is_none();
let servers = doc
.entry("servers")
.or_insert_with(|| {
let mut table = toml_edit::Table::new();
table.set_implicit(true);
toml_edit::Item::Table(table)
})
.as_table_mut()
.with_context(|| format!("`servers` in {REGISTRY_FILE} must be a table"))?;
let mut entry = toml_edit::Table::new();
entry.insert("url", toml_edit::value(url));
if !roots.is_empty() {
let array: toml_edit::Array = roots.iter().map(String::as_str).collect();
entry.insert("roots", toml_edit::value(array));
}
servers.insert(name.as_str(), toml_edit::Item::Table(entry));
write_private(&path, &doc.to_string())?;
if let Some(token) = token {
crate::commands::path_util::create_private_dir(&data_dir.join(TOKEN_DIR))
.context("cannot create the token directory")?;
write_private(&token_path(data_dir, name), token)?;
}
Ok(AddOutcome {
token_discarded,
roots_kept: roots_were_omitted && !roots.is_empty(),
})
}
/// What `add` did beyond writing what it was given.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct AddOutcome {
/// The URL changed without a new token, so the old server's token was
/// removed and the profile refuses events until a token is added.
pub(crate) token_discarded: bool,
/// No roots were given, so the profile kept its registered ones.
pub(crate) roots_kept: bool,
}
/// Remove a profile and its token. Returns whether the profile existed.
pub(crate) fn remove(data_dir: &Path, name: &ProfileName) -> Result<bool> {
let path = registry_path(data_dir);
let mut existed = false;
if let Some(text) = read_bounded(&path)? {
let mut doc = parse_doc(&text)?;
existed = doc
.get_mut("servers")
.and_then(toml_edit::Item::as_table_like_mut)
.and_then(|servers| servers.remove(name.as_str()))
.is_some();
if existed {
write_private(&path, &doc.to_string())?;
}
}
Ok(discard_token(data_dir, name)? || existed)
}
/// Remove every stored profile token, keeping the registry itself.
///
/// `uninstall` removes the hooks, which are the only readers of these
/// tokens; leaving them on disk would strand live credentials the same way
/// the single hook token used to (#552). The registry holds no secrets and
/// survives, so a later reinstall lists each profile with `token: missing`
/// and refuses its events until a token is stored again.
///
/// # Errors
/// Propagates IO failures other than "not found".
pub(crate) fn clear_tokens(data_dir: &Path) -> std::io::Result<()> {
match std::fs::remove_dir_all(data_dir.join(TOKEN_DIR)) {
Ok(()) => Ok(()),
Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()),
Err(e) => Err(e),
}
}
/// Delete a profile's token file. Returns whether one was there.
fn discard_token(data_dir: &Path, name: &ProfileName) -> Result<bool> {
match std::fs::remove_file(token_path(data_dir, name)) {
Ok(()) => Ok(true),
Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(false),
Err(e) => Err(e).context("cannot remove token"),
}
}
/// Atomic replace through the canonical [`ai_memory_wiki::write_atomic`]. Its
/// temp file is created `0600`, so replacing an existing file never inherits
/// wider permissions from it.
fn write_private(path: &Path, contents: &str) -> Result<()> {
ai_memory_wiki::write_atomic(path, contents.as_bytes())
.map(|_| ())
.with_context(|| format!("cannot write {}", path.display()))
}
#[cfg(test)]
mod tests {
use super::*;
fn store_token(data_dir: &Path, name: &ProfileName, token: &str) -> Result<()> {
crate::commands::path_util::create_private_dir(&data_dir.join(TOKEN_DIR))?;
write_private(&token_path(data_dir, name), token)
}
fn name(raw: &str) -> ProfileName {
ProfileName::parse(raw).unwrap()
}
#[test]
fn profile_names_cannot_become_paths() {
for bad in [
"",
"../auth-token",
"a/b",
"Team",
"-x",
"a b",
"a.b",
&"a".repeat(65),
"con",
"nul",
"aux",
"prn",
"com1",
"lpt9",
] {
assert!(
ProfileName::parse(bad).is_none(),
"{bad:?} must be rejected"
);
}
for good in [
"team-a", "b", "0", "team_b-2", "console", "com", "com10", "nul-1",
] {
assert!(
ProfileName::parse(good).is_some(),
"{good:?} must be accepted"
);
}
}
#[test]
fn a_missing_registry_is_empty_and_an_unknown_profile_is_refused() {
let dd = tempfile::tempdir().unwrap();
assert_eq!(load(dd.path()).unwrap(), Registry::default());
assert_eq!(
resolve(dd.path(), "team-a", dd.path(), None),
Err(Rejection::UnknownProfile)
);
}
#[test]
fn a_malformed_registry_refuses_every_selection() {
let dd = tempfile::tempdir().unwrap();
for text in [
"[servers.a]\nurl = \"https://a.example\"\ntoken = \"leak\"\n",
"server_url = \"https://a.example\"\n",
"[servers.a]\nurl = \"ftp://a.example\"\n",
"[servers.a]\nroots = [\"~/x\"]\n",
"[servers.a]\nurl = \"https://a.example\"\nroots = [\"relative\"]\n",
"[servers.\"Bad\"]\nurl = \"https://a.example\"\n",
"not toml [",
] {
std::fs::write(dd.path().join(REGISTRY_FILE), text).unwrap();
store_token(dd.path(), &name("a"), "tok").unwrap();
assert_eq!(
resolve(dd.path(), "a", dd.path(), None),
Err(Rejection::InvalidRegistry),
"{text}"
);
}
}
#[test]
fn an_oversized_registry_is_refused() {
let dd = tempfile::tempdir().unwrap();
let mut text = String::from("[servers.a]\nurl = \"https://a.example\"\n");
text.push_str(&"#".repeat(usize::try_from(MAX_REGISTRY_BYTES).unwrap()));
std::fs::write(dd.path().join(REGISTRY_FILE), text).unwrap();
assert!(load(dd.path()).is_err());
}
/// A platform-absolute path string under `base`.
fn under(base: &Path, rel: &str) -> String {
base.join(rel).to_string_lossy().into_owned()
}
#[test]
fn a_single_profile_without_roots_resolves_with_its_own_token() {
let dd = tempfile::tempdir().unwrap();
add(
dd.path(),
&name("a"),
"https://a.example/",
&[],
Some("tok-a"),
)
.unwrap();
let resolved = resolve(dd.path(), "a", &dd.path().join("anywhere"), None).unwrap();
assert_eq!(resolved.url, "https://a.example");
assert_eq!(resolved.token, "tok-a");
}
#[test]
fn a_profile_without_a_token_is_refused() {
let dd = tempfile::tempdir().unwrap();
add(dd.path(), &name("a"), "https://a.example", &[], None).unwrap();
assert_eq!(
resolve(dd.path(), "a", &dd.path().join("anywhere"), None),
Err(Rejection::NoToken)
);
}
#[test]
fn roots_are_required_once_a_second_profile_exists() {
let dd = tempfile::tempdir().unwrap();
let work = dd.path().join("work");
add(
dd.path(),
&name("a"),
"https://a.example",
&[],
Some("tok-a"),
)
.unwrap();
add(
dd.path(),
&name("b"),
"https://b.example",
&[under(&work, "b")],
Some("tok-b"),
)
.unwrap();
assert_eq!(
resolve(dd.path(), "a", &work.join("a"), None),
Err(Rejection::RootsRequired)
);
assert!(resolve(dd.path(), "b", &work.join("b").join("repo"), None).is_ok());
}
#[test]
fn roots_match_whole_components_and_expand_home() {
let dd = tempfile::tempdir().unwrap();
let home = dd.path().join("home");
add(
dd.path(),
&name("b"),
"https://b.example",
&["~/work/team-b".to_owned()],
Some("tok-b"),
)
.unwrap();
let team_b = home.join("work").join("team-b");
assert!(resolve(dd.path(), "b", &team_b, Some(&home)).is_ok());
assert!(resolve(dd.path(), "b", &team_b.join("x").join("y"), Some(&home)).is_ok());
for outside in [
home.join("work").join("team-bb"),
home.join("work"),
team_b.join("..").join("team-a"),
dd.path().join("elsewhere").join("work").join("team-b"),
] {
assert_eq!(
resolve(dd.path(), "b", &outside, Some(&home)),
Err(Rejection::OutsideRoots),
"{}",
outside.display()
);
}
assert_eq!(
resolve(dd.path(), "b", &team_b, None),
Err(Rejection::OutsideRoots),
"a `~` root cannot match when home is unknown"
);
}
#[test]
fn add_preserves_other_profiles_and_comments_and_remove_deletes_the_token() {
let dd = tempfile::tempdir().unwrap();
let root_a = toml_edit::Value::from(under(dd.path(), "a"));
std::fs::write(
dd.path().join(REGISTRY_FILE),
format!(
"# managed by hand\n[servers.a]\nurl = \"https://a.example\"\nroots = [{root_a}]\n"
),
)
.unwrap();
add(
dd.path(),
&name("b"),
"https://b.example",
&[under(dd.path(), "b")],
Some("tok-b"),
)
.unwrap();
let text = std::fs::read_to_string(dd.path().join(REGISTRY_FILE)).unwrap();
assert!(text.contains("# managed by hand"), "{text}");
assert!(
!text.contains("tok-b"),
"tokens never go in the registry: {text}"
);
let registry = load(dd.path()).unwrap();
assert_eq!(registry.profiles.len(), 2);
assert!(remove(dd.path(), &name("b")).unwrap());
assert!(read_token(dd.path(), &name("b")).is_none());
assert_eq!(load(dd.path()).unwrap().profiles.len(), 1);
assert!(!remove(dd.path(), &name("b")).unwrap());
}
/// Moving a profile to a new URL without a new token must not pair the
/// old server's token with the new server.
#[test]
fn changing_the_url_without_a_token_discards_the_old_token() {
let dd = tempfile::tempdir().unwrap();
add(
dd.path(),
&name("b"),
"https://old.example",
&[],
Some("tok-old"),
)
.unwrap();
let outcome = add(dd.path(), &name("b"), "https://new.example", &[], None).unwrap();
assert!(outcome.token_discarded);
assert_eq!(read_token(dd.path(), &name("b")), None);
assert_eq!(
resolve(dd.path(), "b", dd.path(), None),
Err(Rejection::NoToken)
);
let outcome = add(
dd.path(),
&name("b"),
"https://new.example",
&[],
Some("tok-new"),
)
.unwrap();
assert!(!outcome.token_discarded);
assert_eq!(
resolve(dd.path(), "b", dd.path(), None).unwrap().token,
"tok-new"
);
let outcome = add(dd.path(), &name("b"), "https://new.example/", &[], None).unwrap();
assert!(!outcome.token_discarded, "the same URL keeps its token");
assert_eq!(
read_token(dd.path(), &name("b")).as_deref(),
Some("tok-new")
);
}
/// Re-running `add` without `--root` (for a token rotation) keeps the
/// registered roots instead of lifting the restriction.
#[test]
fn omitted_roots_keep_the_registered_ones() {
let dd = tempfile::tempdir().unwrap();
let root = under(dd.path(), "b");
add(
dd.path(),
&name("b"),
"https://b.example",
std::slice::from_ref(&root),
Some("t1"),
)
.unwrap();
let outcome = add(dd.path(), &name("b"), "https://b.example", &[], Some("t2")).unwrap();
assert!(outcome.roots_kept);
assert_eq!(
load(dd.path()).unwrap().profiles[&name("b")].roots,
vec![root]
);
let other = under(dd.path(), "c");
let outcome = add(
dd.path(),
&name("b"),
"https://b.example",
std::slice::from_ref(&other),
None,
)
.unwrap();
assert!(!outcome.roots_kept);
assert_eq!(
load(dd.path()).unwrap().profiles[&name("b")].roots,
vec![other]
);
}
/// Uninstall removes the tokens and nothing else; the registry stays so
/// the profiles are listed as tokenless rather than forgotten.
#[test]
fn clear_tokens_removes_every_token_and_keeps_the_registry() {
let dd = tempfile::tempdir().unwrap();
add(
dd.path(),
&name("a"),
"https://a.example",
&[],
Some("tok-a"),
)
.unwrap();
add(
dd.path(),
&name("b"),
"https://b.example",
&[under(dd.path(), "b")],
Some("tok-b"),
)
.unwrap();
clear_tokens(dd.path()).unwrap();
assert!(!dd.path().join(TOKEN_DIR).exists());
assert_eq!(read_token(dd.path(), &name("a")), None);
assert_eq!(read_token(dd.path(), &name("b")), None);
assert_eq!(load(dd.path()).unwrap().profiles.len(), 2);
assert_eq!(
resolve(dd.path(), "b", &dd.path().join("b"), None),
Err(Rejection::NoToken),
"the profile is still registered, only its token is gone"
);
clear_tokens(dd.path()).expect("clearing an already-clear store is not an error");
}
#[test]
fn add_refuses_to_build_on_an_invalid_registry() {
let dd = tempfile::tempdir().unwrap();
std::fs::write(dd.path().join(REGISTRY_FILE), "stray = 1\n").unwrap();
assert!(add(dd.path(), &name("a"), "https://a.example", &[], None).is_err());
assert_eq!(
std::fs::read_to_string(dd.path().join(REGISTRY_FILE)).unwrap(),
"stray = 1\n"
);
}
#[cfg(unix)]
#[test]
fn tokens_are_owner_only_even_when_replacing_a_wider_file() {
use std::os::unix::fs::PermissionsExt as _;
let dd = tempfile::tempdir().unwrap();
store_token(dd.path(), &name("a"), "old").unwrap();
let path = token_path(dd.path(), &name("a"));
std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o644)).unwrap();
store_token(dd.path(), &name("a"), "new").unwrap();
assert_eq!(
std::fs::metadata(&path).unwrap().permissions().mode() & 0o777,
0o600
);
assert_eq!(
std::fs::metadata(dd.path().join(TOKEN_DIR))
.unwrap()
.permissions()
.mode()
& 0o777,
0o700
);
assert_eq!(read_token(dd.path(), &name("a")).as_deref(), Some("new"));
}
}
crates/ai-memory-cli/tests/fixtures/opencode2-plugin.mjs
Vendored
+164
View File
@@ -0,0 +1,164 @@
// Execute the generated adapter, rather than only checking template substrings.
// Run by `cargo test` (opencode2_plugin_passes_the_node_host_fixture), or by hand:
// AI_MEMORY_DATA_DIR=<empty dir> node --experimental-strip-types opencode2-plugin.mjs /absolute/plugin.ts
import assert from "node:assert/strict";
import { pathToFileURL } from "node:url";
import { readdirSync, readFileSync } from "node:fs";
import { join } from "node:path";
const requests = [];
globalThis.fetch = async (input, options = {}) => {
// Like a real fetch, a request cancelled before it is sent never arrives.
options.signal?.throwIfAborted();
const url = new URL(input);
requests.push({ url, payload: options.body ? JSON.parse(options.body) : undefined });
return new Response(url.pathname === "/handoff" ? "Remember the verified handoff." : "{}");
};
const { default: plugin } = await import(pathToFileURL(process.argv[2]).href);
function host(directory, records) {
const hooks = new Map();
const events = [];
let wake;
return {
hooks,
emit(event) {
events.push(event);
wake?.();
},
ctx: {
location: { directory },
session: {
async get({ sessionID }) {
const record = records.get(sessionID);
assert.ok(record, `unknown session ${sessionID}`);
return record;
},
async hook(name, callback) {
hooks.set(`session.${name}`, callback);
return { async dispose() {} };
},
},
tool: {
async hook(name, callback) {
hooks.set(`tool.${name}`, callback);
return { async dispose() {} };
},
},
event: {
async *subscribe({ signal }) {
signal.addEventListener("abort", () => wake?.(), { once: true });
while (!signal.aborted) {
if (!events.length) await new Promise((resolve) => { wake = resolve; });
while (events.length) yield events.shift();
}
},
},
},
};
}
async function until(predicate, message) {
const deadline = performance.now() + 5000;
while (!predicate()) {
assert.ok(performance.now() < deadline, message);
await new Promise((resolve) => setTimeout(resolve, 10));
}
}
const root = process.cwd();
const info = (id, parentID) => ({ id, title: id, location: { directory: root }, parentID });
const records = new Map([["root-a", info("root-a")], ["root-b", { ...info("root-b"), location: { directory: root + "/beta" } }],
["child", info("child", "root-a")]]);
const a = host(root, records);
const b = host(root + "/beta", records);
const disposeA = await plugin.setup(a.ctx);
const disposeB = await plugin.setup(b.ctx);
try {
// Existing/resumed sessions may never emit session.created after plugin load.
await a.hooks.get("session.prompt")({ sessionID: "root-a", messageID: "m1", prompt: { text: "remember alpha" } });
const first = { sessionID: "root-a", system: [] };
const next = { sessionID: "root-a", system: [] };
await a.hooks.get("session.context")(first);
await a.hooks.get("session.context")(next);
assert.deepEqual(first.system, next.system);
assert.equal(first.system.length, 1);
assert.equal(requests.filter((r) => r.url.pathname === "/handoff").length, 1, "claim once, inject repeatedly");
await a.hooks.get("session.prompt")({ sessionID: "child", messageID: "c1", prompt: { text: "child work" } });
const child = { sessionID: "child", system: [] };
await a.hooks.get("session.context")(child);
assert.equal(child.system.length, 0);
assert.equal(requests.filter((r) => r.url.pathname === "/handoff").length, 1, "child must not claim");
await a.hooks.get("tool.execute.after")({ sessionID: "root-a", tool: "example", id: "t1", input: {},
status: "completed", result: { content: "content-only tool evidence" } });
for (const type of ["session.execution.succeeded", "session.execution.failed", "session.execution.interrupted"]) {
a.emit({ type, data: { sessionID: "root-a" } });
b.emit({ type, data: { sessionID: "root-a" } });
}
a.emit({ type: "session.execution.succeeded", data: { sessionID: "child" } });
await until(() => requests.filter((r) => r.url.searchParams.get("event") === "stop").length === 4, "completion events delivered");
const stops = requests.filter((r) => r.url.searchParams.get("event") === "stop");
assert.equal(stops.filter((r) => r.payload.turn_checkpoint).length, 3);
assert.equal(stops.find((r) => r.payload.sessionID === "child").payload.agent_id, "root-a");
assert.ok(requests.some((r) => r.payload?.output === "content-only tool evidence"));
assert.ok(!requests.some((r) => r.url.searchParams.get("event") === "session-end"));
await b.hooks.get("session.prompt")({ sessionID: "root-b", messageID: "b1", prompt: { text: "beta" } });
await b.hooks.get("session.context")({ sessionID: "root-b", system: [] });
await disposeA();
assert.ok(requests.some((r) => r.url.searchParams.get("event") === "session-end" && r.payload.sessionID === "root-a"));
assert.ok(!requests.some((r) => r.url.searchParams.get("event") === "session-end" && r.payload.sessionID === "root-b"), "location cleanup must not close another instance");
assert.equal(requests.find((r) => r.url.searchParams.get("event") === "session-end" && r.payload.sessionID === "child").payload.agent_id, "root-a", "child close must retain ancestry to suppress automatic handoffs");
console.log("PASS: resumed sessions, retained handoff, child isolation, terminal events, content-only output, per-location cleanup");
} finally {
await disposeB();
}
const spool = join(process.env.AI_MEMORY_DATA_DIR, "hook-spool");
const realNow = Date.now;
globalThis.fetch = async () => { throw new Error("offline fixture"); };
Date.now = () => 1800000000000;
try {
const offlineA = host(root, records);
const offlineB = host(root + "/beta", records);
const closeA = await plugin.setup(offlineA.ctx);
const closeB = await plugin.setup(offlineB.ctx);
await Promise.all([
offlineA.hooks.get("session.prompt")({ sessionID: "root-a", messageID: "offline-a", prompt: { text: "offline alpha" } }),
offlineB.hooks.get("session.prompt")({ sessionID: "root-b", messageID: "offline-b", prompt: { text: "offline beta" } }),
]);
await Promise.all([closeA(), closeB()]);
const entries = readdirSync(spool).filter((name) => name.endsWith(".json")).map((name) => JSON.parse(readFileSync(join(spool, name), "utf8")));
for (const id of ["root-a", "root-b"]) {
assert.ok(entries.some((entry) => new URL(entry.url).searchParams.get("event") === "session-start" && JSON.parse(entry.body).sessionID === id), `same-millisecond spool must retain ${id}`);
}
assert.ok(entries.every((entry) => new URL(entry.url).searchParams.has("ingest_key")), "stable delivery keys survive spooling");
console.log("PASS: concurrent location spools do not overwrite same-millisecond events");
} finally {
Date.now = realNow;
}
// A server which accepts requests but never answers must not hold plugin unload
// behind every queued request's individual timeout.
globalThis.fetch = (_url, { signal } = {}) => new Promise((_resolve, reject) => {
if (signal.aborted) reject(signal.reason);
else signal.addEventListener("abort", () => reject(signal.reason), { once: true });
});
const stalled = host(root, records);
const closeStalled = await plugin.setup(stalled.ctx);
await stalled.hooks.get("session.prompt")({ sessionID: "root-a", messageID: "stalled", prompt: { text: "bounded shutdown" } });
for (let n = 0; n < 60; n++) {
await stalled.hooks.get("tool.execute.after")({ sessionID: "root-a", tool: "example", id: `pending-${n}`, input: {}, status: "completed", result: { content: "queued evidence" } });
}
stalled.emit({ type: "session.text.ended", data: { sessionID: "root-a", text: "last response" } });
stalled.emit({ type: "session.execution.succeeded", data: { sessionID: "root-a" } });
await new Promise((resolve) => setImmediate(resolve));
const shutdown = performance.now();
// The plugin's drain timers are unref'd; a real host process stays alive.
const host_alive = setInterval(() => {}, 1000);
await closeStalled();
clearInterval(host_alive);
assert.ok(performance.now() - shutdown < 5000, "unload must not wait for the entire network timeout backlog");
console.log("PASS: unload cancels stalled capture and remains bounded");
crates/ai-memory-cli/tests/suite/autoscope_env.rs
+32
View File
@@ -22,6 +22,26 @@ fn bin() -> &'static str {
env!("CARGO_BIN_EXE_ai-memory")
}
/// Logged by `start_watcher` when `serve --no-watcher` is honoured.
const WATCHER_DISABLED: &str = "watcher disabled by --no-watcher";
/// Logged by `start_watcher` when it actually installs an FSEvents/inotify
/// instance. Must not appear in these children: they do not exercise watching.
const WATCHER_STARTED: &str = "starting wiki watcher";
/// Pin that the spawned binary opted out of the wiki watcher. The disable
/// line is logged before the `[auto_scope]` mode needle these tests wait on,
/// so a successful match without it would mean the flag was ignored.
fn assert_watcher_opted_out(stderr: &str) {
assert!(
stderr.contains(WATCHER_DISABLED),
"spawned serve must log that the watcher was opted out.\nstderr:\n{stderr}"
);
assert!(
!stderr.contains(WATCHER_STARTED),
"spawned serve must not install a wiki watcher.\nstderr:\n{stderr}"
);
}
/// Spawn the binary with the given env vars and stream stderr until either
/// `needle` appears in a line or `timeout` elapses. Always kills the child
/// before returning. Returns `(matched_line_or_none, all_stderr_captured)`.
@@ -41,6 +61,11 @@ fn spawn_and_wait_for_log(
"http",
"--bind",
"127.0.0.1:0",
// None of these tests watch the wiki. The watcher is a
// machine-global FSEvents/inotify instance; concurrent serve
// children can exhaust it (#745). Opt out so the child never
// takes one.
"--no-watcher",
"--data-dir",
])
.arg(tmp.path())
@@ -96,6 +121,9 @@ fn spawn_and_wait_for_log(
// Pump thread terminates once stderr closes (post-kill); collect what
// it captured.
let all_stderr = pump.join().unwrap_or_default();
if matched.is_some() {
assert_watcher_opted_out(&all_stderr);
}
(matched, all_stderr)
}
@@ -111,6 +139,10 @@ fn spawn_and_wait_for_exit(envs: &[(&str, &str)], timeout: Duration) -> (bool, S
"http",
"--bind",
"127.0.0.1:0",
// Same opt-out as [`spawn_and_wait_for_log`]: invalid config
// fails before the watcher is installed, but keeping the flag
// here means every serve spawn in this file is hermetic.
"--no-watcher",
"--data-dir",
])
.arg(tmp.path())

Some files were not shown because too many files have changed in this diff Show More