| 1 | /** Additive prompt sections. Rust owns selection, attribution and prompt assembly. */ |
| 2 | import { Service } from '@deepseek-ai/cordis' |
| 3 | import type { RpcPeer } from '../rpc.ts' |
| 4 | import { OwnedRegistrations, type OwnedEntry, type OwnerBase } from './owned.ts' |
| 5 | |
| 6 | export const MAX_PROMPT_SECTION_BYTES = 4 * 1024 |
| 7 | export const MAX_PROMPT_OWNER_BYTES = 32 * 1024 |
| 8 | export const MAX_PROMPT_HOST_BYTES = 128 * 1024 |
| 9 | export const MAX_PROMPT_SECTIONS_PER_OWNER = 128 |
| 10 | export const MAX_PROMPT_SECTIONS_PER_HOST = 1024 |
| 11 | const SECTION_ID = /^[a-z][a-z0-9_-]{0,63}$/u |
| 12 | |
| 13 | export interface PromptSectionDefinition { |
| 14 | readonly id: string |
| 15 | readonly text: string |
| 16 | /** Only Core supplies these values, once at the accepted turn boundary. */ |
| 17 | readonly interpolate?: 'model-cwd' |
| 18 | } |
| 19 | |
| 20 | export interface LocalPromptSection<O extends OwnerBase = OwnerBase> extends OwnedEntry<O> { |
| 21 | definition: PromptSectionDefinition |
| 22 | } |
| 23 | |
| 24 | interface PromptReservation<O extends OwnerBase> { |
| 25 | entry: LocalPromptSection<O> |
| 26 | bytes: number |
| 27 | dispose: () => void |
| 28 | } |
| 29 | |
| 30 | /** Copy plain text now; mutating the author's object cannot change an admitted section. */ |
| 31 | export function normalizePromptSection(value: unknown): PromptSectionDefinition { |
| 32 | if (typeof value !== 'object' || value === null || Array.isArray(value)) { |
| 33 | throw new TypeError('prompt section must be an object with id and text') |
| 34 | } |
| 35 | if (Object.keys(value).some((key) => key !== 'id' && key !== 'text' && key !== 'interpolate')) { |
| 36 | throw new TypeError('prompt section supports only id and text, plus optional model-cwd interpolation') |
| 37 | } |
| 38 | const { id, text, interpolate } = value as { id?: unknown; text?: unknown; interpolate?: unknown } |
| 39 | if (typeof id !== 'string' || !SECTION_ID.test(id)) { |
| 40 | throw new TypeError('prompt section id must be lower case, start with a letter, and use a-z, 0-9, _ or - (at most 64 characters)') |
| 41 | } |
| 42 | if (typeof text !== 'string' || text.trim().length === 0) { |
| 43 | throw new TypeError(`prompt section "${id}" needs non-empty text`) |
| 44 | } |
| 45 | if (/[\u0000-\u0008\u000b\u000c\u000e-\u001f\u007f-\u009f]/u.test(text)) { |
| 46 | throw new TypeError(`prompt section "${id}" text contains control characters`) |
| 47 | } |
| 48 | if (Buffer.byteLength(text, 'utf8') > MAX_PROMPT_SECTION_BYTES) { |
| 49 | throw new RangeError(`prompt section "${id}" exceeds ${MAX_PROMPT_SECTION_BYTES} UTF-8 bytes`) |
| 50 | } |
| 51 | if (interpolate !== undefined && interpolate !== 'model-cwd') throw new TypeError('prompt interpolation supports only model-cwd') |
| 52 | if (interpolate !== undefined) validatePromptTemplate(text) |
| 53 | return Object.freeze({ id, text, ...(interpolate === undefined ? {} : { interpolate }) }) |
| 54 | } |
| 55 | |
| 56 | /** Borrowed strict simple-group semantics: unmatched opens are literal; only |
| 57 | * existing Core turn facts are permitted. Values are never expanded in JS. */ |
| 58 | export function validatePromptTemplate(text: string): void { |
| 59 | for (let open = text.indexOf('{{'); open >= 0;) { |
| 60 | const group = /^\{\{([^{}]*)\}\}/u.exec(text.slice(open)) |
| 61 | if (!group) { |
| 62 | if (text.indexOf('}}', open + 2) >= 0) throw new TypeError('malformed prompt variable reference') |
| 63 | open = text.indexOf('{{', open + 2) |
| 64 | continue |
| 65 | } |
| 66 | if (!/^[a-z][a-z0-9_]*$/u.test(group[1]!)) throw new TypeError('malformed prompt variable reference') |
| 67 | if (group[1] !== 'model' && group[1] !== 'cwd') throw new TypeError(`unknown Core prompt variable "{{${group[1]}}}"; supported variables: model, cwd`) |
| 68 | open = text.indexOf('{{', open + group[0].length) |
| 69 | } |
| 70 | } |
| 71 | |
| 72 | /** Reserve pending registrations too, so a burst cannot bypass byte/count limits. */ |
| 73 | export class PromptSections<O extends OwnerBase> { |
| 74 | private readonly registrations: OwnedRegistrations<O, LocalPromptSection<O>> |
| 75 | private readonly templates: OwnedRegistrations<O, LocalPromptSection<O>> |
| 76 | private readonly owners = new Map<O, Map<string, PromptReservation<O>>>() |
| 77 | private bytes = 0 |
| 78 | private count = 0 |
| 79 | |
| 80 | constructor(rpc: RpcPeer, ownedBy: (owner: O) => Map<number, LocalPromptSection<O>>, warn: (message: string, owner: O) => void) { |
| 81 | this.registrations = new OwnedRegistrations(rpc, 'prompt_section', ownedBy, warn) |
| 82 | this.templates = new OwnedRegistrations(rpc, 'prompt_template', ownedBy, warn) |
| 83 | } |
| 84 | |
| 85 | register(owner: O, definition: PromptSectionDefinition): () => void { |
| 86 | if (owner.state !== 'activating' && owner.state !== 'active') throw new Error('prompt owner is not live') |
| 87 | const section = normalizePromptSection(definition) |
| 88 | const sections = this.owners.get(owner) ?? new Map<string, PromptReservation<O>>() |
| 89 | if (sections.has(section.id)) throw new Error(`prompt section "${section.id}" is already registered; dispose it before registering it again`) |
| 90 | const bytes = Buffer.byteLength(section.text, 'utf8') |
| 91 | const ownerBytes = [...sections.values()].reduce((sum, item) => sum + item.bytes, 0) |
| 92 | if (ownerBytes + bytes > MAX_PROMPT_OWNER_BYTES || this.bytes + bytes > MAX_PROMPT_HOST_BYTES) { |
| 93 | throw new RangeError('prompt section owner or host UTF-8 byte limit reached') |
| 94 | } |
| 95 | if (sections.size >= MAX_PROMPT_SECTIONS_PER_OWNER || this.count >= MAX_PROMPT_SECTIONS_PER_HOST) { |
| 96 | throw new RangeError('prompt section owner or host registration limit reached') |
| 97 | } |
| 98 | const entry: LocalPromptSection<O> = { owner, name: section.id, definition: section, disposed: false } |
| 99 | const undo = (section.interpolate === undefined ? this.registrations : this.templates).add(entry, { name: section.id, description: section.text }) |
| 100 | const record: PromptReservation<O> = { entry, bytes, dispose: () => { |
| 101 | if (sections.get(section.id) !== record) return |
| 102 | sections.delete(section.id) |
| 103 | this.bytes -= bytes |
| 104 | this.count -= 1 |
| 105 | if (sections.size === 0) this.owners.delete(owner) |
| 106 | undo() |
| 107 | } } |
| 108 | sections.set(section.id, record) |
| 109 | this.owners.set(owner, sections) |
| 110 | this.bytes += bytes |
| 111 | this.count += 1 |
| 112 | return record.dispose |
| 113 | } |
| 114 | |
| 115 | /** Release reservations even after a plugin fails or times out during teardown. */ |
| 116 | forget(owner: O) { |
| 117 | for (const record of [...(this.owners.get(owner)?.values() ?? [])]) record.dispose() |
| 118 | this.registrations.forget(owner) |
| 119 | this.templates.forget(owner) |
| 120 | } |
| 121 | } |
| 122 | |
| 123 | export interface PromptHost<O extends OwnerBase> { |
| 124 | ownerOf(ctx: any): O | undefined |
| 125 | promptSections: PromptSections<O> |
| 126 | } |
| 127 | |
| 128 | /** The only author API: an owner-scoped contribution and its idempotent disposer. */ |
| 129 | export function definePromptService<O extends OwnerBase>(host: PromptHost<O>) { |
| 130 | class PromptShim extends Service { |
| 131 | constructor(ctx: any) { super(ctx, 'prompt') } |
| 132 | |
| 133 | registerSection(definition: PromptSectionDefinition): () => void { |
| 134 | const ctx: any = this.ctx |
| 135 | const owner = host.ownerOf(ctx) |
| 136 | if (!owner) throw new Error('prompt.registerSection called outside an extension owner') |
| 137 | const section = normalizePromptSection(definition) |
| 138 | return ctx.effect(() => host.promptSections.register(owner, section), `prompt.registerSection(${JSON.stringify(section.id)})`) |
| 139 | } |
| 140 | } |
| 141 | Object.freeze(PromptShim.prototype) |
| 142 | return PromptShim |
| 143 | } |
| 144 |