diff --git a/.env.example b/.env.example index f1c207c181..77ca0f3a38 100644 --- a/.env.example +++ b/.env.example @@ -98,6 +98,8 @@ # AGENTMEMORY_GRAPH_WEIGHT=0.2 # Graph traversal bonus on smart-search ranking # TOKEN_BUDGET=2000 # Max tokens injected via mem::context per session # MAX_OBS_PER_SESSION=500 # Per-session observation cap before consolidation kicks in +# SUMMARIZE_CHUNK_SIZE=400 # When mem::summarize sees a session larger than this, it chunks observations and map-reduces (chunk-summarize → reduce-merge) to stay within the LLM's context window. Default 400 ≈ 50k tokens per chunk at ~110 tok/obs. Native sessions are capped by MAX_OBS_PER_SESSION; chunking primarily matters for bulk-imported jsonl sessions, which bypass that cap. +# SUMMARIZE_CHUNK_CONCURRENCY=6 # Parallel chunk LLM calls during chunked summarize. Default 6 fits ~100-chunk sessions under iii's 180s function-invocation timeout at typical ~8s/call. High-throughput providers (Novita, DeepInfra, DeepSeek) commonly allow 100+ concurrent — bump this for very large imported sessions. # ----------------------------------------------------------------------------- # 5. Behaviour flags diff --git a/.gitignore b/.gitignore index 9a9260b8ca..585d0f494c 100644 --- a/.gitignore +++ b/.gitignore @@ -13,6 +13,8 @@ dist/ plugin/scripts/*.map plugin/scripts/*.d.mts data/ +data-*/ +agentmemory-debug/ .gstack/ # Lock files — never commit (see feedback_no_lockfiles memory) diff --git a/scripts/backfill-imported-sessions.sh b/scripts/backfill-imported-sessions.sh new file mode 100755 index 0000000000..a247a57e11 --- /dev/null +++ b/scripts/backfill-imported-sessions.sh @@ -0,0 +1,259 @@ +#!/usr/bin/env bash +# Backfill memory artifacts for sessions imported via `agentmemory import-jsonl`. +# +# The import path only persists Session + Observation rows (via synthetic, +# zero-LLM compression) and the deterministic crystal/lesson derivation. +# It does NOT call mem::summarize, so the semantic/procedural/reflect tiers +# of the consolidation pipeline have nothing to roll up. +# +# This script walks every session tagged `jsonl-import` and: +# 1. POSTs /agentmemory/summarize per session (LLM call) +# 2. POSTs /agentmemory/consolidate-pipeline once at the end +# +# Graph extraction (/agentmemory/graph/extract) is intentionally skipped — +# its API takes a per-observation payload, which is cost-prohibitive for +# bulk imports. `reflect` falls back to a no-graph clustering mode. +# +# Usage: +# scripts/backfill-imported-sessions.sh --dry-run +# scripts/backfill-imported-sessions.sh --limit 5 +# scripts/backfill-imported-sessions.sh # process all + +set -euo pipefail + +URL="${AGENTMEMORY_URL:-http://localhost:3111}" +DRY_RUN=0 +LIMIT=0 # 0 = no limit +ONLY_TAG="jsonl-import" +SKIP_CONSOLIDATE=0 +SKIP_AGENTS=0 # drop sessions whose project starts with "agent-" +MAX_OBS=0 # 0 = no cap; skip sessions with more observations than this +DEBUG_ON_ERROR=0 # on failure, dump session metadata + obs to DEBUG_DIR +DEBUG_DIR="${AGENTMEMORY_DEBUG_DIR:-./agentmemory-debug}" +PROJECT_PATTERN="" # jq test() regex against .project; "" means no filter + +# Cost-estimate knobs (defaults tuned for DeepSeek V4 Flash on DeepInfra: +# $0.14 / 1M input, $0.28 / 1M output). Override via env if needed. +COST_IN_PER_1M="${AGENTMEMORY_COST_IN_PER_1M:-0.14}" +COST_OUT_PER_1M="${AGENTMEMORY_COST_OUT_PER_1M:-0.28}" +# Rough token weight per compressed observation, derived from inspecting +# real synthetic-compression payloads in the kv store (mostly 100-300 tok, +# heavy-tailed). Override if your sessions are unusually verbose. +TOKENS_PER_OBS="${AGENTMEMORY_TOKENS_PER_OBS:-200}" +# Reserved per-call output budget (XML summary is small). +TOKENS_OUT_PER_SESSION="${AGENTMEMORY_TOKENS_OUT_PER_SESSION:-500}" + +while [[ $# -gt 0 ]]; do + case "$1" in + --dry-run) DRY_RUN=1; shift ;; + --limit) LIMIT="${2:?--limit needs a number}"; shift 2 ;; + --tag) ONLY_TAG="${2:?--tag needs a value (use empty string for all)}"; shift 2 ;; + --skip-consolidate) SKIP_CONSOLIDATE=1; shift ;; + --skip-agents) SKIP_AGENTS=1; shift ;; + --max-obs) MAX_OBS="${2:?--max-obs needs a number}"; shift 2 ;; + --debug-on-error) DEBUG_ON_ERROR=1; shift ;; + --project-pattern) PROJECT_PATTERN="${2:?--project-pattern needs a regex}"; shift 2 ;; + -h|--help) + sed -n '2,28p' "$0" + exit 0 ;; + *) echo "unknown flag: $1" >&2; exit 2 ;; + esac +done + +for bin in curl jq; do + command -v "$bin" >/dev/null || { echo "missing dependency: $bin" >&2; exit 1; } +done + +# Curl timeout profiles. Metadata reads (livez, sessions list, observations +# pull for debug dumps) should fail fast and retry transient blips. The LLM +# work calls (summarize, consolidate) intentionally have no --retry and a +# wide --max-time: each call can legitimately take minutes for chunked +# summarize on large sessions, and retrying a half-finished LLM job is +# expensive both in dollars and in duplicated server-side work. +META_CURL_OPTS=(--connect-timeout 10 --max-time 30 --retry 2 --retry-delay 1) +WORK_CURL_OPTS=(--connect-timeout 10 --max-time 1800) + +echo "agentmemory backfill — server: $URL" +[[ "$DRY_RUN" == 1 ]] && echo "DRY RUN: no POSTs will be made." + +# --- liveness --- +if ! curl -fsS "${META_CURL_OPTS[@]}" "$URL/agentmemory/livez" >/dev/null; then + echo "server not reachable at $URL (try: npx @agentmemory/agentmemory)" >&2 + exit 1 +fi + +# --- collect session ids --- +sessions_json="$(curl -fsS "${META_CURL_OPTS[@]}" "$URL/agentmemory/sessions")" +filter='.sessions[] | select(.status=="completed")' +if [[ -n "$ONLY_TAG" ]]; then + filter+=" | select((.tags // []) | index(\"$ONLY_TAG\"))" +fi +if [[ "$SKIP_AGENTS" == 1 ]]; then + filter+=' | select((.project // "") | startswith("agent-") | not)' +fi +if [[ -n "$PROJECT_PATTERN" ]]; then + # jq's test() applies a regex against the project string. + filter+=" | select((.project // \"\") | test(\"$PROJECT_PATTERN\"))" +fi +if [[ "$MAX_OBS" -gt 0 ]]; then + filter+=" | select((.observationCount // 0) <= $MAX_OBS)" +fi +filter+=' | "\(.id)\t\(.observationCount // 0)\t\(.project // "")"' + +rows=() +while IFS= read -r line; do + rows+=("$line") +done < <(echo "$sessions_json" | jq -r "$filter") +total="${#rows[@]}" + +if [[ "$total" -eq 0 ]]; then + echo "no sessions matched (tag='$ONLY_TAG'); nothing to do." + exit 0 +fi + +if [[ "$LIMIT" -gt 0 && "$LIMIT" -lt "$total" ]]; then + rows=("${rows[@]:0:$LIMIT}") +fi + +echo "matched $total session(s); will process ${#rows[@]}." +total_obs=0 +for row in "${rows[@]}"; do + obs="$(cut -f2 <<<"$row")" + total_obs=$(( total_obs + obs )) +done +est_in=$(( total_obs * TOKENS_PER_OBS + ${#rows[@]} * 500 )) +est_out=$(( ${#rows[@]} * TOKENS_OUT_PER_SESSION )) +est_cost="$(awk -v i="$est_in" -v o="$est_out" -v ci="$COST_IN_PER_1M" -v co="$COST_OUT_PER_1M" \ + 'BEGIN { printf "%.2f", (i*ci + o*co) / 1000000 }')" + +echo "≈ ${#rows[@]} summarize LLM calls (one per session, covering $total_obs observations)" +printf '≈ %d input tok + %d output tok → $%s (rates: in=$%s/1M out=$%s/1M, %s tok/obs)\n' \ + "$est_in" "$est_out" "$est_cost" "$COST_IN_PER_1M" "$COST_OUT_PER_1M" "$TOKENS_PER_OBS" +echo + +if [[ "$DRY_RUN" == 1 ]]; then + printf '%-40s %10s %s\n' "session" "obs" "project" + for row in "${rows[@]}"; do + id="$(cut -f1 <<<"$row")" + obs="$(cut -f2 <<<"$row")" + proj="$(cut -f3 <<<"$row")" + printf '%-40s %10s %s\n' "$id" "$obs" "$proj" + done + echo + echo "(dry run) next steps if you re-run without --dry-run:" + echo " for each session above: POST $URL/agentmemory/summarize {sessionId}" + if [[ "$SKIP_CONSOLIDATE" == 0 ]]; then + echo " then: POST $URL/agentmemory/consolidate-pipeline {}" + fi + exit 0 +fi + +# --- summarize loop --- +if [[ "$DEBUG_ON_ERROR" == 1 ]]; then + mkdir -p "$DEBUG_DIR" + echo "debug mode: failed calls will dump to $DEBUG_DIR/" + echo +fi + +dump_failure() { + local id="$1" obs="$2" resp="$3" + # Replace anything outside [A-Za-z0-9._-] with `_` before joining with + # DEBUG_DIR. Session IDs from the API are UUIDs in practice, but the + # server doesn't enforce that — a hostile or buggy id containing `/` or + # `..` would otherwise escape the debug directory. + local safe_id + safe_id="$(printf '%s' "$id" | tr -c 'A-Za-z0-9._-' '_')" + local file="$DEBUG_DIR/${safe_id}.json" + # Pull the raw observations (what would have gone into the prompt) so the + # operator can reconstruct the upstream payload locally. We also compute + # narrative size stats so size-related rejections are immediately visible. + # Stream observations through stdin (avoids exec-arg overflow on + # multi-thousand-obs sessions — macOS argv ceiling is ~256k). + # `--get --data-urlencode` percent-encodes the session id so special + # characters can't corrupt the query string. + curl -fsS "${META_CURL_OPTS[@]}" --get \ + --data-urlencode "sessionId=$id" \ + "$URL/agentmemory/observations" \ + | jq \ + --arg id "$id" \ + --argjson obsCount "$obs" \ + --arg url "$URL/agentmemory/summarize" \ + --argjson response "$resp" \ + '. as $root + | .observations as $obs + | { + sessionId: $id, + observationCount: $obsCount, + request: { url: $url, method: "POST", body: { sessionId: $id } }, + response: $response, + observations: $obs, + stats: { + totalNarrativeBytes: ($obs | map(.narrative // "" | length) | add // 0), + maxNarrativeBytes: ($obs | map(.narrative // "" | length) | max // 0), + titleHistogram: ($obs | group_by(.title) | map({title: .[0].title, count: length}) | sort_by(-.count)) + } + }' >"$file" + echo " → $file" +} + +ok=0; skipped=0; failed=0 +i=0 +for row in "${rows[@]}"; do + i=$(( i + 1 )) + id="$(cut -f1 <<<"$row")" + obs="$(cut -f2 <<<"$row")" + + body="$(jq -nc --arg id "$id" '{sessionId:$id}')" + resp="$(curl -sS "${WORK_CURL_OPTS[@]}" -X POST "$URL/agentmemory/summarize" \ + -H 'content-type: application/json' --data "$body" || echo '{"success":false,"error":"curl_failed"}')" + # iii's HTTP layer occasionally returns non-JSON (HTML 5xx, empty body + # on timeout, etc.). Validate before parsing so `set -e` doesn't abort + # the whole backfill loop on a single bad response. + if jq -e . >/dev/null 2>&1 <<<"$resp"; then + status="$(jq -r '.success // false' <<<"$resp")" + err="$(jq -r '.error // ""' <<<"$resp")" + title="$(jq -r '.summary.title // ""' <<<"$resp")" + else + status="false" + err="invalid_json_response" + title="" + fi + + if [[ "$status" == "true" ]]; then + ok=$(( ok + 1 )) + printf '[%3d/%3d] OK %s obs=%-5s %s\n' "$i" "${#rows[@]}" "$id" "$obs" "$title" + elif [[ "$err" == "no_observations" || "$err" == "no_provider" ]]; then + skipped=$(( skipped + 1 )) + printf '[%3d/%3d] SKIP %s obs=%-5s %s\n' "$i" "${#rows[@]}" "$id" "$obs" "$err" + else + failed=$(( failed + 1 )) + printf '[%3d/%3d] FAIL %s obs=%-5s %s\n' "$i" "${#rows[@]}" "$id" "$obs" "$err" + [[ "$DEBUG_ON_ERROR" == 1 ]] && dump_failure "$id" "$obs" "$resp" + fi +done + +echo +echo "summarize: ok=$ok skipped=$skipped failed=$failed" + +# --- consolidate --- +if [[ "$SKIP_CONSOLIDATE" == 1 ]]; then + echo "skipping consolidate-pipeline (--skip-consolidate)" + exit 0 +fi + +if [[ "$ok" -eq 0 ]]; then + echo "no summaries produced; skipping consolidate-pipeline." + exit 0 +fi + +echo +echo "running consolidate-pipeline …" +resp="$(curl -sS "${WORK_CURL_OPTS[@]}" -X POST "$URL/agentmemory/consolidate-pipeline" \ + -H 'content-type: application/json' --data '{}' || echo '{"success":false,"error":"curl_failed"}')" +if jq -e . >/dev/null 2>&1 <<<"$resp"; then + echo "$resp" | jq . +else + echo "consolidate-pipeline returned non-JSON (likely a timeout or upstream error):" + printf '%s\n' "$resp" | head -c 500 + echo +fi diff --git a/src/functions/summarize.ts b/src/functions/summarize.ts index 140e0e121c..80b29a099b 100644 --- a/src/functions/summarize.ts +++ b/src/functions/summarize.ts @@ -7,7 +7,12 @@ import type { } from "../types.js"; import { KV } from "../state/schema.js"; import { StateKV } from "../state/kv.js"; -import { SUMMARY_SYSTEM, buildSummaryPrompt } from "../prompts/summary.js"; +import { + SUMMARY_SYSTEM, + buildSummaryPrompt, + REDUCE_SYSTEM, + buildReducePrompt, +} from "../prompts/summary.js"; import { getXmlTag, getXmlChildren } from "../prompts/xml.js"; import { SummaryOutputSchema } from "../eval/schemas.js"; import { validateOutput } from "../eval/validator.js"; @@ -16,6 +21,169 @@ import type { MetricsStore } from "../eval/metrics-store.js"; import { safeAudit } from "./audit.js"; import { logger } from "../logger.js"; +// Per-chunk observation budget when a session is too large to fit in one +// LLM call. Default ≈ 50k input tokens per chunk at ~110 tok/obs — fits +// comfortably in 128k-window models. Override via SUMMARIZE_CHUNK_SIZE. +const CHUNK_SIZE_DEFAULT = 400; +// Concurrent in-flight chunk calls. 6 keeps a 100-chunk session under +// iii's 180s function-invocation timeout at ~8s/call while staying +// inside generous-but-not-unlimited provider rate limits (well below +// OpenAI free tier's 500 RPM). High-throughput providers +// (Novita / DeepInfra / DeepSeek) typically allow 100+ concurrent — set +// SUMMARIZE_CHUNK_CONCURRENCY higher to cover ~1000+ chunk sessions. +const CHUNK_CONCURRENCY_DEFAULT = 6; +// Bail on the merged summary if more than this fraction of chunks fail +// to parse — a half-blind narrative is worse than a clean error. +const MAX_SKIP_RATIO = 0.5; + +function getChunkSize(): number { + const raw = process.env.SUMMARIZE_CHUNK_SIZE; + if (!raw) return CHUNK_SIZE_DEFAULT; + const n = parseInt(raw, 10); + return Number.isFinite(n) && n > 0 ? n : CHUNK_SIZE_DEFAULT; +} + +function getChunkConcurrency(): number { + const raw = process.env.SUMMARIZE_CHUNK_CONCURRENCY; + if (!raw) return CHUNK_CONCURRENCY_DEFAULT; + const n = parseInt(raw, 10); + return Number.isFinite(n) && n > 0 ? n : CHUNK_CONCURRENCY_DEFAULT; +} + +// One chunk call with retry-once. Returns null when both attempts fail — +// whether by parse failure, provider 4xx (content rejected by upstream +// filters), or transient network/5xx errors that didn't recover on retry. +// All failure modes are equivalent at this layer: the chunk is unusable, +// skip it and let the caller decide via the skip-ratio bailout whether +// the overall summary is still trustworthy. Errors that affect every +// chunk (auth, model down) will trip the bailout naturally. +async function summarizeChunkWithRetry( + provider: MemoryProvider, + chunk: CompressedObservation[], + sessionId: string, + project: string, + idx: number, + total: number, +): Promise { + for (let attempt = 1; attempt <= 2; attempt++) { + try { + const xml = await provider.summarize( + SUMMARY_SYSTEM, + buildSummaryPrompt(chunk), + ); + const parsed = parseSummaryXml(xml, sessionId, project, chunk.length); + if (parsed) return parsed; + logger.warn("Summarize chunk parse failed", { + sessionId, + chunk: `${idx + 1}/${total}`, + attempt, + }); + } catch (err) { + logger.warn("Summarize chunk LLM call failed", { + sessionId, + chunk: `${idx + 1}/${total}`, + attempt, + error: err instanceof Error ? err.message : String(err), + }); + } + } + return null; +} + +// Returns the final summary XML string. For sessions ≤ chunk size, this is +// a single LLM call (legacy behavior). For larger sessions, observations +// are split into chunks processed in parallel batches, each chunk retried +// once on parse failure, persistently-bad chunks skipped, and remaining +// partials merged via a reduce call. +async function produceSummaryXml( + provider: MemoryProvider, + compressed: CompressedObservation[], + sessionId: string, + project: string, +): Promise<{ + response: string; + mode: "single" | "chunked"; + chunks: number; + skipped?: number; +}> { + const chunkSize = getChunkSize(); + if (compressed.length <= chunkSize) { + const response = await provider.summarize( + SUMMARY_SYSTEM, + buildSummaryPrompt(compressed), + ); + return { response, mode: "single", chunks: 1 }; + } + + const chunks: CompressedObservation[][] = []; + for (let i = 0; i < compressed.length; i += chunkSize) { + chunks.push(compressed.slice(i, i + chunkSize)); + } + const concurrency = getChunkConcurrency(); + logger.info("Summarize chunking session", { + sessionId, + chunks: chunks.length, + chunkSize, + concurrency, + totalObservations: compressed.length, + }); + + // Sparse array preserves chunk → index mapping after parallel resolution, + // so the reduce step sees partials in chronological order even when some + // were skipped. + const partialByIdx: Array = new Array(chunks.length).fill(null); + for (let batchStart = 0; batchStart < chunks.length; batchStart += concurrency) { + const batch = chunks.slice(batchStart, batchStart + concurrency); + await Promise.all( + batch.map(async (chunk, j) => { + const idx = batchStart + j; + partialByIdx[idx] = await summarizeChunkWithRetry( + provider, + chunk, + sessionId, + project, + idx, + chunks.length, + ); + }), + ); + } + + const skipped = partialByIdx.filter((p) => p === null).length; + const partials = partialByIdx.filter((p): p is SessionSummary => p !== null); + + if (skipped > Math.floor(chunks.length * MAX_SKIP_RATIO)) { + throw new Error( + `too_many_chunks_skipped: ${skipped}/${chunks.length} chunks failed to parse after retry`, + ); + } + if (skipped > 0) { + logger.warn("Summarize chunks partially skipped", { + sessionId, + skipped, + total: chunks.length, + }); + } + + const reduceInput = partials.map((p) => { + const originalIdx = partialByIdx.indexOf(p); + return { + title: p.title, + narrative: p.narrative, + keyDecisions: p.keyDecisions, + filesModified: p.filesModified, + concepts: p.concepts, + obsRangeStart: originalIdx * chunkSize + 1, + obsRangeEnd: Math.min((originalIdx + 1) * chunkSize, compressed.length), + }; + }); + const response = await provider.summarize( + REDUCE_SYSTEM, + buildReducePrompt(reduceInput), + ); + return { response, mode: "chunked", chunks: chunks.length, skipped }; +} + function parseSummaryXml( xml: string, sessionId: string, @@ -85,8 +253,12 @@ export function registerSummarizeFunction( } try { - const prompt = buildSummaryPrompt(compressed); - const response = await provider.summarize(SUMMARY_SYSTEM, prompt); + const { response, mode, chunks } = await produceSummaryXml( + provider, + compressed, + sessionId, + session.project, + ); if (!response || !response.trim()) { const latencyMs = Date.now() - startMs; if (metricsStore) { @@ -95,8 +267,8 @@ export function registerSummarizeFunction( logger.warn("Empty provider response on summarize", { sessionId, provider: provider.name, - promptBytes: prompt.length, - systemBytes: SUMMARY_SYSTEM.length, + mode, + chunks, observationCount: compressed.length, }); return { success: false, error: "empty_provider_response" }; diff --git a/src/prompts/summary.ts b/src/prompts/summary.ts index f01b28b8ee..bd04021277 100644 --- a/src/prompts/summary.ts +++ b/src/prompts/summary.ts @@ -36,3 +36,52 @@ export function buildSummaryPrompt(observations: Array<{ }) return `Session observations (${observations.length} total):\n\n${lines.join('\n\n---\n\n')}` } + +export const REDUCE_SYSTEM = `You are merging multiple partial summaries of the SAME coding session into one final session summary. The partials are chronological chunks of one continuous session — not separate sessions. + +Output EXACTLY this XML format with no additional text: + + + Short session title (max 100 chars) + 3-5 sentence narrative covering the whole session + + Key technical decision made + + + path/to/modified/file + + + key concept from session + + + +Rules: +- Synthesize a single narrative that reflects the whole arc, not a chunk-by-chunk recap +- Preserve every distinct decision across chunks +- Union (deduplicate) all files and concepts +- Title should capture the session's overall outcome` + +export function buildReducePrompt(partials: Array<{ + title: string + narrative: string + keyDecisions: string[] + filesModified: string[] + concepts: string[] + obsRangeStart: number + obsRangeEnd: number +}>): string { + const sections = partials.map((p, i) => { + const decisions = p.keyDecisions.map((d) => ` - ${d}`).join('\n') + const files = p.filesModified.map((f) => ` - ${f}`).join('\n') + const concepts = p.concepts.join(', ') + return `[Chunk ${i + 1} of ${partials.length} — obs ${p.obsRangeStart}-${p.obsRangeEnd}] +Title: ${p.title} +Narrative: ${p.narrative} +Decisions: +${decisions} +Files: +${files} +Concepts: ${concepts}` + }) + return `Partial summaries (${partials.length} chunks of one session, chronological):\n\n${sections.join('\n\n---\n\n')}` +} diff --git a/test/summarize.test.ts b/test/summarize.test.ts new file mode 100644 index 0000000000..03aa1926a0 --- /dev/null +++ b/test/summarize.test.ts @@ -0,0 +1,417 @@ +import { describe, it, expect, beforeEach, afterEach, vi } from "vitest"; + +vi.mock("../src/logger.js", () => ({ + logger: { info: vi.fn(), warn: vi.fn(), error: vi.fn() }, +})); + +vi.mock("../src/state/schema.js", () => ({ + KV: { + sessions: "sessions", + summaries: "summaries", + observations: (sessionId: string) => `obs:${sessionId}`, + audit: "audit", + }, +})); + +vi.mock("../src/eval/schemas.js", () => ({ + SummaryOutputSchema: {}, +})); + +vi.mock("../src/eval/validator.js", () => ({ + validateOutput: () => ({ valid: true, result: { errors: [] } }), +})); + +vi.mock("../src/eval/quality.js", () => ({ + scoreSummary: () => 100, +})); + +vi.mock("../src/functions/audit.js", () => ({ + safeAudit: vi.fn(), +})); + +import { registerSummarizeFunction } from "../src/functions/summarize.js"; +import type { + CompressedObservation, + Session, + MemoryProvider, +} from "../src/types.js"; + +function mockKV() { + const store = new Map>(); + return { + store, + get: async (scope: string, key: string): Promise => + (store.get(scope)?.get(key) as T) ?? null, + set: async (scope: string, key: string, data: T): Promise => { + if (!store.has(scope)) store.set(scope, new Map()); + store.get(scope)!.set(key, data); + return data; + }, + delete: async (scope: string, key: string): Promise => { + store.get(scope)?.delete(key); + }, + list: async (scope: string): Promise => { + const entries = store.get(scope); + return entries ? (Array.from(entries.values()) as T[]) : []; + }, + }; +} + +function mockSdk() { + const functions = new Map(); + return { + functions, + registerFunction: (id: string, handler: Function) => { + functions.set(id, handler); + }, + registerTrigger: () => {}, + trigger: async () => ({}), + }; +} + +function makeObs(i: number, sessionId: string): CompressedObservation { + return { + id: `obs_${i}`, + sessionId, + timestamp: new Date().toISOString(), + type: "conversation", + title: `obs ${i}`, + facts: [`fact ${i}`], + narrative: `narrative for obs ${i}`, + concepts: [], + files: [`src/file_${i}.ts`], + importance: 5, + }; +} + +function makeProvider(responses: string[]): MemoryProvider & { + calls: Array<{ system: string; user: string }>; +} { + const calls: Array<{ system: string; user: string }> = []; + let i = 0; + return { + name: "test", + calls, + compress: async () => "", + summarize: async (system: string, user: string) => { + calls.push({ system, user }); + const r = responses[i] ?? responses[responses.length - 1]; + i += 1; + return r; + }, + }; +} + +function summaryXml(opts: { + title: string; + narrative?: string; + decisions?: string[]; + files?: string[]; + concepts?: string[]; +}): string { + const d = (opts.decisions ?? []).map((x) => `${x}`).join(""); + const f = (opts.files ?? []).map((x) => `${x}`).join(""); + const c = (opts.concepts ?? []).map((x) => `${x}`).join(""); + return ` +${opts.title} +${opts.narrative ?? "narrative"} +${d} +${f} +${c} +`; +} + +async function setupHandler(opts: { + sessionId: string; + obsCount: number; + provider: MemoryProvider; +}) { + const sdk = mockSdk(); + const kv = mockKV(); + const session: Session = { + id: opts.sessionId, + project: "test-project", + cwd: "/tmp", + startedAt: new Date().toISOString(), + status: "completed", + observationCount: opts.obsCount, + }; + await kv.set("sessions", opts.sessionId, session); + for (let i = 0; i < opts.obsCount; i++) { + const o = makeObs(i, opts.sessionId); + await kv.set(`obs:${opts.sessionId}`, o.id, o); + } + registerSummarizeFunction(sdk as any, kv as any, opts.provider); + const handler = sdk.functions.get("mem::summarize")!; + return { handler, kv }; +} + +describe("mem::summarize chunking", () => { + const ORIGINAL_ENV = { ...process.env }; + + beforeEach(() => { + delete process.env.SUMMARIZE_CHUNK_SIZE; + delete process.env.SUMMARIZE_CHUNK_CONCURRENCY; + }); + + afterEach(() => { + process.env = { ...ORIGINAL_ENV }; + }); + + it("small session takes the single-call path (no chunking, no reduce)", async () => { + const provider = makeProvider([ + summaryXml({ + title: "Small session", + decisions: ["decision A"], + files: ["src/a.ts"], + concepts: ["concept-a"], + }), + ]); + const { handler, kv } = await setupHandler({ + sessionId: "ses_small", + obsCount: 10, + provider, + }); + + const result: any = await handler({ sessionId: "ses_small" }); + + expect(result.success).toBe(true); + expect(provider.calls).toHaveLength(1); + expect(provider.calls[0].user).toContain("Session observations (10 total)"); + const stored: any = await kv.get("summaries", "ses_small"); + expect(stored?.title).toBe("Small session"); + }); + + it("large session map-reduces: N chunk calls + 1 reduce call", async () => { + process.env.SUMMARIZE_CHUNK_SIZE = "100"; + process.env.SUMMARIZE_CHUNK_CONCURRENCY = "1"; // serial keeps call ordering deterministic + const provider = makeProvider([ + summaryXml({ title: "Chunk 1", decisions: ["dA"], files: ["src/a.ts"], concepts: ["ca"] }), + summaryXml({ title: "Chunk 2", decisions: ["dB"], files: ["src/b.ts"], concepts: ["cb"] }), + summaryXml({ title: "Chunk 3", decisions: ["dC"], files: ["src/c.ts"], concepts: ["cc"] }), + summaryXml({ + title: "Merged", + decisions: ["dA", "dB", "dC"], + files: ["src/a.ts", "src/b.ts", "src/c.ts"], + concepts: ["ca", "cb", "cc"], + }), + ]); + const { handler, kv } = await setupHandler({ + sessionId: "ses_large", + obsCount: 250, + provider, + }); + + const result: any = await handler({ sessionId: "ses_large" }); + + expect(result.success).toBe(true); + expect(provider.calls).toHaveLength(4); + // First three are chunk calls (use the summary system prompt). + expect(provider.calls[0].system).toContain("session summarizer"); + expect(provider.calls[2].system).toContain("session summarizer"); + // Last is the reduce call (uses the merge system prompt). + expect(provider.calls[3].system).toContain("merging multiple partial summaries"); + expect(provider.calls[3].user).toContain("Chunk 1 of 3"); + expect(provider.calls[3].user).toContain("Chunk 3 of 3"); + + const stored: any = await kv.get("summaries", "ses_large"); + expect(stored?.title).toBe("Merged"); + // observationCount on the persisted summary should reflect the full session, + // not just the final chunk. + expect(stored?.observationCount).toBe(250); + expect(stored?.keyDecisions).toEqual(["dA", "dB", "dC"]); + }); + + it("SUMMARIZE_CHUNK_SIZE env override is respected", async () => { + process.env.SUMMARIZE_CHUNK_SIZE = "50"; + process.env.SUMMARIZE_CHUNK_CONCURRENCY = "1"; + const provider = makeProvider([ + summaryXml({ title: "chunk" }), + summaryXml({ title: "chunk" }), + summaryXml({ title: "chunk" }), + summaryXml({ title: "chunk" }), + summaryXml({ title: "merged" }), + ]); + const { handler } = await setupHandler({ + sessionId: "ses_env", + obsCount: 175, + provider, + }); + + const result: any = await handler({ sessionId: "ses_env" }); + + expect(result.success).toBe(true); + // 175 obs ÷ 50 = 4 chunks (last chunk has 25) + 1 reduce = 5 calls. + expect(provider.calls).toHaveLength(5); + }); + + it("flaky chunk: parse fails once, retried, then succeeds — no skip", async () => { + process.env.SUMMARIZE_CHUNK_SIZE = "100"; + process.env.SUMMARIZE_CHUNK_CONCURRENCY = "1"; + const provider = makeProvider([ + summaryXml({ title: "ok1" }), + "", // chunk 2 attempt 1: parse-fail + summaryXml({ title: "ok2" }), // chunk 2 attempt 2 (retry): success + summaryXml({ title: "ok3" }), + summaryXml({ title: "merged" }), + ]); + const { handler, kv } = await setupHandler({ + sessionId: "ses_flaky", + obsCount: 250, + provider, + }); + + const result: any = await handler({ sessionId: "ses_flaky" }); + + expect(result.success).toBe(true); + // 3 chunks × 1 attempt + 1 retry on chunk 2 + 1 reduce = 5 calls. + expect(provider.calls).toHaveLength(5); + const stored: any = await kv.get("summaries", "ses_flaky"); + expect(stored?.title).toBe("merged"); + }); + + it("persistently-broken chunk is skipped, reduce still runs on remaining partials", async () => { + process.env.SUMMARIZE_CHUNK_SIZE = "100"; + process.env.SUMMARIZE_CHUNK_CONCURRENCY = "1"; + const provider = makeProvider([ + summaryXml({ title: "ok1" }), + "", "", // chunk 2: both attempts parse-fail + summaryXml({ title: "ok3" }), + summaryXml({ title: "merged-with-skip" }), + ]); + const { handler, kv } = await setupHandler({ + sessionId: "ses_skip", + obsCount: 250, + provider, + }); + + const result: any = await handler({ sessionId: "ses_skip" }); + + expect(result.success).toBe(true); + // 1 ok + (1 + 1 retry skip) + 1 ok + 1 reduce = 5 calls. + expect(provider.calls).toHaveLength(5); + // Reduce input should mention only 2 of 3 chunks (chunk 2 skipped) — + // but the chunk indices in the reduce labels should reflect chunk 1 and 3, + // preserving chronological boundaries. + const reduceCall = provider.calls[4]; + expect(reduceCall.user).toContain("Chunk 1 of 2"); + expect(reduceCall.user).toContain("Chunk 2 of 2"); + expect(reduceCall.user).toContain("obs 1-100"); // first surviving chunk + expect(reduceCall.user).toContain("obs 201-250"); // third surviving chunk (was idx 2, range 201-250) + const stored: any = await kv.get("summaries", "ses_skip"); + expect(stored?.title).toBe("merged-with-skip"); + }); + + it("too many skipped chunks bails out with a clear error", async () => { + process.env.SUMMARIZE_CHUNK_SIZE = "100"; + process.env.SUMMARIZE_CHUNK_CONCURRENCY = "1"; + // 3 chunks, 2 fully broken → >50% skipped → bail. + const provider = makeProvider([ + summaryXml({ title: "ok1" }), + "", "", + "", "", + ]); + const { handler } = await setupHandler({ + sessionId: "ses_too_broken", + obsCount: 250, + provider, + }); + + const result: any = await handler({ sessionId: "ses_too_broken" }); + + expect(result.success).toBe(false); + expect(result.error).toMatch(/too_many_chunks_skipped: 2\/3/); + }); + + it("provider error on one chunk after retry is skipped, not propagated", async () => { + process.env.SUMMARIZE_CHUNK_SIZE = "100"; + process.env.SUMMARIZE_CHUNK_CONCURRENCY = "1"; + let i = 0; + const provider: MemoryProvider & { calls: any[] } = { + name: "test", + calls: [], + compress: async () => "", + summarize: async (system: string, user: string) => { + (provider as any).calls.push({ system, user }); + i += 1; + if (i === 1) return summaryXml({ title: "ok1" }); + // chunk 2: both attempts throw (e.g. provider 400) + if (i === 2 || i === 3) throw new Error("OpenAI API error (400): content rejected"); + if (i === 4) return summaryXml({ title: "ok3" }); + return summaryXml({ title: "merged-with-skip" }); + }, + }; + const { handler, kv } = await setupHandler({ + sessionId: "ses_net", + obsCount: 250, + provider, + }); + + const result: any = await handler({ sessionId: "ses_net" }); + + expect(result.success).toBe(true); + // 1 ok + 2 fail + 1 ok + 1 reduce = 5 calls. + expect((provider as any).calls.length).toBe(5); + const stored: any = await kv.get("summaries", "ses_net"); + expect(stored?.title).toBe("merged-with-skip"); + }); + + it("every chunk failing on provider error trips too_many_chunks_skipped", async () => { + process.env.SUMMARIZE_CHUNK_SIZE = "100"; + process.env.SUMMARIZE_CHUNK_CONCURRENCY = "1"; + // 3 chunks, all chunk calls throw → 3/3 skipped → bail. + const provider: MemoryProvider & { calls: any[] } = { + name: "test", + calls: [], + compress: async () => "", + summarize: async (system: string, user: string) => { + (provider as any).calls.push({ system, user }); + throw new Error("OpenAI API error (400): invalid request"); + }, + }; + const { handler } = await setupHandler({ + sessionId: "ses_all_400", + obsCount: 250, + provider, + }); + + const result: any = await handler({ sessionId: "ses_all_400" }); + + expect(result.success).toBe(false); + expect(result.error).toMatch(/too_many_chunks_skipped: 3\/3/); + }); + + it("chunks run in parallel batches according to SUMMARIZE_CHUNK_CONCURRENCY", async () => { + process.env.SUMMARIZE_CHUNK_SIZE = "100"; + process.env.SUMMARIZE_CHUNK_CONCURRENCY = "2"; + let inflight = 0; + let maxInflight = 0; + const provider: MemoryProvider & { calls: any[] } = { + name: "test", + calls: [], + compress: async () => "", + summarize: async (system: string, user: string) => { + (provider as any).calls.push({ system, user }); + inflight += 1; + maxInflight = Math.max(maxInflight, inflight); + // Yield to event loop so siblings can also enter before we resolve. + await new Promise((r) => setTimeout(r, 5)); + inflight -= 1; + if (system.includes("merging")) return summaryXml({ title: "merged" }); + return summaryXml({ title: "ok" }); + }, + }; + const { handler } = await setupHandler({ + sessionId: "ses_par", + obsCount: 400, // 4 chunks at chunkSize=100 + provider, + }); + + const result: any = await handler({ sessionId: "ses_par" }); + + expect(result.success).toBe(true); + // 4 chunks at concurrency 2 → max 2 in flight at once during the chunk phase. + // Reduce is a single call so doesn't bump it. + expect(maxInflight).toBe(2); + }); +});