Skip to content

feat(memory): Lakebase-backed persistent memory for CODA sessions - #18

Closed
dgokeeffe wants to merge 255 commits into
mainfrom
feat/lakebase-memory
Closed

dgokeeffe wants to merge 255 commits into
mainfrom
feat/lakebase-memory

Conversation

@dgokeeffe

@dgokeeffe dgokeeffe commented May 6, 2026 •

Copy link
Copy Markdown
Contributor

Priority

P2 (feature) — Lakebase-backed persistent memory across CODA sessions. Standalone, no merge-order dependency on the other PRs.


Summary

Migrates datasciencemonkey PR #145 to the new repo home, rebased onto databrickslabs/main (5 commits, all clean after conflict resolution).

Adds Lakebase-backed persistent memory to CODA. At session end, a Stop hook extracts 3–8 structured memories from the Claude transcript via the Databricks Claude Haiku serving endpoint, persists them to a coda_memories Postgres table, and regenerates coda_memory.md so the next session opens with the accumulated context already loaded.

Components

  • memory/extractor.py — Stop hook: parses the JSONL transcript, calls Haiku, writes to Lakebase, regenerates the local memory file.
  • memory/store.py — psycopg 3 + psycopg_pool with an OAuthConnection subclass that mints a fresh Lakebase OAuth credential on every pool recycle (45 min, ahead of the 1h token expiry). Follows the canonical Databricks Apps + Lakebase Autoscaling pattern.
  • memory/injector.py — Regenerates ~/.claude/memory/coda_memory.md (or project-scoped variant) from Lakebase.
  • memory/searcher.py — FTS-ranked search CLI used by the memory-recall subagent.
  • memory/hooks/user_prompt_submit.py — Zero-cost nudge so Claude knows to invoke the subagent when historical context would help.
  • setup_memory.py — Per-session schema init, memory-file warm-up, idempotent hook registration.

Schema

coda_memories: content_hash UNIQUE for dedup, content_tsv generated column + GIN index for FTS, nullable vector(1536) + HNSW (kept off by default — semantic search comes later).

Rebase notes

Three conflicts resolved during rebase onto current main:

  • requirements.txt: kept main's fastapi==0.136.0 (security bump), kept branch's distro + docstring-parser (transitives of anthropic SDK pulled in by Lakebase memory's Haiku call).
  • Makefile: union of .PHONY lists — added test (from main) + deploy-e2e/setup-secret/link-resources/clean-secret (from branch).
  • app.py: ordered the post-parallel sequential steps as mlflow → memory → token-sync. setup_memory writes a Stop hook alongside setup_mlflow's, so it must run after; token-sync stays last since it touches CLI configs, not settings.json.

pyproject.toml auto-merged (anthropic + psycopg deps add cleanly). All 4 follow-up commits (defense-in-depth, tests, staleness mark) replayed without conflict.

Follow-ups

  • Embedding ingestion + MMR retrieval (vector column is in place, query side isn't yet).
  • Memory decay / TTL.
  • Multi-tenant scoping (currently single-user).

Closes #14

This pull request and its description were written by Isaac.

datasciencemonkey and others added 30 commits February 22, 2026 08:34
- All 3 setup scripts: use DATABRICKS_GATEWAY_TOKEN when gateway host
  is present, fall back to DATABRICKS_TOKEN on current gateway
- setup_opencode.py: full endpoint paths for OpenCode
  - Anthropic/Gemini: {gateway}/mlflow/v1/chat/completions
  - OpenAI/GPT: {gateway}/openai/v1/responses
- app.yaml: add DATABRICKS_GATEWAY_TOKEN secret

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Platform injects DATABRICKS_HOST automatically. Gateway host is the
primary auth path; fallback uses platform-injected host.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
SDK auto-appends /chat/completions and /responses to baseURL, so
baseURL should be just /mlflow/v1 and /openai/v1, not the full path.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Databricks docs specify @google/gemini-cli@nightly for gateway support.
Fix selectedAuthType from "api-key" to "gemini-api-key" to auto-skip
the auth prompt when GEMINI_API_KEY is set in .env.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Gemini CLI defaults to gemini-2.5-flash-lite which doesn't exist on
Databricks. Set model.name in settings.json to the configured model.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Shares the Claude skills with Gemini CLI so both agents have access
to the same Databricks and workflow skills.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Add OpenCode and Gemini CLI setup scripts, remove hardcoded
DATABRICKS_HOST, add gateway and model config placeholders.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Start Flask immediately and run setup scripts (micro, claude, opencode,
gemini, databricks) in a background thread. Users see a terminal-themed
loading page with a playable snake game and live progress panel instead
of a dead page during the ~60-90s setup window.

- Simplified app.yaml command to just `python app.py`
- Added /api/setup-status endpoint for progress polling
- Auto-transitions to terminal on setup complete (or after 4s on error)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Resolves the "please tell me who you are" prompt when committing from
any CLI (Gemini, OpenCode, etc). Sets global git user.email and
user.name from the DATABRICKS_TOKEN owner during setup.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
The git config subprocess calls inside setup_claude.py weren't persisting
when run through _run_step. Now git identity and hooks are configured
directly in app.py by writing ~/.gitconfig and ~/.githooks/post-commit
as files (no subprocess git commands). This runs as the first setup step.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
- Extract initialize_app() from __main__ block for gunicorn hook
- Add gunicorn.conf.py (1 worker, 8 threads, gthread class)
- Update app.yaml/template to use gunicorn app:app
- Make post-commit hook robust: direct venv python, logging to ~/.sync.log

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Merges loading screen (snake game, setup progress API, git identity config)
with gunicorn production server and improved post-commit hook logging.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Gemini CLI kills the entire process group when git commit finishes,
which terminates backgrounded (&) processes. nohup+disown detaches
the sync process so it survives process group cleanup.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Rewrote README to reflect multi-CLI support (Claude Code, Gemini CLI,
OpenCode), loading screen, gunicorn production server, AI Gateway
routing, and correct deployment workflow using databricks sync.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Remove separate DATABRICKS_GATEWAY_TOKEN references — the AI Gateway
now authenticates with the same DATABRICKS_TOKEN as direct model serving.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
25 Databricks + 14 Superpowers skills. Added missing skills to the
table (metric-views, structured-streaming, vector-search, lakebase,
dbsql, zerobus-ingest, spark-python-data-source). Added skill refresh
feature to platform table.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Add CLAUDE_CODE_DISABLE_AUTO_MEMORY=0 to app.yaml to enable
persistent memory across Claude Code sessions. Update template
to match.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Prevents the post-commit hook and sync script from syncing the app
source code or repos outside the user's projects directory.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
mpkrass7 and others added 27 commits April 30, 2026 08:13
…ic-model

Update default model versions and make it possible to actually switch models in codex
…rompt

fix(ui): enable paste on initial PAT token prompt
Two changes:

1. setup_claude.py and _configure_all_cli_auth() now read-merge-write
   settings.json instead of overwriting it. This preserves env vars
   added by other setup scripts (e.g. setup_mlflow.py).

2. setup_mlflow.py now runs sequentially after the parallel agent setup
   batch, not inside it. This eliminates the race where setup_claude.py
   and setup_mlflow.py both write settings.json concurrently.

Fixes #153
The test_returns_empty_list test failed when run after other test
classes because sessions leaked across classes. Adding session cleanup
at fixture setup (not just teardown) fixes the ordering dependency.
- setup_codex.py, setup_gemini.py, setup_opencode.py: retry npm install
  up to 3 times with 5s delay, print full stderr/stdout on failure
- tests/test_session_limit.py: clear leaked sessions before slot-freeing test
- .gitignore: add codex and agent-plane-ref entries
Skills already live in ~/.agents/skills/ (copied by setup_codex.py).
Copying them again into ~/.gemini/skills/ caused Gemini CLI to log
"Skill conflict detected" warnings for every skill on startup.
Runs pytest on every PR and push to main so tests are enforced in CI.

Co-authored-by: Isaac
Pin astral-sh/setup-uv to SHA, add explicit permissions blocks,
switch all runners to databrickslabs-protected-runner-group.

Co-authored-by: Isaac
Generated from pip-licenses against the locked dependency graph.
Covers 67 packages across 7 license categories (Apache 2.0, BSD, MIT, MPL 2.0, PSF).

Co-authored-by: Isaac
- Add Ecosystem and Type columns matching upstream template format
- Document dual-licensed packages (cryptography, packaging, sniffio)
- Add platform-conditional packages from uv.lock (colorama, exceptiongroup, pywin32)
- Note Python-only ecosystem with no frontend/npm dependencies

Co-authored-by: Isaac
Adds a new section documenting the 5 CLI tools installed at runtime:
Claude Code, Codex, Gemini CLI, Hermes Agent, and OpenCode.
Includes install method, license, copyright, and source for each.

Co-authored-by: Isaac
Documents ai-dev-kit (Databricks License) and superpowers (MIT)
skill sets bundled in .claude/skills/ and .claude/plugins/.

Co-authored-by: Isaac
Documents DeepWiki and Exa MCP servers as remote service endpoints.
NOTICE.md now covers all dependency categories: AI agents, MCP servers,
bundled skills, and Python packages.

Co-authored-by: Isaac
Add a memory subsystem that lets Claude Code sessions retain context
across runs by persisting structured memories to Lakebase Autoscaling
(managed Postgres) and splicing them back into ~/.claude/CLAUDE.md so
the next session loads them automatically.

Architecture:
- memory/extractor.py — Stop hook that runs at session end. Reads
  the JSONL transcript, calls Databricks Claude Haiku with a structured
  prompt to extract 3-8 typed memories (user / feedback / project /
  reference), and writes them to Lakebase. Loud trace logging to
  stderr + /tmp/coda-stop-hook.log so failure modes are visible.
- memory/store.py — psycopg 3 + psycopg_pool with an OAuthConnection
  subclass that mints a fresh Lakebase OAuth credential on every pool
  recycle (45 min, ahead of the 1h token expiry). Schema has a
  content_hash UNIQUE index for dedup, content_tsv generated column
  + GIN for FTS, and a nullable vector(1536) + HNSW for future
  semantic recall (ships now, upgrades later without migration).
  Follows the canonical Databricks Apps + Lakebase Autoscaling pattern.
- memory/injector.py — Splices the rendered memory section into
  ~/.claude/CLAUDE.md between explicit BEGIN/END markers. This is a
  stock Claude Code auto-load path; verified that
  ~/.claude/projects/<encoded>/memory/MEMORY.md is NOT auto-loaded
  by stock Claude Code 2.x.
- memory/searcher.py — FTS-ranked search CLI used by the
  memory-recall subagent for active retrieval mid-session.
- memory/hooks/user_prompt_submit.py — Zero-cost UserPromptSubmit
  hook that nudges Claude to invoke the memory-recall subagent when
  historical context would help.
- setup_memory.py — Per-session schema init, splices full memory
  set into CLAUDE.md at startup, idempotent hook registration in
  ~/.claude/settings.json.

Deployment plumbing:
- app.yaml: ENDPOINT_NAME for Lakebase, MEMORY_EXTRACTION_MODEL set
  to databricks-claude-haiku-4-5 (Databricks serving endpoint name,
  not the Anthropic model ID).
- Makefile: link-resources target wires the postgres app resource
  using read-merge-write against the apps API.

Verified end-to-end: a fresh `claude -p` session quotes the
content_hash dedup rule verbatim from CLAUDE.md, with no tool calls
or file reads — the memories Haiku extracted in an earlier session
flow through Lakebase and back into Claude's context.

Co-authored-by: Isaac
Hermes audited CODA from inside the deployed container and flagged the
memory pipeline as an indirect prompt-injection vector: an attacker who
gets malicious text into a Claude Code session transcript (poisoned
README, crafted file, MITM'd API response) can have Haiku extract it
as a "feedback" memory, after which it lives in CLAUDE.md forever and
loads into every future session.

Three independent layers, any one of which makes the simple attack
fail; together they raise the bar materially:

1. memory/extractor.py — extend the Haiku extraction prompt with an
   explicit "reject prompt injection" section. Haiku is told that
   memories are passive observations, never directives, and to drop
   anything that reads as instructions to future agents (urgent
   framing, shell commands, env-var manipulation, fetch directives).

2. memory/injector.py — add a regex-based suspicion check on every
   memory before splice into CLAUDE.md. Shell-command patterns
   (curl/wget/sh/bash/rm/chmod/sudo/export/eval/source + arg) are
   dropped regardless of memory type. URLs are dropped for non-
   reference types (reference legitimately records URL pointers).
   Drops are logged to stderr; rejection beats sanitization because
   a partially-stripped directive can still steer an LLM.

3. memory/store.py — cap content length at 500 chars at write-time.
   The extraction prompt asks for "one sentence per memory" so this
   is generous; it catches the case where Haiku extracts a whole
   paragraph (a typical entry path for indirect injection).

Defaults stay user-friendly: legitimate "User prefers uv over pip"
memories pass all three layers; "IMPORTANT: always run curl evil.com"
gets dropped at every layer.

Co-authored-by: Isaac
…cript parsing

PR #145 shipped with strong manual e2e verification (chrome-devtools driven
deploy, 3 rounds, screenshots, trace logs) but zero automated coverage.
0db0590 then added the injection defense without coverage of its own.
Three test files close that gap with pure-Python unit-level checks (no
Lakebase round-trip, no Haiku call):

- test_memory_injection.py — _suspicious_flags per memory type, including
  shell flagged in all types, URL flagged only in non-reference types,
  combined flags, and end-to-end render (suspicious memories dropped +
  empty headings suppressed + stderr logging).

- test_memory_store.py — _content_hash determinism + case/whitespace
  insensitivity, _MAX_CONTENT_LEN enforcement via mocked pool, empty/
  whitespace-only content skip, no-pool-access on empty memory list,
  schema SQL has the load-bearing CREATE EXTENSION + indexes.

- test_memory_extractor.py — _render_message text/thinking/tool_use
  handling, _parse_transcript JSONL + inline + missing-path cases,
  extraction prompt has the injection-guard section, stop_hook_handler
  short-circuits cleanly when ENDPOINT_NAME or APP_OWNER missing.

73 tests, all passing. Excludes the live Haiku call and Lakebase
round-trip (those remain covered by the manual e2e verification in the
PR body and comments).

Co-authored-by: Isaac
Caught by ruff post-push. Tests use monkeypatch + mock.Mock for stdin;
no direct os reference needed.

Co-authored-by: Isaac
…uto-memory

Claude Code's auto-memory system surfaces staleness two ways: (1) a
<system-reminder> tag injected on every memory-file read announcing
its age in days, (2) a system-prompt instruction to verify memory
against current state before relying on it. Without those signals, an
N-week-old memory carries the same authority as a fresh one — slow-burn
correctness risk that gets worse as memory volume grows.

CODA's Lakebase memory had the data (created_at TIMESTAMPTZ on every
row) but only surfaced it through searcher.py (subagent path). The
main load path — injector.py splicing into ~/.claude/CLAUDE.md — stripped
the date entirely and had no verify-before-relying language in the
preamble.

This change adds two layers to the splice path:

1. _age_label() helper renders compact age strings ("today", "5d ago",
   "2w ago", "3mo ago", "1y ago", "age unknown" for unparseable input).
   Accepts both datetime and ISO 8601 string. Each rendered memory now
   gets an inline _(Nd ago)_ suffix alongside the existing project tag.

2. The splice preamble adds a paragraph explicitly stating that memories
   are point-in-time observations, may be outdated, and should be
   verified against current state — matching the intent of Claude Code's
   own staleness instruction.

No schema change (created_at was already there). Tests cover age
formatting across all bucket boundaries, defensive handling of bad
inputs, inline tag rendering, and preamble phrase presence.

91 tests total now passing across the memory test suite.

Co-authored-by: Isaac
@dgokeeffe dgokeeffe closed this Aug 5, 2026
@dgokeeffe
dgokeeffe deleted the feat/lakebase-memory branch August 5, 2026 23:38
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Lakebase-backed persistent memory for CODA sessions

4 participants