neon-branch-per-session

1/**
2 * neon-branch-per-session — Claude Mod
3 *
4 * One Neon database branch per Claude Code session. When the session starts
5 * the mod creates `<prefix>/<session id>` from the project's default branch
6 * (or `parentBranchId`), reads its connection string and sets it as
7 * DATABASE_URL, which every command Claude runs afterwards inherits
8 * (`$.env.set` reaches the process and what it starts). The branch carries an
9 * expiry, so it removes itself even when the session dies without a goodbye.
10 *
11 *   session.start   create (or reuse, on a resumed session) the branch, set the env
12 *   command.run     /neon: show the branch, `keep` it past its expiry, `delete` it
13 *   session.end     onEnd: expire (nothing to do), delete, or keep
14 *
15 * What this is not: a sandbox. Commands still run with the user's own
16 * permissions, and a project whose .env holds another database URL can point
17 * Claude at it. Loaders such as dotenv do not override a variable already in
18 * the environment, which is what makes the branch win in the common case.
19 *
20 * Needs Claude Code >= 2.1.287 and a Neon
21 * API key in the plugin's options. Never hardcode it in this file.
22 *
23 * Options (pluginConfigs["neon-branch-per-session@skills-dir"].options):
24 *   neonApiKey, projectId   required; without both the mod does nothing
25 *   parentBranchId: string  branch to fork (default: the project's default branch)
26 *   database, role: string  connection string target (default neondb / neondb_owner)
27 *   pooled: boolean         pooled connection string (default true)
28 *   expireHours: number     self-delete after this long (default 24, max 720)
29 *   onEnd: "expire" | "delete" | "keep"   what exiting does (default expire)
30 *   namePrefix: string      branch name prefix (default claude)
31 *   interactiveOnly: boolean skip claude -p / SDK runs (default true)
32 */
33import type { Register } from 'claude-code'
34import {
35  branchName,
36  connectionUri,
37  createBranch,
38  deleteBranch,
39  expiryIso,
40  getBranch,
41  hoursLeft,
42  keepBranch,
43  uriHost,
44} from './neon-api.ts'
45import { NeonError } from './neon-api.ts'
46import type { Branch } from './neon-api.ts'
47
48const COMMAND = 'neon'
49// the first release kept one map under this key; it is only read now, for sessions that began before the split
50const LEGACY_KEY = 'neon-branch-per-session:branches'
51// one key per session, so two sessions starting together never overwrite each other's entry
52const sessionKey = (id: string) => `neon-branch-per-session:session:${id}`
53
54type Held = { branch: Branch; host: string; isKept: boolean; sessionId: string }
55
56let held: Held | undefined
57
58const str = (v: unknown, fallback: string) => (typeof v === 'string' && v.trim() ? v.trim() : fallback)
59
60export const register: Register = (on, options) => {
61  const apiKey = str(options.neonApiKey, '')
62  const project = str(options.projectId, '')
63  const parentId = str(options.parentBranchId, '')
64  const database = str(options.database, 'neondb')
65  const role = str(options.role, 'neondb_owner')
66  const pooled = options.pooled !== false
67  const expireHours = typeof options.expireHours === 'number' ? options.expireHours : 24
68  const onEnd = options.onEnd === 'delete' || options.onEnd === 'keep' ? options.onEnd : 'expire'
69  const prefix = str(options.namePrefix, 'claude')
70  const interactiveOnly = options.interactiveOnly !== false
71  const isConfigured = apiKey !== '' && project !== ''
72
73  const statusText = () => {
74    if (!held) return undefined
75    const left = hoursLeft(held.branch.expiresAt, Date.now())
76    const life = held.isKept || left === undefined ? 'kept' : `expires in ${left}h`
77    return `neon: ${held.branch.name} · ${life}`
78  }
79
80  on('session.start', async ($, e, next) => {
81    const r = await next(e)
82    await $.command
83      .register({
84        name: COMMAND,
85        description: "This session's Neon branch: show it, keep it past its expiry, or delete it",
86        argumentHint: '[keep|delete]',
87        immediate: true,
88      })
89      .catch(err => $.ui.log(`neon-branch-per-session: /${COMMAND} not registered: ${err}`))
90
91    if (!isConfigured) {
92      $.ui.log('neon-branch-per-session: set neonApiKey and projectId in the plugin options to create a branch per session', {
93        to: 'debug',
94      })
95      return r
96    }
97    if (interactiveOnly && !e.isInteractive) return r
98
99    const f = (url: string, init?: Parameters<typeof $.http.fetch>[1]) => $.http.fetch(url, init)
100    try {
101      const sessionId = await $.session.id()
102      const name = branchName(prefix, sessionId)
103      const legacy = ((await $.store.get(LEGACY_KEY)) ?? {}) as Record<string, string>
104      // keyed by the full session id: the 8 characters in the branch name are not unique enough to share a branch on
105      // an entry written under the branch name by an earlier version is still honoured
106      const known = ((await $.store.get(sessionKey(sessionId))) as string | undefined) ?? legacy[name]
107      let branch = known ? await getBranch(f, apiKey, project, known) : undefined
108      const isReused = branch !== undefined
109      if (!branch) {
110        const make = (n: string) =>
111          createBranch(f, apiKey, project, { name: n, parentId: parentId || undefined, expiresAt: expiryIso(Date.now(), expireHours) })
112        // another session may already hold the 8-character name: retry once with the tail of the full id
113        branch = await make(name).catch(err => {
114          if (!(err instanceof NeonError) || (err.status !== 409 && err.status !== 422)) throw err
115          return make(branchName(prefix, sessionId, 16))
116        })
117        await $.store.set(sessionKey(sessionId), branch.id)
118      }
119      // held before the URI is read: a branch that exists can always be kept or deleted with /neon
120      held = { branch, host: '', isKept: !branch.expiresAt, sessionId }
121      const uri = await connectionUri(f, apiKey, project, { branchId: branch.id, database, role, pooled })
122      await $.env.set('DATABASE_URL', uri)
123      await $.env.set('NEON_BRANCH', branch.name)
124      held = { ...held, host: uriHost(uri) }
125      $.ui.status(statusText())
126      $.ui.toast(`neon: ${isReused ? 'reusing' : 'created'} branch ${branch.name}; DATABASE_URL points at it`)
127      $.ui.log(`neon-branch-per-session: ${isReused ? 'reused' : 'created'} ${branch.name} (${branch.id}) on ${held.host}`, {
128        to: 'debug',
129      })
130    } catch (err) {
131      $.ui.toast(`neon: ${held ? `branch ${held.branch.name} exists but DATABASE_URL is not set (` : 'no branch for this session ('}${err instanceof Error ? err.message : String(err)})`)
132      $.ui.log(`neon-branch-per-session: ${err instanceof Error ? err.message : String(err)}`)
133    }
134    return r
135  })
136
137  on('command.run', { command: COMMAND }, async ($, e) => {
138    if (!isConfigured) return { text: 'neon: set neonApiKey and projectId in the plugin options' }
139    if (!held) return { text: 'neon: no branch for this session (see the debug log for why)' }
140    const f = (url: string, init?: Parameters<typeof $.http.fetch>[1]) => $.http.fetch(url, init)
141    const arg = e.args.trim().toLowerCase()
142
143    if (arg === 'keep') {
144      try {
145        await keepBranch(f, apiKey, project, held.branch.id)
146        held = { ...held, isKept: true, branch: { ...held.branch, expiresAt: undefined } }
147        $.ui.status(statusText())
148        return { text: `neon: ${held.branch.name} will no longer expire; delete it in the Neon console or with /${COMMAND} delete` }
149      } catch (err) {
150        return { text: `neon: ${err instanceof Error ? err.message : String(err)}` }
151      }
152    }
153
154    if (arg === 'delete') {
155      const gone = held.branch.name
156      try {
157        await deleteBranch(f, apiKey, project, held.branch.id)
158        await $.store.delete(sessionKey(held.sessionId))
159        held = undefined
160        await $.env.set('DATABASE_URL', undefined)
161        await $.env.set('NEON_BRANCH', undefined)
162        $.ui.status(undefined)
163        return { text: `neon: deleted ${gone}; DATABASE_URL is unset for this session` }
164      } catch (err) {
165        return { text: `neon: ${err instanceof Error ? err.message : String(err)}` }
166      }
167    }
168
169    const left = hoursLeft(held.branch.expiresAt, Date.now())
170    const life = held.isKept || left === undefined ? 'kept (no expiry)' : `expires in ${left}h`
171    return {
172      text: `neon: ${held.branch.name} (${held.branch.id}) · ${life} · DATABASE_URL on ${held.host} · /${COMMAND} keep | delete`,
173    }
174  })
175
176  // `clear` and `resume` end a conversation, not the process that owns the branch
177  on('session.end', async ($, e, next) => {
178    const r = await next(e)
179    if (!held || e.reason === 'clear' || e.reason === 'resume') return r
180    const f = (url: string, init?: Parameters<typeof $.http.fetch>[1]) => $.http.fetch(url, init)
181    try {
182      if (onEnd === 'delete' && !held.isKept) {
183        await deleteBranch(f, apiKey, project, held.branch.id)
184        await $.store.delete(sessionKey(held.sessionId))
185      } else if (onEnd === 'keep' && !held.isKept) {
186        await keepBranch(f, apiKey, project, held.branch.id)
187      }
188    } catch (err) {
189      $.ui.log(`neon-branch-per-session: end-of-session ${onEnd} failed: ${err instanceof Error ? err.message : String(err)}`)
190    }
191    return r
192  })
193}
194

1/**
2 * neon-api.ts — the Neon REST calls the mod makes, over a `fetch` it is given.
3 *
4 * Reference: https://api-docs.neon.tech (base https://console.neon.tech/api/v2,
5 * `Authorization: Bearer <key>`). The five calls used:
6 *   POST   /projects/{p}/branches                 create; `branch.expires_at` (RFC 3339, at most 30 days out)
7 *   GET    /projects/{p}/branches/{b}             does it still exist
8 *   PATCH  /projects/{p}/branches/{b}             `{ branch: { expires_at: null } }` clears the expiry
9 *   DELETE /projects/{p}/branches/{b}
10 *   GET    /projects/{p}/connection_uri           ?branch_id&database_name&role_name&pooled -> { uri }
11 *
12 * No `$` in here: the module passes `(url, init) => $.http.fetch(url, init)`.
13 * The API key is only ever sent in the Authorization header; nothing here
14 * returns or logs it, and the connection string is returned, never logged.
15 */
16
17export const API = 'https://console.neon.tech/api/v2'
18
19export type Fetch = (
20  url: string,
21  init?: { method?: string; headers?: Record<string, string>; body?: string },
22) => Promise<{ status: number; ok: boolean; text: string }>
23
24export type Branch = { id: string; name: string; expiresAt?: string }
25
26export class NeonError extends Error {
27  status: number
28  constructor(message: string, status: number) {
29    super(message)
30    this.status = status
31  }
32}
33
34const MAX_HOURS = 720
35
36/** `<prefix>/<first `length` characters of the session id>`, with anything a branch name should not hold dropped. */
37export function branchName(prefix: string, sessionId: string, length = 8): string {
38  const clean = (s: string) => s.replace(/[^A-Za-z0-9._-]/g, '')
39  const id = clean(sessionId).slice(0, length) || 'session'
40  const head = clean(prefix)
41  return head ? `${head}/${id}` : id
42}
43
44/** RFC 3339 expiry `hours` from `now`, held to Neon's 1 hour to 30 days window. */
45export function expiryIso(now: number, hours: number): string {
46  const h = Number.isFinite(hours) ? Math.min(MAX_HOURS, Math.max(1, hours)) : 24
47  return new Date(now + h * 3_600_000).toISOString().replace(/\.\d{3}Z$/, 'Z')
48}
49
50export function hoursLeft(expiresAt: string | undefined, now: number): number | undefined {
51  if (!expiresAt) return undefined
52  const ms = Date.parse(expiresAt) - now
53  return Number.isNaN(ms) ? undefined : Math.max(0, Math.round(ms / 3_600_000))
54}
55
56/** The host of a connection string, for showing where it points without showing the credentials. */
57export function uriHost(uri: string): string {
58  const m = /@([^/:?]+)/.exec(uri)
59  return m?.[1] ?? 'unknown host'
60}
61
62function headers(key: string, withBody: boolean): Record<string, string> {
63  const h: Record<string, string> = { authorization: `Bearer ${key}`, accept: 'application/json' }
64  if (withBody) h['content-type'] = 'application/json'
65  return h
66}
67
68async function call(f: Fetch, key: string, method: string, path: string, body?: unknown) {
69  const res = await f(`${API}${path}`, {
70    method,
71    headers: headers(key, body !== undefined),
72    body: body === undefined ? undefined : JSON.stringify(body),
73  })
74  if (!res.ok) {
75    let detail = ''
76    try {
77      detail = (JSON.parse(res.text) as { message?: string }).message ?? ''
78    } catch {
79      detail = res.text.slice(0, 120)
80    }
81    throw new NeonError(`Neon ${method} ${path.split('?')[0].replace(/\/(br|ep)-[^/]+/g, '/$1-…')} answered ${res.status}${detail ? `: ${detail}` : ''}`, res.status)
82  }
83  return res.text ? (JSON.parse(res.text) as Record<string, unknown>) : {}
84}
85
86type BranchBody = { id?: string; name?: string; expires_at?: string }
87const toBranch = (b: BranchBody | undefined): Branch => {
88  if (!b?.id) throw new Error('Neon answered without a branch id')
89  return { id: b.id, name: b.name ?? b.id, expiresAt: b.expires_at }
90}
91
92export async function createBranch(
93  f: Fetch,
94  key: string,
95  project: string,
96  opts: { name: string; parentId?: string; expiresAt: string },
97): Promise<Branch> {
98  const res = await call(f, key, 'POST', `/projects/${encodeURIComponent(project)}/branches`, {
99    branch: {
100      name: opts.name,
101      expires_at: opts.expiresAt,
102      ...(opts.parentId ? { parent_id: opts.parentId } : {}),
103    },
104    endpoints: [{ type: 'read_write' }],
105  })
106  return toBranch(res.branch as BranchBody | undefined)
107}
108
109/** undefined when the branch is gone. */
110export async function getBranch(f: Fetch, key: string, project: string, id: string): Promise<Branch | undefined> {
111  try {
112    const res = await call(f, key, 'GET', `/projects/${encodeURIComponent(project)}/branches/${encodeURIComponent(id)}`)
113    return toBranch(res.branch as BranchBody | undefined)
114  } catch (err) {
115    if (err instanceof NeonError && err.status === 404) return undefined
116    throw err
117  }
118}
119
120export async function keepBranch(f: Fetch, key: string, project: string, id: string): Promise<void> {
121  await call(f, key, 'PATCH', `/projects/${encodeURIComponent(project)}/branches/${encodeURIComponent(id)}`, {
122    branch: { expires_at: null },
123  })
124}
125
126export async function deleteBranch(f: Fetch, key: string, project: string, id: string): Promise<void> {
127  await call(f, key, 'DELETE', `/projects/${encodeURIComponent(project)}/branches/${encodeURIComponent(id)}`)
128}
129
130export async function connectionUri(
131  f: Fetch,
132  key: string,
133  project: string,
134  opts: { branchId: string; database: string; role: string; pooled: boolean },
135): Promise<string> {
136  const q = new URLSearchParams({
137    branch_id: opts.branchId,
138    database_name: opts.database,
139    role_name: opts.role,
140    pooled: String(opts.pooled),
141  })
142  const res = await call(f, key, 'GET', `/projects/${encodeURIComponent(project)}/connection_uri?${q}`)
143  if (typeof res.uri !== 'string' || !res.uri) throw new Error('Neon answered without a connection URI')
144  return res.uri
145}
146