refactor(integrations): shared agent plugin runtimes and native adapters (#7203)
This commit is contained in:
Kartik
@@ -3,367 +3,52 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "core"))
|
||||
_here = Path(__file__).resolve()
|
||||
_bundled_core = _here.parents[2] / "core"
|
||||
_core_dir = _bundled_core if (_bundled_core / "memory_core.py").is_file() else _bundled_core / "python"
|
||||
sys.path.insert(0, str(_core_dir))
|
||||
sys.path.insert(0, str(_here.parent))
|
||||
|
||||
import telemetry
|
||||
from memory_core import (
|
||||
EvidenceStore,
|
||||
api_key,
|
||||
cache_plugin_api_key,
|
||||
checkpoint_session,
|
||||
clear_stale_api_key_cache,
|
||||
data_dir,
|
||||
detached_process_kwargs,
|
||||
format_context,
|
||||
record_session_start,
|
||||
import telemetry # noqa: E402
|
||||
from memory_core import ( # noqa: E402
|
||||
configure_harness,
|
||||
record_sidekick_start,
|
||||
record_sidekick_stop,
|
||||
record_stop,
|
||||
record_tool,
|
||||
record_user_prompt,
|
||||
search_memories,
|
||||
)
|
||||
from transcript import record_stop # noqa: E402
|
||||
import hook_runner # noqa: E402
|
||||
|
||||
configure_harness("claude-code", data_dir_name="claude-code-plugin", source_tag="claude_code_plugin")
|
||||
telemetry.init(harness="claude-code", source_tag="CLAUDE_CODE_PLUGIN")
|
||||
|
||||
|
||||
STALE_RUNNING_SECONDS = 300
|
||||
PENDING_EXPIRY_SECONDS = 7 * 24 * 60 * 60
|
||||
PENDING_LAUNCH_LIMIT = 5
|
||||
def _sidekick_start(store, hook_input):
|
||||
context = record_sidekick_start(store, hook_input)
|
||||
if context:
|
||||
return {
|
||||
"hookSpecificOutput": {
|
||||
"hookEventName": "SubagentStart",
|
||||
"additionalContext": context,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def read_hook_input() -> dict:
|
||||
try:
|
||||
value = json.load(sys.stdin)
|
||||
return value if isinstance(value, dict) else {}
|
||||
except (json.JSONDecodeError, OSError):
|
||||
return {}
|
||||
|
||||
|
||||
def first_prompt_memory_output(store: EvidenceStore, hook_input: dict) -> dict:
|
||||
"""Search once before Claude handles the first prompt in a session."""
|
||||
repo, session_id, prompt, is_first_prompt = record_user_prompt(store, hook_input)
|
||||
if not is_first_prompt:
|
||||
return {}
|
||||
try:
|
||||
minimum_query_chars = int(os.environ.get("MEM0_CODE_MIN_QUERY_CHARS", "20"))
|
||||
except ValueError:
|
||||
minimum_query_chars = 20
|
||||
if len(prompt.strip()) < max(minimum_query_chars, 1):
|
||||
return {}
|
||||
result = search_memories(
|
||||
store,
|
||||
repo,
|
||||
session_id,
|
||||
prompt,
|
||||
top_k=5,
|
||||
operation="first-prompt-search",
|
||||
timeout=2,
|
||||
)
|
||||
if not result.memories:
|
||||
return {}
|
||||
context = format_context(
|
||||
result.memories,
|
||||
"Mem0 found these relevant memories from earlier work in this repository:",
|
||||
)
|
||||
telemetry.record(
|
||||
"context_injected",
|
||||
repo=repo,
|
||||
session_id=session_id,
|
||||
trigger="first-prompt",
|
||||
memory_count=len(result.memories),
|
||||
context_chars=len(context),
|
||||
prompt_chars=len(prompt),
|
||||
)
|
||||
return {
|
||||
"hookSpecificOutput": {
|
||||
"hookEventName": "UserPromptSubmit",
|
||||
"additionalContext": context,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _launch_handoff(handoff_path: Path) -> bool:
|
||||
running_path = handoff_path.with_suffix(".running")
|
||||
try:
|
||||
handoff_path.replace(running_path)
|
||||
except OSError:
|
||||
return False
|
||||
worker = Path(__file__).resolve().parents[2] / "core" / "flush_worker.py"
|
||||
log_path = data_dir() / "flush-worker.log"
|
||||
log_handle = open(log_path, "a", encoding="utf-8")
|
||||
try:
|
||||
subprocess.Popen(
|
||||
[sys.executable, str(worker), str(running_path)],
|
||||
stdin=subprocess.DEVNULL,
|
||||
stdout=log_handle,
|
||||
stderr=log_handle,
|
||||
close_fds=True,
|
||||
**detached_process_kwargs(),
|
||||
)
|
||||
finally:
|
||||
log_handle.close()
|
||||
return True
|
||||
|
||||
|
||||
def recover_pending_handoffs() -> int:
|
||||
pending_dir = data_dir() / "pending"
|
||||
pending_dir.mkdir(parents=True, exist_ok=True)
|
||||
now = time.time()
|
||||
for running in pending_dir.glob("*.running"):
|
||||
try:
|
||||
if now - running.stat().st_mtime > STALE_RUNNING_SECONDS:
|
||||
running.replace(running.with_suffix(".json"))
|
||||
except OSError:
|
||||
continue
|
||||
recoverable = []
|
||||
for handoff in pending_dir.glob("*.json"):
|
||||
try:
|
||||
age = now - handoff.stat().st_mtime
|
||||
except OSError:
|
||||
continue
|
||||
if age > PENDING_EXPIRY_SECONDS:
|
||||
handoff.unlink(missing_ok=True)
|
||||
continue
|
||||
recoverable.append((age, handoff))
|
||||
recoverable.sort(key=lambda item: item[0], reverse=True)
|
||||
launched = 0
|
||||
for _, handoff in recoverable[:PENDING_LAUNCH_LIMIT]:
|
||||
launched += int(_launch_handoff(handoff))
|
||||
return launched
|
||||
|
||||
|
||||
def refresh_pending_handoffs() -> None:
|
||||
"""Hold unsent packets while paused instead of letting them expire."""
|
||||
pending_dir = data_dir() / "pending"
|
||||
if not pending_dir.is_dir():
|
||||
return
|
||||
for pattern in ("*.json", "*.running"):
|
||||
for handoff in pending_dir.glob(pattern):
|
||||
try:
|
||||
os.utime(handoff)
|
||||
except OSError:
|
||||
continue
|
||||
|
||||
|
||||
def hand_off_flush(
|
||||
hook_input: dict, reason: str, *, wait_for_inflight: bool = False
|
||||
) -> None:
|
||||
"""Persist hook input and detach delivery from Claude's shutdown lifecycle."""
|
||||
pending_dir = data_dir() / "pending"
|
||||
pending_dir.mkdir(parents=True, exist_ok=True)
|
||||
material = (
|
||||
f"{hook_input.get('cwd', '')}\0{hook_input.get('session_id', '')}\0{reason}"
|
||||
)
|
||||
digest = hashlib.sha256(material.encode()).hexdigest()[:24]
|
||||
handoff_path = pending_dir / f"{digest}-{uuid.uuid4().hex[:8]}.json"
|
||||
temporary_path = handoff_path.with_suffix(".tmp")
|
||||
temporary_path.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"hook_input": hook_input,
|
||||
"reason": reason,
|
||||
"wait_for_inflight": wait_for_inflight,
|
||||
}
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
temporary_path.replace(handoff_path)
|
||||
_launch_handoff(handoff_path)
|
||||
|
||||
|
||||
def automatic_flush_enabled() -> bool:
|
||||
return os.environ.get("MEM0_CODE_AUTO_FLUSH", "true").lower() in {
|
||||
"1",
|
||||
"true",
|
||||
"yes",
|
||||
"on",
|
||||
}
|
||||
|
||||
|
||||
def schedule_periodic_checkpoint(
|
||||
store: EvidenceStore,
|
||||
hook_input: dict,
|
||||
repo,
|
||||
session_id: str,
|
||||
) -> bool:
|
||||
"""Start one background extraction when a complete block is ready."""
|
||||
if (
|
||||
not automatic_flush_enabled()
|
||||
or not api_key()
|
||||
or not store.checkpoint_due(repo.identity, session_id)
|
||||
):
|
||||
return False
|
||||
if store.prepare_flush(repo, session_id, "periodic") is None:
|
||||
return False
|
||||
hand_off_flush(hook_input, "periodic")
|
||||
return True
|
||||
|
||||
|
||||
DEFAULT_IDLE_FLUSH_SECONDS = 300
|
||||
|
||||
|
||||
def _idle_flush_seconds() -> int:
|
||||
try:
|
||||
return max(
|
||||
int(os.environ.get("MEM0_CODE_IDLE_FLUSH_SECONDS", str(DEFAULT_IDLE_FLUSH_SECONDS))),
|
||||
0,
|
||||
)
|
||||
except ValueError:
|
||||
return DEFAULT_IDLE_FLUSH_SECONDS
|
||||
|
||||
|
||||
def schedule_idle_flush(
|
||||
store: EvidenceStore,
|
||||
hook_input: dict,
|
||||
repo,
|
||||
session_id: str,
|
||||
) -> bool:
|
||||
"""Launch a delayed background flush for sessions that may never end."""
|
||||
delay = _idle_flush_seconds()
|
||||
if delay <= 0 or not automatic_flush_enabled() or not api_key():
|
||||
return False
|
||||
if store.has_inflight_flush(repo.identity, session_id):
|
||||
return False
|
||||
if not store.has_unflushed_events(repo.identity, session_id):
|
||||
return False
|
||||
pending_dir = data_dir() / "pending"
|
||||
pending_dir.mkdir(parents=True, exist_ok=True)
|
||||
material = f"idle\0{hook_input.get('cwd', '')}\0{hook_input.get('session_id', '')}"
|
||||
digest = hashlib.sha256(material.encode()).hexdigest()[:24]
|
||||
for old in pending_dir.glob(f"idle-{digest}*"):
|
||||
old.unlink(missing_ok=True)
|
||||
handoff_path = pending_dir / f"idle-{digest}-{uuid.uuid4().hex[:8]}.json"
|
||||
temporary_path = handoff_path.with_suffix(".tmp")
|
||||
temporary_path.write_text(
|
||||
json.dumps({
|
||||
"hook_input": hook_input,
|
||||
"reason": "idle",
|
||||
"delay_seconds": delay,
|
||||
}),
|
||||
encoding="utf-8",
|
||||
)
|
||||
temporary_path.replace(handoff_path)
|
||||
_launch_handoff(handoff_path)
|
||||
return True
|
||||
|
||||
|
||||
def log_failure(exc: Exception) -> None:
|
||||
try:
|
||||
log_path = data_dir() / "plugin-errors.log"
|
||||
with log_path.open("a", encoding="utf-8") as handle:
|
||||
handle.write(f"{time.time():.3f} {type(exc).__name__}: {exc}\n")
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"action",
|
||||
choices=[
|
||||
"session-start",
|
||||
"user-prompt",
|
||||
"post-tool",
|
||||
"post-tool-failure",
|
||||
"sidekick-start",
|
||||
"sidekick-stop",
|
||||
"stop",
|
||||
"flush",
|
||||
],
|
||||
)
|
||||
parser.add_argument("--reason", default="manual")
|
||||
parser.add_argument("--plugin-data-dir", default="")
|
||||
args = parser.parse_args()
|
||||
if args.plugin_data_dir:
|
||||
os.environ["MEM0_CODE_DATA_DIR"] = args.plugin_data_dir
|
||||
cache_plugin_api_key()
|
||||
if args.action == "session-start":
|
||||
clear_stale_api_key_cache()
|
||||
hook_input = read_hook_input()
|
||||
store = EvidenceStore()
|
||||
try:
|
||||
if store.is_paused():
|
||||
if args.action == "session-start":
|
||||
refresh_pending_handoffs()
|
||||
telemetry.record("session_start", paused=True)
|
||||
telemetry.spawn_flush()
|
||||
return 0
|
||||
if args.action == "session-start":
|
||||
if telemetry.is_first_run():
|
||||
telemetry.record("install")
|
||||
recovered = recover_pending_handoffs()
|
||||
record_session_start(store, hook_input)
|
||||
if recovered:
|
||||
telemetry.record("handoff_recovered", count=recovered)
|
||||
telemetry.spawn_flush()
|
||||
elif args.action == "user-prompt":
|
||||
output = first_prompt_memory_output(store, hook_input)
|
||||
if output:
|
||||
print(json.dumps(output))
|
||||
elif args.action == "post-tool":
|
||||
record_tool(store, hook_input)
|
||||
elif args.action == "post-tool-failure":
|
||||
record_tool(store, hook_input, failed=True)
|
||||
elif args.action == "sidekick-start":
|
||||
context = record_sidekick_start(store, hook_input)
|
||||
if context:
|
||||
print(
|
||||
json.dumps(
|
||||
{
|
||||
"hookSpecificOutput": {
|
||||
"hookEventName": "SubagentStart",
|
||||
"additionalContext": context,
|
||||
}
|
||||
}
|
||||
)
|
||||
)
|
||||
elif args.action == "sidekick-stop":
|
||||
record_sidekick_stop(store, hook_input)
|
||||
elif args.action == "stop":
|
||||
repo, session_id = record_stop(store, hook_input)
|
||||
if not schedule_periodic_checkpoint(store, hook_input, repo, session_id):
|
||||
schedule_idle_flush(store, hook_input, repo, session_id)
|
||||
elif args.action == "flush":
|
||||
automatic = args.reason in {"session-end", "pre-compact"}
|
||||
if automatic and not automatic_flush_enabled():
|
||||
return 0
|
||||
if args.reason == "session-end":
|
||||
# In print mode, SessionEnd can arrive before the Stop hook has
|
||||
# recorded Claude's final response. Read any remaining visible
|
||||
# transcript messages before preparing the final extraction.
|
||||
record_stop(store, hook_input)
|
||||
if os.environ.get("MEM0_CODE_SYNC_FLUSH") == "1":
|
||||
print(json.dumps(checkpoint_session(store, hook_input, args.reason)))
|
||||
else:
|
||||
session_id = str(hook_input.get("session_id") or "unknown-session")
|
||||
repo = store.repo_for_session(session_id, hook_input.get("cwd"))
|
||||
already_running = store.has_inflight_flush(repo.identity, session_id)
|
||||
if already_running and args.reason == "session-end":
|
||||
hand_off_flush(
|
||||
hook_input, args.reason, wait_for_inflight=True
|
||||
)
|
||||
elif not already_running and store.prepare_flush(
|
||||
repo, session_id, args.reason
|
||||
) is not None:
|
||||
hand_off_flush(hook_input, args.reason)
|
||||
finally:
|
||||
store.close()
|
||||
return 0
|
||||
def _sidekick_stop(store, hook_input):
|
||||
record_sidekick_stop(store, hook_input)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
raise SystemExit(main())
|
||||
except Exception as exc:
|
||||
# Memory must never prevent the coding agent from continuing.
|
||||
log_failure(exc)
|
||||
raise SystemExit(0)
|
||||
hook_runner.entry_point(
|
||||
record_stop_fn=record_stop,
|
||||
extra_actions={
|
||||
"post-tool-failure": lambda s, h: record_tool(s, h, failed=True),
|
||||
"sidekick-start": _sidekick_start,
|
||||
"sidekick-stop": _sidekick_stop,
|
||||
},
|
||||
data_dir_env="MEM0_CODE_DATA_DIR",
|
||||
automatic_flush_reasons={"session-end", "pre-compact"},
|
||||
)
|
||||
|
||||
@@ -0,0 +1,362 @@
|
||||
"""Claude Code transcript parsing for Mem0 memory extraction."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import html
|
||||
import json
|
||||
import re
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
_here = Path(__file__).resolve()
|
||||
sys.path.insert(0, str(_here.parents[2] / "core" / "python"))
|
||||
|
||||
from memory_core import ( # noqa: E402
|
||||
EvidenceStore,
|
||||
RepoContext,
|
||||
_session_id,
|
||||
redact,
|
||||
)
|
||||
|
||||
|
||||
def _message_content_text(content: Any) -> str:
|
||||
"""Return visible text from one Claude transcript message."""
|
||||
if isinstance(content, str):
|
||||
return redact(content).strip()
|
||||
if not isinstance(content, list):
|
||||
return ""
|
||||
parts = []
|
||||
for block in content:
|
||||
if not isinstance(block, dict) or block.get("type") != "text":
|
||||
continue
|
||||
text = redact(block.get("text", "")).strip()
|
||||
if text:
|
||||
parts.append(text)
|
||||
return "\n\n".join(parts)
|
||||
|
||||
|
||||
def _transcript_rows(
|
||||
path: str, offset: int = 0
|
||||
) -> tuple[list[dict[str, Any]], int, bool]:
|
||||
"""Parse transcript rows from a byte offset, returning rows, end offset, and whether the offset was honored."""
|
||||
if not path:
|
||||
return [], 0, False
|
||||
rows = []
|
||||
try:
|
||||
resolved = Path(path).expanduser()
|
||||
if not 0 <= offset <= resolved.stat().st_size:
|
||||
offset = 0
|
||||
end = offset
|
||||
with resolved.open("rb") as handle:
|
||||
handle.seek(offset)
|
||||
for line in handle:
|
||||
if not line.endswith(b"\n"):
|
||||
break
|
||||
end += len(line)
|
||||
try:
|
||||
row = json.loads(line.decode("utf-8", errors="replace"))
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
if isinstance(row, dict) and row.get("uuid"):
|
||||
rows.append(row)
|
||||
except OSError:
|
||||
return [], offset, False
|
||||
return rows, end, offset > 0
|
||||
|
||||
|
||||
def _active_transcript_chain(
|
||||
rows: list[dict[str, Any]], session_id: str
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Follow the current Claude conversation branch from its latest record."""
|
||||
by_uuid = {str(row["uuid"]): row for row in rows if row.get("uuid")}
|
||||
leaf = next(
|
||||
(
|
||||
row
|
||||
for row in reversed(rows)
|
||||
if not row.get("isSidechain")
|
||||
and str(row.get("sessionId") or "") == session_id
|
||||
),
|
||||
None,
|
||||
)
|
||||
if leaf is None:
|
||||
return []
|
||||
|
||||
chain = []
|
||||
seen = set()
|
||||
current = leaf
|
||||
while current is not None:
|
||||
uuid = str(current.get("uuid") or "")
|
||||
if not uuid or uuid in seen:
|
||||
break
|
||||
seen.add(uuid)
|
||||
chain.append(current)
|
||||
current = by_uuid.get(str(current.get("parentUuid") or ""))
|
||||
chain.reverse()
|
||||
return chain
|
||||
|
||||
|
||||
def _human_prompt_text(row: dict[str, Any]) -> str:
|
||||
if row.get("type") != "user":
|
||||
return ""
|
||||
origin = row.get("origin") or {}
|
||||
if isinstance(origin, dict) and origin.get("kind") not in {None, "human"}:
|
||||
return ""
|
||||
message = row.get("message") or {}
|
||||
content = message.get("content") if isinstance(message, dict) else None
|
||||
if not isinstance(content, str):
|
||||
return ""
|
||||
text = redact(content).strip()
|
||||
if text.startswith("<!-- attach -->"):
|
||||
text = text.removeprefix("<!-- attach -->").strip()
|
||||
ignored_prefixes = (
|
||||
"<local-command-caveat>",
|
||||
"<local-command-stdout>",
|
||||
"<local-command-stderr>",
|
||||
"<command-name>",
|
||||
"<system-reminder>",
|
||||
"<task-notification>",
|
||||
)
|
||||
return "" if text.startswith(ignored_prefixes) else text
|
||||
|
||||
|
||||
def _xml_value(text: str, tag: str) -> str:
|
||||
match = re.search(fr"<{tag}>(.*?)</{tag}>", text, re.DOTALL)
|
||||
return html.unescape(match.group(1).strip()) if match else ""
|
||||
|
||||
|
||||
def _agent_assignment(tool_input: dict[str, Any]) -> str:
|
||||
prompt = redact(tool_input.get("prompt", "")).strip()
|
||||
if not prompt:
|
||||
return ""
|
||||
agent_type = redact(tool_input.get("subagent_type", "agent")).strip() or "agent"
|
||||
description = redact(tool_input.get("description", "")).strip()
|
||||
heading = f"Subagent assignment ({agent_type}"
|
||||
if description:
|
||||
heading += f": {description}"
|
||||
return f"{heading}):\n{prompt}"
|
||||
|
||||
|
||||
def _agent_response(tool_input: dict[str, Any], result: str) -> str:
|
||||
result = redact(result).strip()
|
||||
if not result or result.startswith("Async agent launched successfully."):
|
||||
return ""
|
||||
agent_type = redact(tool_input.get("subagent_type", "agent")).strip() or "agent"
|
||||
description = redact(tool_input.get("description", "")).strip()
|
||||
heading = f"Subagent response ({agent_type}"
|
||||
if description:
|
||||
heading += f": {description}"
|
||||
return f"{heading}):\n{result}"
|
||||
|
||||
|
||||
def _tool_result_text(block: dict[str, Any]) -> str:
|
||||
return _message_content_text(block.get("content"))
|
||||
|
||||
|
||||
def transcript_extraction_messages(
|
||||
transcript_path: str,
|
||||
session_id: str,
|
||||
*,
|
||||
previous_leaf_uuid: str = "",
|
||||
prompt_hint: str = "",
|
||||
fallback_assistant_message: str = "",
|
||||
label_final_response: bool = False,
|
||||
start_offset: int = 0,
|
||||
) -> tuple[list[dict[str, str]], str, int]:
|
||||
"""Read the meaningful part of the current Claude exchange."""
|
||||
rows, end_offset, resumed = _transcript_rows(transcript_path, start_offset)
|
||||
chain = _active_transcript_chain(rows, session_id)
|
||||
if not chain:
|
||||
if resumed:
|
||||
return [], previous_leaf_uuid, end_offset
|
||||
fallback = redact(fallback_assistant_message).strip()
|
||||
return (
|
||||
([{"role": "assistant", "content": f"Main Claude response:\n{fallback}"}]
|
||||
if fallback
|
||||
else []),
|
||||
"",
|
||||
end_offset,
|
||||
)
|
||||
|
||||
leaf_uuid = str(chain[-1].get("uuid") or "")
|
||||
if previous_leaf_uuid and leaf_uuid == previous_leaf_uuid:
|
||||
return [], leaf_uuid, end_offset
|
||||
start = 0
|
||||
if previous_leaf_uuid:
|
||||
for index, row in enumerate(chain):
|
||||
if str(row.get("uuid") or "") == previous_leaf_uuid:
|
||||
start = index + 1
|
||||
break
|
||||
else:
|
||||
previous_leaf_uuid = ""
|
||||
if not previous_leaf_uuid and not resumed:
|
||||
prompt_hint = redact(prompt_hint).strip()
|
||||
candidates = [
|
||||
index
|
||||
for index, row in enumerate(chain)
|
||||
if _human_prompt_text(row)
|
||||
and (
|
||||
not prompt_hint
|
||||
or _human_prompt_text(row) == prompt_hint
|
||||
)
|
||||
]
|
||||
task_notifications = [
|
||||
index
|
||||
for index, row in enumerate(chain)
|
||||
if isinstance((row.get("message") or {}).get("content"), str)
|
||||
and (row.get("message") or {})["content"].startswith("<task-notification>")
|
||||
]
|
||||
if candidates:
|
||||
start = candidates[-1]
|
||||
elif task_notifications:
|
||||
start = task_notifications[-1]
|
||||
|
||||
tool_uses: dict[str, tuple[str, dict[str, Any]]] = {}
|
||||
for row in chain:
|
||||
message = row.get("message") or {}
|
||||
content = message.get("content") if isinstance(message, dict) else None
|
||||
if not isinstance(content, list):
|
||||
continue
|
||||
for block in content:
|
||||
if not isinstance(block, dict) or block.get("type") != "tool_use":
|
||||
continue
|
||||
tool_id = str(block.get("id") or "")
|
||||
tool_input = block.get("input") or {}
|
||||
if tool_id and isinstance(tool_input, dict):
|
||||
tool_uses[tool_id] = (str(block.get("name") or ""), tool_input)
|
||||
|
||||
output: list[dict[str, str]] = []
|
||||
|
||||
def append(role: str, content: str) -> None:
|
||||
content = redact(content).strip()
|
||||
if content:
|
||||
output.append({"role": role, "content": content})
|
||||
|
||||
for row in chain[start:]:
|
||||
message = row.get("message") or {}
|
||||
if not isinstance(message, dict):
|
||||
continue
|
||||
role = str(message.get("role") or "")
|
||||
content = message.get("content")
|
||||
|
||||
if role == "user" and isinstance(content, str):
|
||||
if content.startswith("<task-notification>"):
|
||||
if _xml_value(content, "status") != "completed":
|
||||
continue
|
||||
tool_id = _xml_value(content, "tool-use-id")
|
||||
tool = tool_uses.get(tool_id)
|
||||
result = _xml_value(content, "result")
|
||||
if tool and tool[0] == "Agent" and result:
|
||||
assignment = _agent_assignment(tool[1])
|
||||
response = _agent_response(tool[1], result)
|
||||
append("assistant", assignment)
|
||||
append("assistant", response)
|
||||
continue
|
||||
human = _human_prompt_text(row)
|
||||
if human:
|
||||
append("user", human)
|
||||
continue
|
||||
|
||||
if not isinstance(content, list):
|
||||
continue
|
||||
for block in content:
|
||||
if not isinstance(block, dict):
|
||||
continue
|
||||
block_type = block.get("type")
|
||||
if role == "assistant" and block_type == "text":
|
||||
append("assistant", str(block.get("text") or ""))
|
||||
continue
|
||||
if role != "user" or block_type != "tool_result":
|
||||
continue
|
||||
|
||||
tool_id = str(block.get("tool_use_id") or "")
|
||||
tool = tool_uses.get(tool_id)
|
||||
if not tool:
|
||||
continue
|
||||
name, tool_input = tool
|
||||
result = _tool_result_text(block)
|
||||
failed = bool(block.get("is_error"))
|
||||
if name == "Agent" and not failed:
|
||||
response = _agent_response(tool_input, result)
|
||||
if response:
|
||||
append("assistant", _agent_assignment(tool_input))
|
||||
append("assistant", response)
|
||||
elif name == "AskUserQuestion" and result and not failed:
|
||||
append("user", f"User answers to Claude's questions:\n{result}")
|
||||
elif name == "ExitPlanMode" and not failed:
|
||||
plan = redact(tool_input.get("plan", "")).strip()
|
||||
if plan:
|
||||
append("assistant", f"Approved implementation plan:\n{plan}")
|
||||
|
||||
fallback = redact(fallback_assistant_message).strip()
|
||||
if fallback:
|
||||
labeled = f"Main Claude response:\n{fallback}"
|
||||
for message in reversed(output):
|
||||
if message["role"] == "assistant" and message["content"] == fallback:
|
||||
message["content"] = labeled
|
||||
break
|
||||
else:
|
||||
append("assistant", labeled)
|
||||
elif label_final_response:
|
||||
last_message = chain[-1].get("message") or {}
|
||||
last_content = (
|
||||
last_message.get("content") if isinstance(last_message, dict) else None
|
||||
)
|
||||
final_parts = [
|
||||
redact(block.get("text", "")).strip()
|
||||
for block in (last_content if isinstance(last_content, list) else [])
|
||||
if isinstance(block, dict)
|
||||
and block.get("type") == "text"
|
||||
and redact(block.get("text", "")).strip()
|
||||
]
|
||||
for start in range(len(output) - len(final_parts), -1, -1):
|
||||
candidate = output[start : start + len(final_parts)]
|
||||
if final_parts and [item["content"] for item in candidate] == final_parts:
|
||||
candidate[0]["content"] = (
|
||||
f"Main Claude response:\n{candidate[0]['content']}"
|
||||
)
|
||||
break
|
||||
return output, leaf_uuid, end_offset
|
||||
|
||||
|
||||
def record_stop(
|
||||
store: EvidenceStore, hook_input: dict[str, Any]
|
||||
) -> tuple[RepoContext, str]:
|
||||
session_id = _session_id(hook_input)
|
||||
repo = store.repo_for_session(session_id, hook_input.get("cwd"))
|
||||
message = redact(hook_input.get("last_assistant_message", "")).strip()
|
||||
previous_stop = store.latest_event_payload(
|
||||
repo.identity, session_id, "assistant_stop"
|
||||
)
|
||||
latest_prompt = store.latest_event_payload(repo.identity, session_id, "user_prompt")
|
||||
transcript_path = str(hook_input.get("transcript_path") or "")
|
||||
previous_offset = previous_stop.get("transcript_offset")
|
||||
start_offset = (
|
||||
previous_offset
|
||||
if isinstance(previous_offset, int)
|
||||
and str(previous_stop.get("transcript_path") or "") == transcript_path
|
||||
else 0
|
||||
)
|
||||
transcript_messages, leaf_uuid, end_offset = transcript_extraction_messages(
|
||||
transcript_path,
|
||||
session_id,
|
||||
previous_leaf_uuid=str(previous_stop.get("transcript_leaf_uuid") or ""),
|
||||
prompt_hint=str(latest_prompt.get("text") or ""),
|
||||
fallback_assistant_message=message,
|
||||
label_final_response=True,
|
||||
start_offset=start_offset,
|
||||
)
|
||||
if transcript_messages:
|
||||
payload: dict[str, Any] = {
|
||||
"text": message,
|
||||
"transcript_messages": transcript_messages,
|
||||
}
|
||||
if leaf_uuid:
|
||||
payload["transcript_leaf_uuid"] = leaf_uuid
|
||||
if transcript_path:
|
||||
payload["transcript_path"] = transcript_path
|
||||
payload["transcript_offset"] = end_offset
|
||||
store.record_event(
|
||||
repo, session_id, "assistant_stop", payload
|
||||
)
|
||||
return repo, session_id
|
||||
Reference in New Issue
Block a user