feat: wire RetainOrchestrator + dedup-safe emit() writes by Gradata · Pull Request #98 · Gradata/gradata

Gradata · 2026-04-17T02:12:55Z

Summary

Closes the event-persistence dup/drift risk: emit() and RetainOrchestrator now share a common dedup identity and the orphaned orchestrator class is reachable from code.

What changed

UNIQUE index on events(ts, type, source) in _ensure_table() — same key the orchestrator's cursor already used, now enforced by the storage layer.
emit() uses INSERT OR IGNORE with id backfill: on a dedup collision the existing row's id is returned so callers keep working.
Module singleton get_retain_orchestrator(brain_dir) + flush_retain(brain_dir) for batch callers. Until now the orchestrator class had zero production callers.
session_close calls flush_retain() as a last-chance safety flush.

Why

RetainOrchestrator was built, tested, then left unwired — a classic "finished upgrade, never deployed" orphan. Before this PR, emit() could double-write on partial-failure retry (JSONL ok, SQLite fail → next retry writes JSONL again). The UNIQUE index + INSERT OR IGNORE make that idempotent.

Verification

161 event-related tests pass locally.
Live smoke test: identical (ts,type,source) INSERT OR IGNORE → 1 row; RetainOrchestrator batch queue of 3 events (1 duplicate) → 2 rows via Phase 1 JSONL dedup + Phase 2 UNIQUE.

Test plan

Unit: existing event tests still pass
Smoke: dup INSERT collapses to 1
Smoke: batch flush respects both dedup layers
Check downstream callers that rely on event["id"] still get valid ids after INSERT OR IGNORE collision

Generated with Gradata

greptile-apps

Gradata has reached the 50-review limit for trial accounts. To continue receiving code reviews, upgrade your plan.

coderabbitai · 2026-04-17T02:13:08Z

Caution

Review failed

Pull request was closed or merged during review

📝 Walkthrough

Deduplication enforcement: add UNIQUE index on events(ts, type, source) to enforce deduplication at the storage layer
Idempotent writes: emit() now uses INSERT OR IGNORE to avoid duplicate inserts on retries/partial failures
ID backfill on collision: when INSERT is ignored, emit() queries the existing row and returns its id so callers keep a valid event["id"]
New public API: get_retain_orchestrator(brain_dir) — returns/creates a cached RetainOrchestrator singleton
New public API: flush_retain(brain_dir) — flushes retained events (returns {"written", "errors", "phases"} or zero-summary)
Session cleanup: session_close() now calls flush_retain() as a last-chance safety flush for pending batched events
Risk mitigation: closes a duplicate/drift window where JSONL writes + SQLite failures could lead to duplicate rows; UNIQUE + INSERT OR IGNORE make writes idempotent

Walkthrough

Adds database-level deduplication by creating a UNIQUE index on (ts, type, source), makes emit() use INSERT OR IGNORE with id backfill on conflicts, introduces module-level RetainOrchestrator helpers (get_retain_orchestrator, flush_retain), and flushes queued events at session close.

Changes

Cohort / File(s)	Summary
Event deduplication & batching `src/gradata/_events.py`	Create `UNIQUE` index on `(ts, type, source)` (suppressed OperationalError); change `emit()` to `INSERT OR IGNORE` and backfill `event["id"]` from existing row when insert is ignored; add `_ORCHESTRATORS` cache and public helpers `get_retain_orchestrator(brain_dir)` and `flush_retain(brain_dir)`.
Session-close flush hook `src/gradata/hooks/session_close.py`	Add private `_flush_retain_queue(brain_dir)` that imports/calls `flush_retain()` and logs if events were written; call this helper from `main()` after tree consolidation to persist any pending batched events.

Sequence Diagram

sequenceDiagram
    participant Client as Client Code
    participant Emit as emit()
    participant DB as SQLite DB
    participant Orch as RetainOrchestrator
    participant Hook as Session Close Hook

    Client->>Emit: emit(event)
    Emit->>DB: INSERT OR IGNORE (ts,type,source,...)
    alt inserted
        DB-->>Emit: rowcount=1
        Emit->>Emit: set event["id"]=lastrowid
    else ignored (duplicate)
        DB-->>Emit: rowcount=0
        Emit->>DB: SELECT id WHERE (ts,type,source)
        DB-->>Emit: return id
        Emit->>Emit: set event["id"]=existing id
    end
    Emit->>Orch: queue event
    Note over Orch: events batched in memory
    Client->>Hook: session closes
    Hook->>Hook: _flush_retain_queue(brain_dir)
    Hook->>Orch: flush_retain(brain_dir)
    Orch->>DB: persist batched events (INSERT OR IGNORE)
    DB-->>Orch: confirmation
    Orch-->>Hook: return {written, errors, phases}

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Suggested labels

bug

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 60.42% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately captures the main change: implementing RetainOrchestrator integration and making emit() writes deduplication-safe through UNIQUE constraints and INSERT OR IGNORE logic.
Description check	✅ Passed	The description is directly related to the changeset, providing clear context on the deduplication strategy, UNIQUE index addition, INSERT OR IGNORE implementation, new helper functions, and session_close integration.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

Create stacked PR
Commit on current branch

🧪 Generate unit tests (beta)

Create PR with unit tests
Commit unit tests in branch feat/wire-retain-orchestrator

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

coderabbitai

Actionable comments posted: 5

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (3)

src/gradata/enhancements/meta_rules.py (1)

621-636: ⚠️ Potential issue | 🟠 Major

Sanitize category too; the prompt is still partially raw.

The descriptions are filtered, but Line 629 still embeds category directly into the LLM prompt. If category names come from lesson data, this leaves a second prompt-injection path open.

🛡️ Suggested fix

     from gradata.enhancements._sanitize import sanitize_lesson_content
 
+    safe_category = sanitize_lesson_content(category, "llm_prompt")
     safe_descriptions = [sanitize_lesson_content(d, "llm_prompt") for d in descriptions]
     bullet_text = "\n".join(f"- {d}" for d in safe_descriptions)
     prompt = (
-        f'Given these {len(descriptions)} learned rules in the "{category}" category:\n'
+        f'Given these {len(descriptions)} learned rules in the "{safe_category}" category:\n'
         f"{bullet_text}\n\n"
         "Synthesize them into 1-3 high-level actionable directives.\n"

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/gradata/enhancements/meta_rules.py` around lines 621 - 636, The prompt
still injects the raw variable category; sanitize it like the descriptions by
passing category through sanitize_lesson_content (same "llm_prompt" context)
before embedding in the prompt construction so the LLM never receives raw
lesson-derived strings—update the code around safe_descriptions / prompt to
create a sanitized_category and use that in the f-string when building prompt;
keep sanitize_lesson_content as the single sanitizer reference.

src/gradata/hooks/inject_brain_rules.py (1)

212-247: ⚠️ Potential issue | 🟠 Major

This only closes one of the XML injection paths.

Lines 224-247 escape the <brain-rules> content, but the same response later concatenates raw lesson/meta text into <mandatory-directives>, <mandatory-reminder>, and <brain-meta-rules>. A crafted description or principle can still terminate those tags and inject prompt content, so the sanitization needs to cover every emitted XML-like block in this file.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@src/gradata/hooks/inject_brain_rules.py` around lines 212 - 247, The review
points out sanitization is only applied to cluster_lines and individual_lines;
extend XML-escaping everywhere raw lesson/meta text is embedded by calling
sanitize_lesson_content(..., "xml") for every field used to build XML blocks
(e.g., any uses of cluster.summary, cluster.category, r.description, r.category
and also any meta/principle/mandatory text that is concatenated into
<mandatory-directives>, <mandatory-reminder>, and <brain-meta-rules>). Locate
builders that produce cluster_lines, individual_lines and the routines that emit
mandatory-directives/mandatory-reminder/brain-meta-rules and ensure each input
string is passed through sanitize_lesson_content before concatenation so no
unescaped "</...>" or XML injection can occur.

src/gradata/enhancements/llm_synthesizer.py (1)

73-95: ⚠️ Potential issue | 🟠 Major

Sanitize theme before putting it into the prompt.

Only the lesson descriptions go through llm_prompt; Line 89 still inserts theme raw. If that label is derived from lesson/category data, a crafted value can still steer the model and bypass the new output gate.

🛡️ Suggested fix

     from gradata.enhancements._sanitize import sanitize_lesson_content
 
+    safe_theme = sanitize_lesson_content(theme, "llm_prompt")
     bullets = []
     for lesson in lessons[:10]:  # Cap at 10 to limit prompt size
         desc = lesson.description
         if desc:
             safe_desc = sanitize_lesson_content(desc, "llm_prompt")
             bullets.append(f"- {safe_desc}")
@@
     bullet_text = "\n".join(bullets)
     prompt = (
-        f"Given these {len(bullets)} user corrections all related to \"{theme}\":\n"
+        f"Given these {len(bullets)} user corrections all related to \"{safe_theme}\":\n"
         f"{bullet_text}\n\n"
         "Write ONE actionable behavioral principle (1-2 sentences) that captures the pattern.\n"

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/gradata/enhancements/llm_synthesizer.py` around lines 73 - 95, The prompt
currently embeds the raw theme string, which can be used for prompt injection;
sanitize the theme with the same sanitizer before interpolation (e.g., call
sanitize_lesson_content(theme, "llm_prompt") and store as safe_theme) and
replace the direct usage of theme in the prompt construction with safe_theme;
update any references around the bullets/prompt creation in llm_synthesizer.py
(where sanitize_lesson_content is already imported and bullets are built) to use
the sanitized value.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@src/gradata/_events.py`:
- Around line 603-626: Remove the unnecessary string quotes from type
annotations now that from __future__ import annotations is present: change the
annotated types in the module-level _ORCHESTRATORS declaration (dict[str,
RetainOrchestrator]), the get_retain_orchestrator signature and return type
(brain_dir: str | Path -> -> RetainOrchestrator), the flush_retain signature and
return type (brain_dir: str | Path -> -> dict), and the quoted annotation in
RetainOrchestrator.__init__ (replace "Path" / "str" style quotes with unquoted
types); keep the same types, only remove the surrounding quotes so ruff UP037 is
satisfied.
- Around line 138-145: The current branch that handles the INSERT OR IGNORE
fallback can leave event["id"] as None if the SELECT returns no row (race where
another process deleted the row); update the branch in the block that checks
cursor.rowcount != 1 (the INSERT fallback using existing =
conn.execute(...).fetchone()) to detect existing is None and fail fast: either
raise a clear exception (e.g., RuntimeError) or at minimum log an error
including ts, event_type, source and the problematic event reference before
raising, so event["id"] is never silently set to None and downstream code using
event["id"] (e.g., correction_event_ids and failure["correction_event_id"])
won't be broken. Ensure the change touches the same branch that assigns
event["id"] from existing and includes the SELECT parameters (ts, event_type,
source) in the diagnostic message.

In `@src/gradata/enhancements/_sanitize.py`:
- Around line 173-182: The current _neutralize_llm_prompt function only
substitutes the trigger token, leaving the rest of the injected clause in place;
update _neutralize_llm_prompt (and its use of _PROMPT_INJECTION_PATTERNS and
_FILTER_PLACEHOLDER) to redact the whole injected clause/line rather than just
the trigger token by using pattern.sub with a replacer that returns
_FILTER_PLACEHOLDER for the entire match plus any contiguous clause up to the
next sentence terminator or newline (e.g., extend the match span to include
following text until a punctuation like .!? or a newline) or, if patterns
already capture the rest, return _FILTER_PLACEHOLDER for the full capture;
ensure spacing is normalized so outputs like "[FILTERED] DAN" or
"[FILTERED]override" do not leave instruction remnants.
- Around line 219-224: The branch that currently logs a warning and returns raw
text in sanitize_lesson_content should fail closed instead: replace the
logging+return behavior with raising an explicit exception (e.g., ValueError or
RuntimeError) that includes the unexpected context value so callers cannot
accidentally receive unsanitized content; keep or add a logging.error with the
same message if needed, but do not return text when context is unknown.

In `@tests/test_sanitize_lesson_content.py`:
- Around line 189-242: The tests currently only check for "[FILTERED]" presence
which allows dangerous tails to remain; update each failing test (e.g.,
test_ignore_previous_instructions_filtered, test_you_are_now_filtered,
test_system_colon_marker_filtered, test_act_as_filtered,
test_jailbreak_keyword_filtered, test_dan_mode_filtered,
test_rephrased_override_filtered) to also assert that the sanitized output
(sanitize_lesson_content(payload, "llm_prompt")) does not contain the original
malicious remainder—e.g., check sanitized_result.lower() does not include key
substrings like "ignore previous instructions", "you are now", "dan", "system:",
"override all previous rules", "jailbreak", "override your safety", etc., or
assert that sanitized_result.replace("[FILTERED]","").strip() does not contain
any of those tokens—so the tests verify the dangerous tail is removed, not just
that a filter marker was inserted.

---

Outside diff comments:
In `@src/gradata/enhancements/llm_synthesizer.py`:
- Around line 73-95: The prompt currently embeds the raw theme string, which can
be used for prompt injection; sanitize the theme with the same sanitizer before
interpolation (e.g., call sanitize_lesson_content(theme, "llm_prompt") and store
as safe_theme) and replace the direct usage of theme in the prompt construction
with safe_theme; update any references around the bullets/prompt creation in
llm_synthesizer.py (where sanitize_lesson_content is already imported and
bullets are built) to use the sanitized value.

In `@src/gradata/enhancements/meta_rules.py`:
- Around line 621-636: The prompt still injects the raw variable category;
sanitize it like the descriptions by passing category through
sanitize_lesson_content (same "llm_prompt" context) before embedding in the
prompt construction so the LLM never receives raw lesson-derived strings—update
the code around safe_descriptions / prompt to create a sanitized_category and
use that in the f-string when building prompt; keep sanitize_lesson_content as
the single sanitizer reference.

In `@src/gradata/hooks/inject_brain_rules.py`:
- Around line 212-247: The review points out sanitization is only applied to
cluster_lines and individual_lines; extend XML-escaping everywhere raw
lesson/meta text is embedded by calling sanitize_lesson_content(..., "xml") for
every field used to build XML blocks (e.g., any uses of cluster.summary,
cluster.category, r.description, r.category and also any
meta/principle/mandatory text that is concatenated into <mandatory-directives>,
<mandatory-reminder>, and <brain-meta-rules>). Locate builders that produce
cluster_lines, individual_lines and the routines that emit
mandatory-directives/mandatory-reminder/brain-meta-rules and ensure each input
string is passed through sanitize_lesson_content before concatenation so no
unescaped "</...>" or XML injection can occur.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro

Run ID: 59092688-5f68-4147-8253-0994fc5e3891

📥 Commits

Reviewing files that changed from the base of the PR and between ae423a7 and 547da55.

📒 Files selected for processing (8)

src/gradata/_events.py
src/gradata/enhancements/_sanitize.py
src/gradata/enhancements/llm_synthesizer.py
src/gradata/enhancements/meta_rules.py
src/gradata/enhancements/rule_to_hook.py
src/gradata/hooks/inject_brain_rules.py
src/gradata/hooks/session_close.py
tests/test_sanitize_lesson_content.py

📜 Review details

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)

GitHub Check: test (3.12)

🧰 Additional context used

📓 Path-based instructions (3)

src/gradata/**/*.py

⚙️ CodeRabbit configuration file

src/gradata/**/*.py: This is the core SDK. Check for: type safety (from future import annotations required), no print()
statements (use logging), all functions accepting BrainContext where DB access occurs, no hardcoded paths. Severity
scoring must clamp to [0,1]. Confidence values must be in [0.0, 1.0].

Files:

src/gradata/enhancements/llm_synthesizer.py
src/gradata/enhancements/meta_rules.py
src/gradata/hooks/inject_brain_rules.py
src/gradata/hooks/session_close.py
src/gradata/enhancements/rule_to_hook.py
src/gradata/enhancements/_sanitize.py
src/gradata/_events.py

src/gradata/hooks/**

⚙️ CodeRabbit configuration file

src/gradata/hooks/**: JavaScript hooks for Claude Code integration. Check for: no shell injection (no execSync with user
input), temp files must use per-user subdirectory, HTTP calls must have timeouts, errors must be silent (never block
the tool chain).

Files:

src/gradata/hooks/inject_brain_rules.py
src/gradata/hooks/session_close.py

tests/**

⚙️ CodeRabbit configuration file

tests/**: Test files. Verify: no hardcoded paths, assertions check specific values not just truthiness,
parametrized tests preferred for boundary conditions, floating point comparisons use pytest.approx.

Files:

tests/test_sanitize_lesson_content.py

🪛 GitHub Actions: SDK CI

src/gradata/enhancements/llm_synthesizer.py