mirror of
https://github.com/akitaonrails/ai-memory.git
synced 2026-10-02 03:24:46 +08:00
Merge remote-tracking branch 'origin/main' into audit/pr-223-fix
# Conflicts: # CHANGELOG.md
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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
@@ -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.
|
||||
|
||||
|
||||
@@ -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`.
|
||||
|
||||
@@ -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(),
|
||||
)),
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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> {
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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
@@ -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
@@ -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.
|
||||
|
||||
|
||||
@@ -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")
|
||||
|
||||
|
||||
@@ -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) {
|
||||
|
||||
Reference in New Issue
Block a user