Merge remote-tracking branch 'origin/main' into audit/pr-223-fix

# Conflicts:
#	CHANGELOG.md
This commit is contained in:
AkitaOnRails
2026-07-23 11:34:13 -03:00
23 changed files with 1103 additions and 90 deletions
+16
View File
@@ -8,6 +8,22 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
## [Unreleased]
### Added
- Opt-in assistant/Stop capture for Claude Code (#196). When BOTH the server
(`capture_assistant = true` / `AI_MEMORY_CAPTURE_ASSISTANT=true`) and the
client (`install-hooks --agent claude-code --capture-assistant`) opt in, a
Claude Code `Stop` event carries a sanitized, 2 KB-capped excerpt of the
assistant's final turn as the Stop body. The client sanitizes with the
built-in patterns and truncates before the excerpt ever touches the spool or
wire, splicing a versioned `_ai_memory_assistant` marker into the body and a
`capture_assistant=1` flag onto the event URL; the server re-scrubs with its
configured `[sanitize]` patterns and re-enforces the 2 KB cap at the
persistence boundary (never trusting the client's length). Off by default,
and any gate failure (server off, wrong agent/event, malformed or
future-versioned marker, empty excerpt) degrades to an empty Stop with the
same `202` response. Assistant text is privacy-sensitive — see `SECURITY.md`
for what it can contain and where it flows. Script-fallback installs cannot
sanitize the field and drop the whole Stop event instead; move to a native
install to capture it.
- Managed workstreams now support Kimi Code through `ai-memory run kimi`.
Returning runs resume the linked native session with `--session <id>`;
fresh sessions are discovered post-exit by exact checkout match through
+1 -1
View File
@@ -22,7 +22,7 @@
| macOS | Supported | Workspace tests run in CI; tagged releases publish native `ai-memory-macos-aarch64.tar.gz` and `ai-memory-macos-x86_64.tar.gz` binaries. The native binary is the recommended path on Apple Silicon. See [`docs/macos.md`](docs/macos.md). |
| Windows via WSL2 | Supported | Use the Linux install path inside WSL2 when the agent runs there. |
| Native Windows | Experimental | Tagged releases publish `ai-memory-windows-x86_64.zip` with `ai-memory.exe`; Docker Desktop wrapper and source builds are also available. Local supported profiles default to host-native hook commands; Claude Code may use its Windows exec form, while other agents use native single command strings matching their hook schema. PowerShell/Git Bash scripts are compatibility fallbacks. See [`docs/windows.md`](docs/windows.md). |
| Claude Code | Supported | MCP config + lifecycle hooks; native commands enforce capture exclusions. |
| Claude Code | Supported | MCP config + lifecycle hooks; native commands enforce capture exclusions. Optionally captures the assistant's final turn on `Stop` when installed with `--capture-assistant` and the server enables `capture_assistant` (double opt-in, off by default). |
| Codex | Supported | MCP config + lifecycle hooks; native commands enforce capture exclusions. No automatic true session-end hook, so run `ai-memory finalize-session` when you need a final summary/handoff. |
| Devin CLI | Supported | MCP config + lifecycle hooks. Hooks use Devin's `PostCompaction` event, inject handoffs via `hookSpecificOutput.additionalContext`, and omit subagent events because Devin does not expose them. |
| OpenCode | Supported | Remote MCP config + generated TypeScript plugin; generated plugin enforces capture exclusions. |
+26 -1
View File
@@ -43,6 +43,27 @@ the project is and is not designed to defend against.
`(workspace_id, project_id)`. A purge operation for project A cannot
delete files that also belong to project B.
- **Assistant/Stop capture is opt-in and sanitized (#196).** The assistant's
final turn is never persisted by default. Storing it requires a **double
opt-in** — `capture_assistant` on the server and `install-hooks
--capture-assistant` on the client. When enabled, be aware that:
- The excerpt is sanitized twice — the client scrubs with the built-in
patterns *before* it reaches the spool or wire, and the server re-scrubs
with its configured `[sanitize]` patterns before storing. Operator
`extra_patterns` run only on the server side, so a secret matched only by an
`extra_patterns` rule may still sit in the excerpt on the client spool/wire
before it reaches the server. Client-side redactions are irreversible: the
server's `allowlist` cannot restore text the client already replaced with
`[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 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.
- The excerpt can quote code, secrets, or content from paths ai-memory never
sees; the `Sanitizer` is a best-effort credential strip, not a guarantee
(see the injection note below).
### Out of scope for v1
- **Multi-tenant authentication and authorisation.** There is one bearer
@@ -53,7 +74,11 @@ the project is and is not designed to defend against.
controls, etc.).
- **MCP tool-call injection via agent output.** The privacy strip
(`Sanitizer`) removes obvious credential patterns from hook payloads, but
it is not a comprehensive injection fence.
it is not a comprehensive injection fence. This applies to the opt-in
assistant/Stop excerpt too: it is untrusted text (it can echo whatever a
tool put in the assistant's response) that, once captured, flows into the
consolidation/reviewer prompts. Enabling `capture_assistant` widens this
surface — leave it off unless you accept that trade-off.
- **Denial of service.** The server is not hardened against a malicious local
actor hammering it with requests.
+14
View File
@@ -1348,6 +1348,13 @@ pub struct HookArgs {
/// Inspect capture policy without spooling, draining, or contacting the server.
#[arg(long)]
pub check_capture: bool,
/// Opt in to assistant/Stop capture: on a Claude Code `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
/// `capture_assistant`. No effect on other events (#196).
#[arg(long)]
pub capture_assistant: bool,
}
/// Arguments for hidden `hook-drain`.
@@ -1416,6 +1423,13 @@ pub struct InstallHooksArgs {
/// identical to prior behavior.
#[arg(long, value_enum, default_value_t = ProjectStrategyArg::Basename)]
pub project_strategy: 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.
#[arg(long)]
pub capture_assistant: bool,
}
/// Arguments for `install-mcp`.
+38 -12
View File
@@ -378,16 +378,29 @@ where
return Ok(());
}
};
// Defensively drop any raw assistant-message field (e.g. Claude Code's
// `last_assistant_message` on Stop) BEFORE it can reach the local spool or
// the wire (#196). Optional capture remains disabled, so stripping it
// pre-spool closes a raw-text exposure with no persisted-data behavior
// change. Reserialize only when we actually removed something, so unrelated
// events keep byte-exact spool
// bodies (see `native_hook_accepts_plain_and_bom_prefixed_json`).
if ai_memory_hooks::strip_assistant_message_raw(&mut json) {
payload = serde_json::to_string(&json)?;
}
// Assistant/Stop capture (#196). On an opted-in install
// (`install-hooks --capture-assistant`), extract the assistant message,
// sanitize + cap it, and splice the versioned `_ai_memory_assistant` marker
// into the body. Otherwise just strip any raw assistant field defensively —
// the field is never persisted until the server-gated opt-in accepts it.
// Reserialize only when the JSON actually changed, so unrelated events keep
// byte-exact spool bodies (see `native_hook_accepts_plain_and_bom_prefixed_json`).
let capture_assistant = if args.capture_assistant {
let transform = ai_memory_hooks::transform_for_client(
&mut json,
AgentKind::from_wire(&args.agent),
ai_memory_hooks::HookEvent::parse(&args.event),
);
if transform.changed {
payload = serde_json::to_string(&json)?;
}
transform.captured
} else {
if ai_memory_hooks::strip_assistant_message_raw(&mut json) {
payload = serde_json::to_string(&json)?;
}
false
};
let (policy_cwd, canonical_session_id) = hook_context(&args.agent, &json);
let policy = policy_cwd.as_deref().map(capture_policy);
let tool_event = is_tool_event(&args.event);
@@ -461,9 +474,17 @@ where
.ok()
.flatten()
.is_some();
// Append the opt-in capture flag outside the marker query so the server can
// gate on it (#196). Only present when a valid protocol was actually spliced
// into this event's body, so the flag and the marker always travel together.
let capture_qs = if capture_assistant {
"&capture_assistant=1"
} else {
""
};
let event_url = format!(
"{base}/hook?event={}&agent={}{}",
args.event, args.agent, hook_qs
"{base}/hook?event={}&agent={}{}{}",
args.event, args.agent, hook_qs, capture_qs
);
let entry = hook_spool::entry_for(
event_url,
@@ -691,6 +712,7 @@ mod tests {
auth_token: None,
project_strategy: None,
check_capture: false,
capture_assistant: false,
}
}
@@ -1161,6 +1183,7 @@ mod tests {
auth_token: None,
project_strategy: None,
check_capture: false,
capture_assistant: false,
};
run_with_payload(
@@ -1202,6 +1225,7 @@ mod tests {
auth_token: None,
project_strategy: None,
check_capture: false,
capture_assistant: false,
};
run_with_payload(
@@ -1242,6 +1266,7 @@ mod tests {
auth_token: None,
project_strategy: None,
check_capture: false,
capture_assistant: false,
};
run_with_payload(Some(data_dir), args, "{}".into(), &mut stdout, |_path| {
@@ -1268,6 +1293,7 @@ mod tests {
auth_token: None,
project_strategy: None,
check_capture: false,
capture_assistant: false,
};
run_with_payload(
@@ -944,6 +944,42 @@ mod tests {
assert_eq!(loaded.body, body, "clean body must stay byte-exact");
}
#[test]
fn drain_preserves_capture_flag_and_protocol_round_trip() {
// An opted-in Stop entry carries `capture_assistant=1` on the URL and the
// sanitized `_ai_memory_assistant` marker in the body. The drain must
// preserve BOTH through spool → load → batch: the load-time strip only
// targets the raw `last_assistant_message`, never the protocol (#196).
let tmp = tempfile::tempdir().unwrap();
let spool = spool_dir(tmp.path());
let entry = entry_for(
"https://x/hook?event=stop&agent=claude-code&capture_assistant=1".into(),
r#"{"session_id":"s","_ai_memory_assistant":{"version":1,"excerpt":"done"}}"#.into(),
None,
false,
);
enqueue(&spool, &entry).unwrap();
let path = list_entries(&spool).unwrap().into_iter().next().unwrap();
let mut result = DrainResult::default();
let loaded = load_live_entry(&path, &mut result).expect("entry is live");
assert!(loaded.url.contains("capture_assistant=1"), "flag dropped");
assert!(
loaded.body.contains("_ai_memory_assistant"),
"protocol dropped"
);
let batch = batch_payload(&[(path, loaded)]).expect("batch built");
assert!(
batch.contains("capture_assistant=1"),
"batch lost the capture flag: {batch}"
);
assert!(
batch.contains("_ai_memory_assistant"),
"batch lost the protocol: {batch}"
);
}
#[cfg(unix)]
#[test]
fn enqueue_creates_spool_dir_private() {
@@ -219,6 +219,18 @@ pub fn run(config: &Config, args: InstallHooksArgs) -> Result<()> {
"[ai-memory] selected shell/PowerShell compatibility path does not enforce capture-policy v1; use a native platform selection or generated integration."
);
}
// Assistant/Stop capture is Claude Code + native-platform only (#196). No
// silent fallback: bail so an operator on a script-fallback platform or a
// different agent is told the flag has no effect instead of installing a
// command whose capture would be silently dropped.
if args.capture_assistant && !capture_assistant_allowed(args.agent) {
anyhow::bail!(
"--capture-assistant requires --agent claude-code on a native hook platform \
(PosixNative/WindowsNative). The current selection uses the script fallback or a \
different agent, where the opt-in cannot take effect. Remove --capture-assistant or \
switch to a native Claude Code install."
);
}
if args.apply {
return match args.agent {
AgentChoice::OpenCode => apply_to_opencode_plugin(&server_url, auth, &args),
@@ -290,6 +302,7 @@ pub fn run(config: &Config, args: InstallHooksArgs) -> Result<()> {
&config.data_dir,
strategy,
&settings_path,
args.capture_assistant,
)
}
AgentChoice::Codex => {
@@ -790,6 +803,14 @@ fn overlay_event_hooks(
/// Mutate `~/.claude/settings.json` in place: replace the hook entries
/// ai-memory cares about (`CLAUDE_CODE_EVENTS`); preserve every other hook the
/// user has wired up to other tools.
/// Whether `--capture-assistant` may take effect for this agent + platform
/// (#196): Claude Code on a native hook platform only. Any other agent or a
/// script-fallback platform cannot honor the opt-in, so the installer bails
/// instead of enabling it silently.
fn capture_assistant_allowed(agent: AgentChoice) -> bool {
matches!(agent, AgentChoice::ClaudeCode) && local_hook_policy_v1_supported()
}
fn apply_to_claude_code_settings(
hooks_dir: &Path,
server_url: &str,
@@ -810,6 +831,7 @@ fn apply_to_claude_code_settings(
auth_token,
Some(data_dir),
strategy,
args.capture_assistant,
);
let our_hooks = payload
.get("hooks")
@@ -2912,6 +2934,7 @@ fn render_claude_code(
data_dir: &Path,
project_strategy: Option<&str>,
settings_path: &Path,
capture_assistant: bool,
) -> Result<()> {
// Soft check: warn (don't bail) if a script is missing. The user
// may be running this command inside docker against a host path
@@ -2936,6 +2959,7 @@ fn render_claude_code(
auth_token,
Some(data_dir),
project_strategy,
capture_assistant,
);
let serialized =
serde_json::to_string_pretty(&payload).context("serializing claude code hook config")?;
@@ -3227,6 +3251,37 @@ mod tests {
use std::process::Command;
use tempfile::TempDir;
#[test]
fn capture_assistant_allowed_only_for_claude_native() {
use crate::cli::AgentChoice::*;
// Every non-Claude agent is rejected regardless of platform (#196): the
// opt-in cannot take effect for them, so the installer must bail.
for agent in [
Codex,
Cursor,
GeminiCli,
OpenCode,
Pi,
Omp,
Openclaw,
AntigravityCli,
Grok,
Zero,
Devin,
KimiCode,
] {
assert!(
!capture_assistant_allowed(agent),
"{agent:?} must not allow --capture-assistant"
);
}
// Claude Code tracks the native-platform gate exactly.
assert_eq!(
capture_assistant_allowed(ClaudeCode),
local_hook_policy_v1_supported()
);
}
#[cfg(unix)]
fn bash_program_for_installer_test() -> Option<std::path::PathBuf> {
Some(std::path::PathBuf::from("bash"))
@@ -3365,6 +3420,7 @@ mod tests {
fn default_hook_args() -> InstallHooksArgs {
InstallHooksArgs {
agent: AgentChoice::OpenCode,
capture_assistant: false,
hooks_dir: None,
server_url: None,
auth_token: None,
@@ -3684,6 +3740,7 @@ mod tests {
.unwrap();
let args = InstallHooksArgs {
agent: AgentChoice::Zero,
capture_assistant: false,
config_file: Some(path.clone()),
..default_hook_args()
};
@@ -4503,6 +4560,7 @@ model = "gpt-5"
let tmp = TempDir::new().unwrap();
let args = InstallHooksArgs {
agent: AgentChoice::Omp,
capture_assistant: false,
hooks_dir: None,
server_url: Some("http://127.0.0.1:49374".into()),
auth_token: None,
@@ -4531,6 +4589,7 @@ model = "gpt-5"
let path = tmp.path().join("extensions").join("ai-memory.ts");
let args = InstallHooksArgs {
agent: AgentChoice::Pi,
capture_assistant: false,
hooks_dir: None,
server_url: Some("http://127.0.0.1:49374".into()),
auth_token: None,
@@ -5303,6 +5362,7 @@ command = "AI_MEMORY_HOOK_URL=http://old:1 /old/ai-memory/hooks/kimi-code/sessio
config_tmp.path(),
&InstallHooksArgs {
agent: AgentChoice::Devin,
capture_assistant: false,
hooks_dir: Some(hooks_tmp.path().to_path_buf()),
server_url: Some("http://127.0.0.1:49374".to_string()),
auth_token: None,
@@ -5365,6 +5425,7 @@ command = "AI_MEMORY_HOOK_URL=http://old:1 /old/ai-memory/hooks/kimi-code/sessio
config_tmp.path(),
&InstallHooksArgs {
agent: AgentChoice::Devin,
capture_assistant: false,
hooks_dir: Some(hooks_tmp.path().to_path_buf()),
server_url: Some("http://127.0.0.1:49374".to_string()),
auth_token: None,
@@ -5416,6 +5477,7 @@ command = "AI_MEMORY_HOOK_URL=http://old:1 /old/ai-memory/hooks/kimi-code/sessio
let args_v1 = InstallHooksArgs {
agent: AgentChoice::Devin,
capture_assistant: false,
hooks_dir: Some(hooks_tmp.path().to_path_buf()),
server_url: Some("http://127.0.0.1:49374".to_string()),
auth_token: None,
@@ -5460,6 +5522,7 @@ command = "AI_MEMORY_HOOK_URL=http://old:1 /old/ai-memory/hooks/kimi-code/sessio
let args_config = InstallHooksArgs {
agent: AgentChoice::Devin,
capture_assistant: false,
hooks_dir: Some(hooks_tmp.path().to_path_buf()),
server_url: Some("http://127.0.0.1:49374".to_string()),
auth_token: None,
@@ -5529,6 +5592,7 @@ command = "AI_MEMORY_HOOK_URL=http://old:1 /old/ai-memory/hooks/kimi-code/sessio
config_tmp.path(),
&InstallHooksArgs {
agent: AgentChoice::Devin,
capture_assistant: false,
hooks_dir: Some(hooks_tmp.path().to_path_buf()),
server_url: Some("http://127.0.0.1:49374".to_string()),
auth_token: None,
@@ -261,6 +261,7 @@ pub(crate) fn build_claude_code_payload_with_data_dir(
auth_token: Option<&str>,
data_dir: Option<&Path>,
project_strategy: Option<&str>,
capture_assistant: bool,
) -> serde_json::Value {
build_hook_payload_for_platform(
&CLAUDE_CODE_EVENTS,
@@ -274,7 +275,8 @@ pub(crate) fn build_claude_code_payload_with_data_dir(
data_dir,
project_strategy,
)
.allow_claude_windows_exec(),
.allow_claude_windows_exec()
.with_capture_assistant(capture_assistant),
)
}
@@ -776,6 +778,9 @@ struct HookCommandContext<'a> {
/// snippets keep command-string script fallback even when the platform env
/// is overridden to `windows-native`.
claude_windows_exec_allowed: bool,
/// Bake `--capture-assistant` onto the native `stop` command only (#196).
/// Set exclusively by `install-hooks --agent claude-code --capture-assistant`.
capture_assistant: bool,
}
impl<'a> HookCommandContext<'a> {
@@ -791,6 +796,7 @@ impl<'a> HookCommandContext<'a> {
data_dir,
project_strategy,
claude_windows_exec_allowed: false,
capture_assistant: false,
}
}
@@ -798,6 +804,22 @@ impl<'a> HookCommandContext<'a> {
self.claude_windows_exec_allowed = true;
self
}
const fn with_capture_assistant(mut self, on: bool) -> Self {
self.capture_assistant = on;
self
}
}
/// Bake `--capture-assistant` onto the native `stop` command only (#196), so the
/// other events stay byte-identical. Empty for every other event, and for a
/// context that did not opt in.
fn native_capture_assistant_arg(context: HookCommandContext<'_>, event: &str) -> &'static str {
if context.capture_assistant && event == "stop" {
" --capture-assistant"
} else {
""
}
}
impl HookCommandPlatform {
@@ -1067,6 +1089,7 @@ fn hook_command(
context.project_strategy,
NativeQuote::Windows,
));
cmd.push_str(native_capture_assistant_arg(context, event));
cmd
}
HookCommandPlatform::PosixNative => {
@@ -1095,6 +1118,7 @@ fn hook_command(
context.project_strategy,
NativeQuote::Posix,
));
cmd.push_str(native_capture_assistant_arg(context, event));
cmd
}
}
@@ -1151,6 +1175,11 @@ fn windows_native_exec_spec_with_exe(
args.push("--project-strategy".to_string());
args.push(strategy.to_string());
}
// Native-exec is the PRIMARY Windows Claude Code path, so the flag must be
// baked here too — not only in the string form above (#196).
if context.capture_assistant && event == "stop" {
args.push("--capture-assistant".to_string());
}
HookHandlerSpec::Exec {
command: plain_windows_path_arg(exe),
args,
@@ -1880,6 +1909,65 @@ check(activeKeep.disposition === "keep" && activeKeep.protocol?.version === 1 &&
assert!(args.iter().all(|arg| !arg.contains(r"\\?\")));
}
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"),
Path::new(script),
"http://h:49374",
None,
HookCommandContext::new(
HookCommandPlatform::WindowsNative,
"claude-code",
None,
None,
)
.with_capture_assistant(capture_assistant),
);
let HookHandlerSpec::Exec { args, .. } = spec else {
panic!("expected exec spec")
};
args
}
#[test]
fn capture_assistant_exec_flag_only_on_stop_and_only_when_opted_in() {
// Bare filenames so `file_stem()` yields the event on any host (a
// backslash path is a single component on Unix).
// Opted in: only the `stop` command carries the flag (#196, D1).
assert!(exec_args("stop.sh", true).contains(&"--capture-assistant".to_string()));
assert!(
!exec_args("session-start.sh", true).contains(&"--capture-assistant".to_string()),
"flag must not leak onto non-stop commands"
);
// Not opted in: even the stop command stays byte-identical to before.
assert!(
!exec_args("stop.sh", false).contains(&"--capture-assistant".to_string()),
"flag must be absent without opt-in"
);
}
#[test]
fn capture_assistant_string_form_only_on_stop() {
// POSIX + Windows string forms: the flag rides only the stop command.
for platform in [
HookCommandPlatform::PosixNative,
HookCommandPlatform::WindowsNative,
] {
let ctx = HookCommandContext::new(platform, "claude-code", None, None)
.with_capture_assistant(true);
let stop = hook_command(Path::new("stop.sh"), "http://h", None, ctx);
let start = hook_command(Path::new("session-start.sh"), "http://h", None, ctx);
assert!(
stop.contains("--capture-assistant"),
"{platform:?} stop missing flag: {stop}"
);
assert!(
!start.contains("--capture-assistant"),
"{platform:?} session-start must not carry flag: {start}"
);
}
}
#[test]
fn windows_native_exec_form_is_claude_only_and_guarded() {
fn handler_for(
@@ -396,6 +396,7 @@ pub async fn run(config: &Config, args: ServeArgs) -> Result<()> {
DEFAULT_HOOK_INGEST_MAX_IN_FLIGHT,
)),
consolidate_on_session_end: config.consolidate_on_session_end,
capture_assistant_enabled: config.capture_assistant,
subagent_sessions: std::sync::Arc::new(tokio::sync::Mutex::new(
ai_memory_hooks::SubagentSessionSet::default(),
)),
+9
View File
@@ -96,6 +96,14 @@ pub struct Config {
/// and via manual `memory_consolidate`. Set with
/// `AI_MEMORY_CONSOLIDATE_ON_SESSION_END=true`.
pub consolidate_on_session_end: bool,
/// Server-side opt-in for assistant/Stop capture (#196). When true, the
/// server honors a client's sanitized `_ai_memory_assistant` marker on a
/// `Stop` event and persists the excerpt as the Stop body. Off by default;
/// when off the marker is stripped and the Stop stays empty. The client half
/// of the double opt-in is baked separately by
/// `install-hooks --capture-assistant`. Set with
/// `AI_MEMORY_CAPTURE_ASSISTANT=true`.
pub capture_assistant: bool,
/// Optional embedding provider (`openai`, `voyage`, `google` / `gemini`).
pub embedding_provider: Option<String>,
/// Optional embedding model override.
@@ -361,6 +369,7 @@ impl Default for Config {
llm_base_url: None,
llm_compat_strict: false,
consolidate_on_session_end: false,
capture_assistant: false,
embedding_provider: None,
embedding_model: None,
embedding_dim: None,
@@ -25,6 +25,16 @@ log_level = "info"
# hook_rate_per_sec = 0.0
# hook_rate_burst = 0.0
# Server-side half of the assistant/Stop capture opt-in (#196). When true, the
# server persists a Claude Code Stop event's final assistant turn (sanitized and
# capped at 2 KB) as the Stop body. OFF by default: assistant text is privacy-
# sensitive — it can quote code, secrets, or content from paths ai-memory never
# sees, and it flows into consolidation/reviewer prompts (and out to a cloud LLM
# provider if one is configured). This is a DOUBLE opt-in: the client must also
# be installed with `ai-memory install-hooks --agent claude-code
# --capture-assistant`. Env equivalent: `AI_MEMORY_CAPTURE_ASSISTANT=true`.
# capture_assistant = false
# M8 retention-sweep parameters.
#
# Formula:
+84 -9
View File
@@ -13,18 +13,26 @@ fn run_hook(data_dir: &Path, payload: &[u8]) -> Output {
}
fn run_hook_event(data_dir: &Path, event: &str, payload: &[u8]) -> Output {
run_hook_full(data_dir, event, payload, false)
}
fn run_hook_full(data_dir: &Path, event: &str, payload: &[u8], capture_assistant: bool) -> Output {
let mut args = vec![
"hook".to_string(),
"--event".to_string(),
event.to_string(),
"--agent".to_string(),
"claude-code".to_string(),
"--server-url".to_string(),
"http://127.0.0.1:1".to_string(),
];
if capture_assistant {
args.push("--capture-assistant".to_string());
}
let mut child = Command::new(bin())
.args(["--data-dir"])
.arg(data_dir)
.args([
"hook",
"--event",
event,
"--agent",
"claude-code",
"--server-url",
"http://127.0.0.1:1",
])
.args(&args)
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
@@ -58,6 +66,17 @@ fn spooled_body(data_dir: &Path) -> String {
entry["body"].as_str().expect("spooled body").to_owned()
}
fn spooled_entry(data_dir: &Path) -> serde_json::Value {
// Use the filtered helper: boundary events spawn a detached drainer, so the
// spool directory can also hold `.drain.lock` / `.json.tmp` while that child
// is alive. Counting raw directory entries races with it (Windows loses most
// often).
let entries = spool_entries(data_dir);
assert_eq!(entries.len(), 1);
serde_json::from_slice(&std::fs::read(entries[0].path()).expect("read spool entry"))
.expect("parse spool entry")
}
#[test]
fn native_hook_accepts_plain_and_bom_prefixed_json() {
let payload = br#"{"session_id":"windows-test","cwd":"C:\\dev\\project","tool_name":"Read","tool_input":{"file_path":"README.md"}}"#;
@@ -158,3 +177,59 @@ fn spool_files_never_leak_the_assistant_field_on_disk() {
"raw assistant field key leaked into the spool file bytes"
);
}
#[test]
fn opted_in_stop_splices_protocol_and_capture_flag() {
// With --capture-assistant, a Stop event spools the sanitized protocol
// marker (NOT the raw field) and carries capture_assistant=1 on the URL so
// the server can gate on it (#196).
let tmp = tempfile::tempdir().expect("tempdir");
let payload = br#"{"session_id":"opt-in","last_assistant_message":"the fix is here"}"#;
let output = run_hook_full(tmp.path(), "stop", payload, true);
assert!(output.status.success());
assert_eq!(output.stdout, b"{}\n");
let entry = spooled_entry(tmp.path());
let body = entry["body"].as_str().expect("spooled body");
let url = entry["url"].as_str().expect("spooled url");
assert!(
!body.contains("last_assistant_message"),
"raw field survived: {body}"
);
assert!(
body.contains("_ai_memory_assistant"),
"protocol marker missing: {body}"
);
assert!(
body.contains("the fix is here"),
"excerpt missing from protocol: {body}"
);
assert!(
url.contains("capture_assistant=1"),
"capture flag missing from url: {url}"
);
}
#[test]
fn opted_in_non_stop_event_is_inert() {
// The flag is a no-op on non-Stop events: no protocol, no capture flag, and
// (absent any assistant field) the body is byte-exact.
let tmp = tempfile::tempdir().expect("tempdir");
let payload = br#"{"session_id":"opt-in","prompt":"hello"}"#;
let output = run_hook_full(tmp.path(), "user-prompt-submit", payload, true);
assert!(output.status.success());
let entry = spooled_entry(tmp.path());
let url = entry["url"].as_str().expect("spooled url");
assert!(
!url.contains("capture_assistant=1"),
"capture flag leaked onto a non-stop event: {url}"
);
assert_eq!(
entry["body"].as_str().expect("spooled body").as_bytes(),
payload,
"unrelated event body must stay byte-exact"
);
}
+421 -8
View File
@@ -1,20 +1,204 @@
//! Raw assistant-message stripping (issue #196).
//! Assistant-message capture plumbing (issue #196).
//!
//! Some agent harnesses attach the assistant's final turn to their `Stop`
//! lifecycle event — Claude Code sends it as a top-level `last_assistant_message`
//! string. The text can contain code, secrets, or content from paths ai-memory
//! never otherwise sees, so the raw field is removed before spooling, transport,
//! tracing, or storage. Optional capture remains disabled.
//! string. That text is high-value for recall but privacy-sensitive: it can
//! quote code, secrets, or content from paths ai-memory never sees. Capturing
//! it is therefore an explicit, double opt-in feature (server config +
//! `install-hooks --capture-assistant`).
//!
//! This module owns the single source of truth for WHICH raw field carries the
//! assistant message per agent/event, the unconditional strip that keeps that
//! raw field off the local spool, the wire, tracing, and storage, and the
//! opt-in path that re-introduces a sanitized, capped excerpt: the client-side
//! [`transform_for_client`] and the server-side [`apply_assistant_backstop`].
/// Raw top-level field names known to carry assistant messages.
use ai_memory_core::{AgentKind, Sanitizer};
use serde::{Deserialize, Serialize};
use crate::payload::{HookEnvelope, HookEvent, truncate_utf8_bytes};
/// Synthetic body key carrying the opt-in, sanitized assistant excerpt from the
/// client to the server. Distinct from the raw `last_assistant_message` field,
/// which is always stripped: this one is the deliberate, capped protocol.
pub const ASSISTANT_MARKER_KEY: &str = "_ai_memory_assistant";
/// Protocol version for the opt-in `_ai_memory_assistant` body marker the client
/// attaches when capture is enabled. Bumping it invalidates markers a stale
/// server would otherwise trust.
pub const ASSISTANT_PROTOCOL_VERSION: u8 = 1;
/// Hard ceiling on the raw assistant-message string the client reads before
/// sanitizing/truncating. Oversized input is treated as absent.
pub const ASSISTANT_MESSAGE_MAX_INPUT_BYTES: usize = 64 * 1024;
/// Byte cap on the sanitized excerpt the opt-in path persists. Kept equal to
/// `truncate_excerpt`'s existing 2 KB excerpt contract so Stop bodies do not
/// become a second, larger excerpt norm.
pub const ASSISTANT_EXCERPT_MAX_BYTES: usize = 2_000;
/// Every raw top-level field name known to carry an assistant message across
/// supported agents. The unconditional strip removes each of these; the closed
/// per-agent table below decides which is a *candidate* for opt-in capture. The
/// union is kept tiny on purpose — one entry per distinct wire spelling.
const ASSISTANT_MESSAGE_FIELDS: &[&str] = &["last_assistant_message"];
/// The raw field that carries the assistant's final message for `(agent, event)`,
/// or `None` when the pair has no verified assistant-message field.
///
/// Closed table: only `ClaudeCode + Stop` is supported today. Extend
/// deliberately — a new entry opts an agent/event into capture and MUST have its
/// field name present in [`ASSISTANT_MESSAGE_FIELDS`] so the strip covers it
/// (enforced by `closed_table_fields_are_all_stripped`).
#[must_use]
pub fn assistant_message_field(agent: AgentKind, event: HookEvent) -> Option<&'static str> {
match (agent, event) {
(AgentKind::ClaudeCode, HookEvent::Stop) => Some("last_assistant_message"),
_ => None,
}
}
/// The opt-in assistant-capture wire protocol carried in the body under
/// [`ASSISTANT_MARKER_KEY`]. `deny_unknown_fields` + the explicit version gate
/// means a malformed or future-versioned marker is rejected (dropped, never
/// persisted) rather than partially trusted.
#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(deny_unknown_fields)]
pub struct AssistantCaptureProtocol {
/// Wire version. Must equal [`ASSISTANT_PROTOCOL_VERSION`].
pub version: u8,
/// Client-sanitized, byte-capped assistant excerpt. The server re-scrubs it
/// with its configured `Sanitizer` at the persistence boundary.
pub excerpt: String,
}
impl AssistantCaptureProtocol {
/// Parse a marker value, accepting only the current protocol version. Any
/// unknown field, wrong type, or version mismatch yields `None`.
#[must_use]
fn parse(value: &serde_json::Value) -> Option<Self> {
let parsed: Self = serde_json::from_value(value.clone()).ok()?;
(parsed.version == ASSISTANT_PROTOCOL_VERSION).then_some(parsed)
}
}
/// Outcome of the client-side assistant transform.
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
pub struct ClientAssistantTransform {
/// The raw JSON was mutated (raw field removed and/or protocol spliced in),
/// so the caller must reserialize the spool/wire payload.
pub changed: bool,
/// A valid protocol was spliced in, so the caller appends
/// `&capture_assistant=1` to the event URL.
pub captured: bool,
}
/// Client-side transform for an install that opted into assistant capture.
///
/// Reads the candidate assistant message for `(agent, event)`, unconditionally
/// strips the raw field, and — when the value is a non-empty, in-bounds string —
/// sanitizes it with the built-in `Sanitizer`, truncates it to
/// [`ASSISTANT_EXCERPT_MAX_BYTES`] on a UTF-8 boundary, and splices the versioned
/// [`AssistantCaptureProtocol`] into the body under [`ASSISTANT_MARKER_KEY`].
///
/// Scrub happens BEFORE truncation (a secret straddling the cap must be redacted
/// before it can be cut). Non-string, empty, or oversized values yield no
/// protocol — the raw field is still stripped, so the event degrades to an empty
/// Stop rather than leaking anything.
pub fn transform_for_client(
raw: &mut serde_json::Value,
agent: AgentKind,
event: HookEvent,
) -> ClientAssistantTransform {
// Read the eligible value BEFORE stripping removes the field.
let candidate = assistant_message_field(agent, event).and_then(|field| {
raw.as_object()
.and_then(|object| object.get(field))
.and_then(serde_json::Value::as_str)
.filter(|value| !value.is_empty() && value.len() <= ASSISTANT_MESSAGE_MAX_INPUT_BYTES)
.map(str::to_string)
});
let mut changed = strip_assistant_message_raw(raw);
let Some(value) = candidate else {
return ClientAssistantTransform {
changed,
captured: false,
};
};
let scrubbed = Sanitizer::builtin().scrub(&value);
let excerpt = truncate_utf8_bytes(&scrubbed, ASSISTANT_EXCERPT_MAX_BYTES);
if excerpt.is_empty() {
return ClientAssistantTransform {
changed,
captured: false,
};
}
let protocol = AssistantCaptureProtocol {
version: ASSISTANT_PROTOCOL_VERSION,
excerpt,
};
let mut captured = false;
if let Some(object) = raw.as_object_mut()
&& let Ok(marker) = serde_json::to_value(&protocol)
{
object.insert(ASSISTANT_MARKER_KEY.to_string(), marker);
changed = true;
captured = true;
}
ClientAssistantTransform { changed, captured }
}
/// Server-side backstop for the opt-in assistant excerpt (#196).
///
/// Always consumes the [`ASSISTANT_MARKER_KEY`] from `env.raw` so it can never
/// persist. Populates `env.body_excerpt` with the excerpt ONLY when every gate
/// holds: the server enabled capture, the client requested it (`capture_assistant`
/// query flag), the agent/event is a supported candidate, and the marker parses
/// as a current-version protocol with a non-empty excerpt. Any failure leaves
/// `body_excerpt` as `None` — an empty Stop — so a forged or stale marker cannot
/// inject content. Infallible: never returns an error, so a capture decision can
/// never turn into a batch fail-fast.
pub fn apply_assistant_backstop(env: &mut HookEnvelope, server_enabled: bool) {
// Consume the marker unconditionally: it must never survive into the stored
// raw, whether or not the gates below accept it.
let marker = env
.raw
.as_object_mut()
.and_then(|object| object.remove(ASSISTANT_MARKER_KEY));
let eligible = server_enabled
&& env.capture_assistant_requested
&& assistant_message_field(env.agent, env.event).is_some();
if !eligible {
return;
}
if let Some(marker) = marker
&& let Some(protocol) = AssistantCaptureProtocol::parse(&marker)
&& !protocol.excerpt.is_empty()
{
// Re-enforce the excerpt cap at the persistence boundary — never trust
// the client's length. Symmetric to the server re-scrub in `process()`:
// a forged or buggy client that satisfies the gates cannot inject an
// oversized Stop body (the request body limit alone is 10 MiB). Almost
// always a no-op, since a well-behaved client already truncated.
env.body_excerpt = Some(truncate_utf8_bytes(
&protocol.excerpt,
ASSISTANT_EXCERPT_MAX_BYTES,
));
}
}
/// Unconditionally remove every known assistant-message field from a raw hook
/// payload's top-level object, returning whether anything was removed.
///
/// This defense is applied on both sides of the wire (client pre-spool and
/// server pre-envelope) and for every agent/event. Only top-level keys are
/// inspected, matching where supported harnesses place the field.
/// This is a defense applied on BOTH sides of the wire (client pre-spool and
/// server pre-envelope) and for EVERY agent/event, not just the supported pair:
/// a raw assistant-message field must never reach the spool, the wire, tracing,
/// or storage unless the explicit opt-in path re-introduces it as a sanitized,
/// capped excerpt. Only top-level keys are inspected — the same scope
/// as `body_is_subagent`, and where every supported harness places the field.
pub fn strip_assistant_message_raw(raw: &mut serde_json::Value) -> bool {
let Some(object) = raw.as_object_mut() else {
return false;
@@ -32,6 +216,66 @@ pub fn strip_assistant_message_raw(raw: &mut serde_json::Value) -> bool {
mod tests {
use super::*;
/// Only `ClaudeCode + Stop` is a capture candidate; every other agent/event
/// pair across the full agent surface must return `None`.
#[test]
fn only_claude_stop_is_a_capture_candidate() {
let events = [
HookEvent::SessionStart,
HookEvent::UserPrompt,
HookEvent::PreToolUse,
HookEvent::PostToolUse,
HookEvent::PreCompact,
HookEvent::PostCompaction,
HookEvent::Notification,
HookEvent::Stop,
HookEvent::SessionEnd,
HookEvent::SubagentStart,
HookEvent::SubagentStop,
HookEvent::Other,
];
for agent in AgentKind::ALL {
for event in events {
let expected = agent == AgentKind::ClaudeCode && event == HookEvent::Stop;
assert_eq!(
assistant_message_field(agent, event).is_some(),
expected,
"agent={agent:?} event={event:?}"
);
}
}
}
/// The strip must cover every field the closed table can name, or an opted-in
/// agent/event could carry a raw field the strip misses.
#[test]
fn closed_table_fields_are_all_stripped() {
let events = [
HookEvent::SessionStart,
HookEvent::UserPrompt,
HookEvent::PreToolUse,
HookEvent::PostToolUse,
HookEvent::PreCompact,
HookEvent::PostCompaction,
HookEvent::Notification,
HookEvent::Stop,
HookEvent::SessionEnd,
HookEvent::SubagentStart,
HookEvent::SubagentStop,
HookEvent::Other,
];
for agent in AgentKind::ALL {
for event in events {
if let Some(field) = assistant_message_field(agent, event) {
assert!(
ASSISTANT_MESSAGE_FIELDS.contains(&field),
"closed-table field {field:?} is not in the strip set"
);
}
}
}
}
/// The strip is unconditional: it removes the raw field regardless of agent
/// and reports the removal, so a client that never verified the agent still
/// cannot leak it.
@@ -56,4 +300,173 @@ mod tests {
assert!(!strip_assistant_message_raw(&mut array));
assert_eq!(array, serde_json::json!(["last_assistant_message"]));
}
fn stop_env(raw: serde_json::Value, requested: bool) -> HookEnvelope {
HookEnvelope::from_query_and_body(
crate::payload::HookQuery {
event: "stop".into(),
agent: Some("claude-code".into()),
capture_assistant: requested.then(|| "1".to_string()),
..Default::default()
},
raw,
)
}
#[test]
fn client_transform_splices_sanitized_capped_excerpt() {
let mut raw = serde_json::json!({
"session_id": "s1",
"last_assistant_message": "fixed the bug"
});
let out = transform_for_client(&mut raw, AgentKind::ClaudeCode, HookEvent::Stop);
assert!(out.changed && out.captured);
assert!(raw.get("last_assistant_message").is_none());
let marker = raw.get(ASSISTANT_MARKER_KEY).expect("protocol spliced");
let protocol = AssistantCaptureProtocol::parse(marker).expect("valid v1 protocol");
assert_eq!(protocol.version, ASSISTANT_PROTOCOL_VERSION);
assert_eq!(protocol.excerpt, "fixed the bug");
}
#[test]
fn client_transform_scrubs_before_truncating() {
// A built-in secret pattern must be redacted in the excerpt.
let secret = "AKIA".to_string() + &"A".repeat(16); // AWS access key id shape
let mut raw = serde_json::json!({ "last_assistant_message": format!("key {secret}") });
let out = transform_for_client(&mut raw, AgentKind::ClaudeCode, HookEvent::Stop);
assert!(out.captured);
let excerpt = raw[ASSISTANT_MARKER_KEY]["excerpt"].as_str().unwrap();
assert!(!excerpt.contains(&secret), "secret survived: {excerpt}");
assert!(excerpt.contains("[REDACTED]"), "not redacted: {excerpt}");
}
#[test]
fn client_transform_truncates_multibyte_within_cap() {
let big = "é".repeat(ASSISTANT_EXCERPT_MAX_BYTES); // 2 bytes each → 2x the cap
let mut raw = serde_json::json!({ "last_assistant_message": big });
let out = transform_for_client(&mut raw, AgentKind::ClaudeCode, HookEvent::Stop);
assert!(out.captured);
let excerpt = raw[ASSISTANT_MARKER_KEY]["excerpt"].as_str().unwrap();
assert!(excerpt.len() <= ASSISTANT_EXCERPT_MAX_BYTES, "over cap");
assert!(excerpt.ends_with('…'), "truncation marker missing");
// Valid UTF-8 (no split codepoint): re-parsing as str succeeds by construction.
assert!(std::str::from_utf8(excerpt.as_bytes()).is_ok());
}
#[test]
fn client_transform_strips_but_omits_protocol_for_empty_or_oversized() {
for value in ["", &"x".repeat(ASSISTANT_MESSAGE_MAX_INPUT_BYTES + 1)] {
let mut raw = serde_json::json!({ "last_assistant_message": value });
let out = transform_for_client(&mut raw, AgentKind::ClaudeCode, HookEvent::Stop);
assert!(!out.captured, "value {:?} must not capture", value.len());
assert!(
raw.get("last_assistant_message").is_none(),
"raw not stripped"
);
assert!(
raw.get(ASSISTANT_MARKER_KEY).is_none(),
"protocol spliced anyway"
);
}
}
#[test]
fn client_transform_ignores_non_candidate_agent_event() {
// Non-Claude agent: raw field still stripped defensively, no protocol.
let mut raw = serde_json::json!({ "last_assistant_message": "hi" });
let out = transform_for_client(&mut raw, AgentKind::Codex, HookEvent::Stop);
assert!(!out.captured);
assert!(raw.get("last_assistant_message").is_none());
assert!(raw.get(ASSISTANT_MARKER_KEY).is_none());
}
#[test]
fn backstop_populates_body_when_all_gates_pass() {
let mut raw = serde_json::json!({ "last_assistant_message": "done" });
transform_for_client(&mut raw, AgentKind::ClaudeCode, HookEvent::Stop);
let mut env = stop_env(raw, true);
apply_assistant_backstop(&mut env, true);
assert_eq!(env.body_excerpt.as_deref(), Some("done"));
assert!(
env.raw.get(ASSISTANT_MARKER_KEY).is_none(),
"marker persisted"
);
}
#[test]
fn backstop_drops_when_server_off_or_flag_absent() {
let make = || {
let mut raw = serde_json::json!({ "last_assistant_message": "done" });
transform_for_client(&mut raw, AgentKind::ClaudeCode, HookEvent::Stop);
raw
};
// Server disabled.
let mut env = stop_env(make(), true);
apply_assistant_backstop(&mut env, false);
assert!(env.body_excerpt.is_none());
assert!(
env.raw.get(ASSISTANT_MARKER_KEY).is_none(),
"marker still consumed"
);
// Query flag absent (client did not request).
let mut env = stop_env(make(), false);
apply_assistant_backstop(&mut env, true);
assert!(env.body_excerpt.is_none());
}
#[test]
fn backstop_caps_oversized_excerpt_server_side() {
// A forged/buggy client can satisfy the gates with an excerpt far above
// the 2 KB cap (the client-side truncation is not trustworthy). The
// server must re-enforce the cap at the persistence boundary.
let huge = "a".repeat(ASSISTANT_EXCERPT_MAX_BYTES * 500);
let raw = serde_json::json!({
ASSISTANT_MARKER_KEY: { "version": 1, "excerpt": huge }
});
let mut env = stop_env(raw, true);
apply_assistant_backstop(&mut env, true);
let body = env.body_excerpt.expect("gates pass, body set");
assert!(
body.len() <= ASSISTANT_EXCERPT_MAX_BYTES,
"server did not cap oversized excerpt: {} bytes",
body.len()
);
}
#[test]
fn backstop_rejects_forged_empty_and_bad_version() {
for marker in [
serde_json::json!({ "version": 1, "excerpt": "" }),
serde_json::json!({ "version": 2, "excerpt": "future" }),
serde_json::json!({ "version": 1, "excerpt": "x", "extra": true }),
serde_json::json!({ "excerpt": "no version" }),
] {
let raw = serde_json::json!({ ASSISTANT_MARKER_KEY: marker });
let mut env = stop_env(raw, true);
apply_assistant_backstop(&mut env, true);
assert!(
env.body_excerpt.is_none(),
"forged marker accepted: {:?}",
env.body_excerpt
);
}
}
#[test]
fn backstop_ignores_non_stop_event() {
let raw = serde_json::json!({
ASSISTANT_MARKER_KEY: { "version": 1, "excerpt": "x" }
});
let mut env = HookEnvelope::from_query_and_body(
crate::payload::HookQuery {
event: "user-prompt-submit".into(),
agent: Some("claude-code".into()),
capture_assistant: Some("1".into()),
..Default::default()
},
raw,
);
apply_assistant_backstop(&mut env, true);
assert!(env.body_excerpt.is_none());
}
}
+6 -1
View File
@@ -32,7 +32,12 @@ pub mod workstream;
// Re-export the sanitizer types from core so callers that grew up
// pointing at this crate's `sanitize` module keep working.
pub use ai_memory_core::{SanitizeConfig, Sanitized, Sanitizer};
pub use assistant_capture::strip_assistant_message_raw;
// Client-side symbols used by the CLI crate; the server-side `apply_assistant_backstop`
// and the protocol/table internals stay crate-private (router reaches them via
// `crate::assistant_capture`).
pub use assistant_capture::{
ClientAssistantTransform, strip_assistant_message_raw, transform_for_client,
};
pub use capture_policy::{
CaptureConfig, CaptureDecision, CaptureDisposition, CapturePolicy, CaptureProtocol,
CaptureSource, ExtractionState, PolicyState, ToolFamily,
+41 -18
View File
@@ -58,6 +58,12 @@ pub struct HookQuery {
/// Invocation-scoped `ai-memory run` lease. Absent for every direct
/// harness launch, preserving legacy capture and handoff behavior.
pub managed_run: Option<String>,
/// Client-side opt-in for assistant/Stop capture, baked onto the native
/// `stop` hook command by `install-hooks --capture-assistant`. A truthy
/// value tells the server the client deliberately attached a sanitized
/// `_ai_memory_assistant` excerpt; the server still gates on its own
/// `capture_assistant` config before persisting it (#196).
pub capture_assistant: Option<String>,
}
/// Coalesced view of an incoming hook event after light parsing of the
@@ -99,6 +105,11 @@ pub struct HookEnvelope {
pub extension: Option<String>,
/// Optional source event name from the extension vocabulary.
pub source_event: Option<String>,
/// Whether the client requested assistant/Stop capture for this event
/// (the `capture_assistant` query flag baked onto the native `stop`
/// command). The server still gates on its own `capture_assistant` config
/// before honoring it (#196).
pub capture_assistant_requested: bool,
/// Optional title hint extracted from the body.
pub title_hint: Option<String>,
/// Optional body excerpt extracted from the agent's raw payload.
@@ -125,6 +136,10 @@ impl std::fmt::Debug for HookEnvelope {
&self.recall_default_global_requested,
)
.field("managed_run", &self.managed_run)
.field(
"capture_assistant_requested",
&self.capture_assistant_requested,
)
.field("extension", &self.extension)
.field("source_event", &self.source_event)
.field(
@@ -358,6 +373,7 @@ impl HookEnvelope {
let drop_subagent_requested = query_flag_truthy(query.drop_subagent.as_deref());
let recall_default_global_requested = query_flag_truthy(query.default_global.as_deref());
let managed_run = query.managed_run.filter(|value| !value.trim().is_empty());
let capture_assistant_requested = query_flag_truthy(query.capture_assistant.as_deref());
let extension = normalize_extension_name(query.extension.as_deref());
let source_event = extension.as_ref().and_then(|_| {
let raw_source = query
@@ -415,6 +431,7 @@ impl HookEnvelope {
drop_subagent_requested,
recall_default_global_requested,
managed_run,
capture_assistant_requested,
extension,
source_event,
title_hint,
@@ -734,25 +751,31 @@ fn truncate_for_title(s: &str) -> String {
}
fn truncate_excerpt(s: &str) -> String {
const MAX: usize = 2_000;
if s.len() <= MAX {
s.to_string()
} else {
// Reserve the ellipsis within the byte cap, not beyond it.
let limit = MAX - '…'.len_utf8();
let mut buf = String::with_capacity(MAX);
let mut end = 0;
for (idx, ch) in s.char_indices() {
let next = idx + ch.len_utf8();
if next > limit {
break;
}
end = next;
}
buf.push_str(&s[..end]);
buf.push('…');
buf
truncate_utf8_bytes(s, 2_000)
}
/// Truncate `s` to at most `max` bytes on a UTF-8 char boundary, reserving the
/// ellipsis within the cap (never beyond it). Shared by the tool-excerpt cap
/// and the opt-in assistant excerpt cap (#196) so there is one UTF-8-safe
/// truncation, two named caps.
pub(crate) fn truncate_utf8_bytes(s: &str, max: usize) -> String {
if s.len() <= max {
return s.to_string();
}
// Reserve the ellipsis within the byte cap, not beyond it.
let limit = max.saturating_sub('…'.len_utf8());
let mut buf = String::with_capacity(max);
let mut end = 0;
for (idx, ch) in s.char_indices() {
let next = idx + ch.len_utf8();
if next > limit {
break;
}
end = next;
}
buf.push_str(&s[..end]);
buf.push('…');
buf
}
fn normalize_extension_name(value: Option<&str>) -> Option<String> {
+179 -23
View File
@@ -409,6 +409,12 @@ pub struct HookState {
/// session close stays cheap; the LLM checkpoint otherwise happens on
/// PreCompact and via manual `memory_consolidate`.
pub consolidate_on_session_end: bool,
/// Opt-in (`AI_MEMORY_CAPTURE_ASSISTANT`): when true, the server honors the
/// client's `_ai_memory_assistant` protocol on a `Stop` event and persists
/// the sanitized excerpt as the Stop body. Off by default; when off the
/// marker is stripped and the Stop stays empty. Double opt-in: the client
/// must also have been installed with `--capture-assistant` (#196).
pub capture_assistant_enabled: bool,
/// Scoped session keys known to be subagents (seeded by `SubagentStart` / any
/// marker-bearing event). For a project that opted into
/// `drop_subagent_captures` (via its `.ai-memory.toml`, forwarded as the
@@ -443,7 +449,11 @@ async fn handle_hook(
// `Value` before it becomes a `HookEnvelope`, so the field can never reach
// `body_excerpt`, tracing, or the store — regardless of client version.
crate::assistant_capture::strip_assistant_message_raw(&mut body);
let env = HookEnvelope::from_query_and_body(query, body);
let mut env = HookEnvelope::from_query_and_body(query, body);
// Consume the opt-in `_ai_memory_assistant` marker and, when both opt-ins are
// on for a supported Stop, populate the Stop body with the sanitized excerpt.
// Any gate failure leaves an empty Stop with the same 202 "queued" response.
crate::assistant_capture::apply_assistant_backstop(&mut env, state.capture_assistant_enabled);
let Some(env) = inspect_capture_envelope(env) else {
return (StatusCode::ACCEPTED, "capture policy dropped");
};
@@ -605,7 +615,11 @@ async fn handle_hook_batch(
// per item before the envelope is built (#196).
crate::assistant_capture::strip_assistant_message_raw(&mut item.body);
let query = parse_hook_query(&item.url);
let env = HookEnvelope::from_query_and_body(query, item.body);
let mut env = HookEnvelope::from_query_and_body(query, item.body);
crate::assistant_capture::apply_assistant_backstop(
&mut env,
state.capture_assistant_enabled,
);
let Some(env) = inspect_capture_envelope(env) else {
// A protocol-directed drop is committed from the spool's point of
// view, but intentionally spends neither ingress capacity nor a
@@ -2247,7 +2261,7 @@ mod tests {
use std::sync::{Arc, Mutex};
use ai_memory_consolidate::{AutoImproveReviewConfig, run_auto_improve_review};
use ai_memory_core::Sanitizer;
use ai_memory_core::{SanitizeConfig, Sanitizer};
use ai_memory_llm::{ChatRequest, ChatResponse, LlmProvider, LlmResult};
use ai_memory_store::Store;
use ai_memory_wiki::Wiki;
@@ -2312,6 +2326,7 @@ mod tests {
project_cache: Arc::new(tokio::sync::Mutex::new(ProjectCacheStore::default())),
active_project: ActiveProject::new(),
consolidate_on_session_end: false,
capture_assistant_enabled: false,
subagent_sessions: Arc::new(tokio::sync::Mutex::new(SubagentSessionSet::default())),
ingest_rate: Arc::new(tokio::sync::Mutex::new(IngestRateLimiter::disabled())),
home_dir: None,
@@ -2970,6 +2985,98 @@ mod tests {
);
}
/// Build a Stop batch item as an opted-in client would: raw field stripped,
/// sanitized `_ai_memory_assistant` marker spliced in, `capture_assistant=1`
/// on the URL.
fn opted_in_stop_item(session_id: &str, message: &str) -> HookBatchItem {
let mut body = serde_json::json!({
"session_id": session_id,
"last_assistant_message": message,
});
let out = crate::assistant_capture::transform_for_client(
&mut body,
ai_memory_core::AgentKind::ClaudeCode,
HookEvent::Stop,
);
assert!(out.captured, "test fixture must produce a protocol");
HookBatchItem {
url: "http://h/hook?event=stop&agent=claude-code&capture_assistant=1".into(),
body,
}
}
fn stop_session_id(session_id: &str) -> ai_memory_core::SessionId {
resolve_session_id(&HookEnvelope::from_query_and_body(
HookQuery {
event: "stop".into(),
agent: Some("claude-code".into()),
session_id: Some(session_id.into()),
..Default::default()
},
serde_json::json!({ "session_id": session_id }),
))
.unwrap()
}
#[tokio::test]
async fn assistant_capture_round_trips_when_both_opt_ins_on() {
let tmp = TempDir::new().unwrap();
let mut state = make_state(&tmp).await;
state.capture_assistant_enabled = true;
let state = Arc::new(state);
let items = vec![opted_in_stop_item("cap-on", "the fix is in config.rs")];
let response = handle_hook_batch(State(state.clone()), None, Json(items))
.await
.into_response();
assert_eq!(response.status(), StatusCode::OK);
let observations = state
.reader
.observations_for_session(stop_session_id("cap-on"))
.await
.unwrap();
let stop = observations
.iter()
.find(|o| o.kind == ai_memory_core::ObservationKind::Stop)
.expect("Stop persisted");
assert_eq!(
stop.body, "the fix is in config.rs",
"excerpt must be persisted as the Stop body"
);
// The synthetic marker must not survive anywhere on disk.
assert!(!any_file_contains(tmp.path(), b"_ai_memory_assistant"));
}
#[tokio::test]
async fn assistant_capture_stays_empty_when_server_disabled() {
let tmp = TempDir::new().unwrap();
// make_state defaults capture_assistant_enabled = false.
let state = Arc::new(make_state(&tmp).await);
let items = vec![opted_in_stop_item("cap-off", "should not persist")];
let response = handle_hook_batch(State(state.clone()), None, Json(items))
.await
.into_response();
assert_eq!(response.status(), StatusCode::OK);
let observations = state
.reader
.observations_for_session(stop_session_id("cap-off"))
.await
.unwrap();
let stop = observations
.iter()
.find(|o| o.kind == ai_memory_core::ObservationKind::Stop)
.expect("Stop still persisted, just empty");
assert!(
stop.body.is_empty(),
"server-off Stop must be empty, got: {:?}",
stop.body
);
assert!(!any_file_contains(tmp.path(), b"should not persist"));
}
/// `pre-tool-use` query+agent for building an env to recompute a SessionId.
fn grok_tool_query() -> HookQuery {
HookQuery {
@@ -6485,10 +6592,17 @@ mod tests {
}
#[tokio::test]
async fn privacy_protocol_evidence_never_persists_protected_tool_sentinel() {
const SENTINEL: &str = "PHASE3_PROTECTED_PATH_AND_CONTENT_7f6c";
async fn privacy_protocol_and_assistant_capture_sentinels_never_reach_storage_or_reviewer() {
const TOOL_SENTINEL: &str = "PHASE3_PROTECTED_PATH_AND_CONTENT_7f6c";
const ASSISTANT_SENTINEL: &str = "ASSISTANT_PRIVATE_RESULT_9c2e";
let tmp = TempDir::new().unwrap();
let state = make_state(&tmp).await;
let mut state = make_state(&tmp).await;
state.sanitizer = Sanitizer::new(&SanitizeConfig {
extra_patterns: vec![ASSISTANT_SENTINEL.into()],
allowlist: Vec::new(),
})
.unwrap();
state.capture_assistant_enabled = true;
let session_id = "privacy-evidence";
for (event, body) in [
@@ -6527,7 +6641,7 @@ mod tests {
},
serde_json::json!({
"session_id": session_id, "tool_name": "Write",
"tool_input": { "file_path": SENTINEL }, "tool_response": SENTINEL,
"tool_input": { "file_path": TOOL_SENTINEL }, "tool_response": TOOL_SENTINEL,
"_ai_memory_capture": capture_protocol("drop", "inactive", "unknown", 99, "extracted"),
}),
);
@@ -6551,6 +6665,33 @@ mod tests {
);
process(&state, metadata, None).await.unwrap();
let mut assistant_body = serde_json::json!({
"session_id": session_id,
"cwd": "/repo",
"last_assistant_message": format!("completed safely: {ASSISTANT_SENTINEL}"),
});
let transformed = crate::assistant_capture::transform_for_client(
&mut assistant_body,
AgentKind::ClaudeCode,
HookEvent::Stop,
);
assert!(transformed.captured);
assert!(assistant_body.get("last_assistant_message").is_none());
let mut assistant = HookEnvelope::from_query_and_body(
HookQuery {
event: "stop".into(),
agent: Some("claude-code".into()),
capture_assistant: Some("true".into()),
..Default::default()
},
assistant_body,
);
crate::assistant_capture::apply_assistant_backstop(
&mut assistant,
state.capture_assistant_enabled,
);
process(&state, assistant, None).await.unwrap();
process(
&state,
HookEnvelope::from_query_and_body(
@@ -6581,41 +6722,54 @@ mod tests {
.unwrap()
.unwrap();
let observations = state.reader.observations_for_session(sid).await.unwrap();
assert!(observations.iter().all(|observation| {
observation.body.is_empty() || !observation.body.contains(SENTINEL)
}));
for sentinel in [TOOL_SENTINEL, ASSISTANT_SENTINEL] {
assert!(observations.iter().all(|observation| {
observation.body.is_empty() || !observation.body.contains(sentinel)
}));
}
assert!(observations.iter().any(|observation| {
observation.title == "file" && observation.body == "tool_family: file\noutcome: unknown"
}));
assert!(
state
.reader
.search_observations_for_project(workspace_id, project_id, SENTINEL.into(), 10)
.await
.unwrap()
.is_empty()
);
assert!(observations.iter().any(|observation| {
observation.kind == ObservationKind::Stop
&& observation.body == "completed safely: [REDACTED]"
}));
for sentinel in [TOOL_SENTINEL, ASSISTANT_SENTINEL] {
assert!(
state
.reader
.search_observations_for_project(workspace_id, project_id, sentinel.into(), 10,)
.await
.unwrap()
.is_empty()
);
}
let page = state
.reader
.page_body_by_ids(workspace_id, project_id, &format!("sessions/{sid}.md"))
.await
.unwrap()
.unwrap();
assert!(!page.body.contains(SENTINEL));
assert!(!page.body.contains(TOOL_SENTINEL));
assert!(!page.body.contains(ASSISTANT_SENTINEL));
let handoff = state
.reader
.latest_open_handoff(workspace_id, project_id, None)
.await
.unwrap()
.unwrap();
assert!(!handoff.summary.contains(SENTINEL));
assert!(!handoff.summary.contains(TOOL_SENTINEL));
assert!(!handoff.summary.contains(ASSISTANT_SENTINEL));
assert!(
!state
.wiki
.recent_checkpoints(20)
.unwrap()
.iter()
.any(|entry| entry.summary.contains(SENTINEL))
.any(|entry| {
entry.summary.contains(TOOL_SENTINEL)
|| entry.summary.contains(ASSISTANT_SENTINEL)
})
);
let llm: &'static RecordingLlm = Box::leak(Box::new(RecordingLlm(Mutex::new(None))));
@@ -6641,9 +6795,11 @@ mod tests {
.take()
.expect("review called recording LLM");
let request_text = format!("{:?}{:?}", request.system, request.messages);
assert!(!request_text.contains(SENTINEL));
assert!(!request_text.contains(TOOL_SENTINEL));
assert!(!request_text.contains(ASSISTANT_SENTINEL));
let report_text = serde_json::to_string(&report).unwrap();
assert!(!report_text.contains(SENTINEL));
assert!(!report_text.contains(TOOL_SENTINEL));
assert!(!report_text.contains(ASSISTANT_SENTINEL));
// Review is read-only: no pending sidecar or approved page is created.
assert!(
state
@@ -158,6 +158,7 @@ impl MultiUserHarness {
DEFAULT_HOOK_INGEST_MAX_IN_FLIGHT,
)),
consolidate_on_session_end: false,
capture_assistant_enabled: false,
subagent_sessions: Arc::new(tokio::sync::Mutex::new(SubagentSessionSet::default())),
ingest_rate: Arc::new(tokio::sync::Mutex::new(
ai_memory_hooks::IngestRateLimiter::disabled(),
@@ -157,6 +157,7 @@ impl Harness {
DEFAULT_HOOK_INGEST_MAX_IN_FLIGHT,
)),
consolidate_on_session_end: false,
capture_assistant_enabled: false,
subagent_sessions: Arc::new(tokio::sync::Mutex::new(SubagentSessionSet::default())),
ingest_rate: Arc::new(tokio::sync::Mutex::new(
ai_memory_hooks::IngestRateLimiter::disabled(),
+3
View File
@@ -340,6 +340,9 @@ that touch the relevant area.
(agentmemory #221 / #143.)
6. **Privacy strip is a typed boundary.** `Sanitized<NewObservation>`
has no other constructor than `sanitize()`. (design-decisions §14.)
The opt-in assistant/Stop excerpt (#196) enters through this same
boundary: the client sanitizes it before it reaches the wire, and the
server re-scrubs it here with its configured patterns before the write.
7. **JSON-schema structured outputs only.** Native provider JSON
modes; no XML, no Instructor wrapping. (agentmemory #492 / #539,
cognee #2840.)
+38 -13
View File
@@ -393,21 +393,46 @@ capability output reflects the selected integration. See the canonical
[capture exclusions reference](marker-file.md#capture-exclusions).
Some agent harnesses attach the assistant's final turn to their `Stop` event —
Claude Code sends it as a raw `last_assistant_message`. That text is never
persisted, and the native hook binary strips the raw field before it can reach
the local spool or the wire; the server strips it defensively on arrival too.
Optional assistant/Stop capture proposed in issue #196 remains disabled.
Claude Code sends it as a raw `last_assistant_message`. By default that text is
never persisted: the native hook binary strips the raw field before it can reach
the local spool or the wire, and the server strips it defensively on arrival.
**Opt-in capture (#196).** You can opt in to storing a sanitized, 2 KB-capped
excerpt of the assistant's final turn as the Stop body. It is a **double
opt-in** — enable the server first, then the client:
1. **Server:** set `capture_assistant = true` in the live
`<data_dir>/config.toml` (or the service's configured TOML file), or set
`AI_MEMORY_CAPTURE_ASSISTANT=true`, then restart `ai-memory serve`.
2. **Client:** re-install the Claude Code hooks with the flag:
```bash
ai-memory install-hooks --agent claude-code --capture-assistant --apply
```
The client sanitizes (built-in patterns) and truncates the excerpt before it
touches the spool or wire; the server re-scrubs with its `[sanitize]` patterns
before storing. If either side is off — or the marker is malformed — the Stop
stays empty. Re-running `install-hooks` without `--capture-assistant` removes
the flag (idempotent). `--capture-assistant` is Claude Code + native-platform
only; on any other agent or the script fallback the installer refuses it rather
than enabling something that cannot take effect. Assistant text is
privacy-sensitive — read the `SECURITY.md` notes on what it can contain and where
it flows (consolidation/reviewer prompts, and out to a cloud LLM provider if one
is configured) before enabling it.
Upgrading the binary is sufficient for native Claude Code installs, and pending
spooled events drain with the field stripped as well. Installs that run the
spooled events drain with the raw field stripped as well. Installs that run the
`.sh`/`.ps1` script fallback (the Docker script bundle or an explicit
`AI_MEMORY_HOOK_PLATFORM=posix`) still POST the raw field on the local wire
until they move to native commands. The Docker wrapper deliberately keeps
script commands because a binary path inside its helper container is not valid
on the host; running `install-hooks` through that wrapper refreshes the scripts
but does not convert them. To close the local-wire exposure, install a native
ai-memory client on the agent host, then use that native executable to run
`install-hooks --agent claude-code --apply`. If the script fallback is retained,
the server still strips the field immediately on receipt before persistence.
`AI_MEMORY_HOOK_PLATFORM=posix`) cannot sanitize the assistant text, so a `Stop`
payload still carrying the raw field is dropped whole by the script rather than
POSTed verbatim. The Docker wrapper deliberately keeps script commands because a
binary path inside its helper container is not valid on the host; running
`install-hooks` through that wrapper refreshes the scripts but does not convert
them. To capture assistant text safely, install a native ai-memory client on the
agent host, then use that native executable to run
`install-hooks --agent claude-code --apply`. Even if the script fallback is
retained, the server still strips any raw field on receipt before persistence.
Native `ai-memory hook --event ...` commands spool events locally. Session start
does a short bounded cleanup drain before fetching a handoff; cancellation-prone
+7 -3
View File
@@ -207,9 +207,13 @@ its existing tool-response/error excerpt and caps the complete rendered body at
body, and association is only by matching agent-provided call IDs. User-prompt stores its prompt
text, notification stores its message/text, and post-compaction stores its
summary; other event bodies are currently empty unless explicitly supported.
Stop and assistant-message capture remain disabled and deferred. The metadata
header is closed; the PostToolUse response/error excerpt remains the existing
bounded content capture.
Stop/assistant-message capture is disabled by default and never persisted; it is
available only through the explicit double opt-in described in the install guide
(`install-hooks --capture-assistant` on the client plus `capture_assistant` on
the server), where the excerpt is sanitized on both sides and capped. It is not
gated by this marker file — assistant text is not path-attributable, so a
`.ai-memory.toml` cannot narrow it. The metadata header is closed; the
PostToolUse response/error excerpt remains the existing bounded content capture.
Capture exclusions are evaluated only where paths have a proven schema, so they
do not claim to filter those other bodies.
+11
View File
@@ -13,6 +13,17 @@ _lib_dir="$(dirname "$0")"
SERVER="${AI_MEMORY_HOOK_URL:-http://127.0.0.1:49374}"
PAYLOAD=$(cat)
# Assistant/Stop capture (#196) is native-only: the script fallback cannot
# sanitize the assistant message, so if the raw payload still carries the field
# we drop the whole Stop rather than POST it verbatim. A literal substring check
# (POSIX `case`, no bash-isms per _lib.sh) — conservative on purpose: any Stop
# mentioning the key is dropped. Move to a native install to capture it safely.
case "$PAYLOAD" in
*'"last_assistant_message"'*)
printf '{}\n'
exit 0
;;
esac
CWD=$(ai_memory_extract_cwd "$PAYLOAD")
QS=$(ai_memory_marker_qs "$CWD")
+7
View File
@@ -237,6 +237,13 @@ function Invoke-AiMemoryHook {
$Server = if ($env:AI_MEMORY_HOOK_URL) { $env:AI_MEMORY_HOOK_URL } else { "http://127.0.0.1:49374" }
$Payload = Read-AiMemoryStdin
# Assistant/Stop capture (#196) is native-only. Scoped to claude-code + stop
# so a PostToolUse whose tool output legitimately contains the literal string
# is unaffected. If a Stop payload still carries the raw field, drop the whole
# event rather than POST it verbatim (the script fallback cannot sanitize it).
if ($Agent -eq "claude-code" -and $Event -eq "stop" -and $Payload -and $Payload.Contains('"last_assistant_message"')) {
return
}
$Cwd = Resolve-AiMemoryCwd -Payload $Payload -Agent $Agent
$QS = Get-AiMemoryMarkerQuery -Cwd $Cwd
if ($env:AI_MEMORY_RUN_ID) {