diff --git a/packages/core/system-prompt/src/index.ts b/packages/core/system-prompt/src/index.ts index 12171f1b2e..c46c38515c 100644 --- a/packages/core/system-prompt/src/index.ts +++ b/packages/core/system-prompt/src/index.ts @@ -6,8 +6,8 @@ import { Context, Service } from 'cordis' import z from 'schemastery' -import { scopeOf, scopeTarget } from '@deepseek-ai/dsh-scope' -import type { ScopeKey, Scoped } from '@deepseek-ai/dsh-scope' +import { AnonymousEntries, NamedEntries, ScopedLayers, scopeTarget } from '@deepseek-ai/dsh-scope' +import type { ScopeKey, ScopeLayer, Scoped } from '@deepseek-ai/dsh-scope' import type { ToolSchema } from '@deepseek-ai/dsh-llm' declare module 'cordis' { @@ -209,6 +209,39 @@ function interpolate(section: AssembledSection, variables: Record ToolProviderResult + +/** One prompt-variable provider stored in a prompt layer. */ +type VariableProvider = (context: AssembleContext) => string | undefined + +/** All prompt registrations owned by one global or scoped layer. */ +class PromptLayer implements ScopeLayer { + readonly sections: NamedEntries + readonly toolProviders = new AnonymousEntries() + readonly variables: NamedEntries + + /** + * Create one prompt layer with diagnostics specific to its ownership scope. + * @param scope - the scoped owner, or `undefined` for global registrations. + */ + constructor(scope: ScopeKey | undefined) { + this.sections = new NamedEntries(name => new Error(scope === undefined + ? `prompt section "${name}" is already registered (for a per-agent override, register through that agent's \`agent.ctx\` instead)` + : `prompt section "${name}" is already registered in this scope`)) + this.variables = new NamedEntries(name => new Error(scope === undefined + ? `prompt variable "${name}" is already registered (for a per-agent value, register through that agent's \`agent.ctx\` instead)` + : `prompt variable "${name}" is already registered in this scope`)) + } + + /** @returns whether this layer owns no prompt registrations. */ + isEmpty(): boolean { + return this.sections.isEmpty() + && this.toolProviders.isEmpty() + && this.variables.isEmpty() + } +} + /** Registry service for the prompt inputs assembled before each model step. */ export class SystemPrompt extends Service { static Config: z = z.object({ @@ -217,13 +250,10 @@ export class SystemPrompt extends Service { toolOrder: z.array(z.string()).default(undefined as unknown as string[]), }) - private sections: PromptSection[] = [] - private toolProviders: ((context: AssembleContext) => ToolProviderResult)[] = [] - private variableProviders = new Map string | undefined>() - /** Per-scope layers (`@deepseek-ai/dsh-scope`); entries drop when a layer empties, so a disposed scope leaves no residue. */ - private scopedSections = new Map() - private scopedToolProviders = new Map ToolProviderResult)[]>() - private scopedVariableProviders = new Map string | undefined>>() + private readonly layers = new ScopedLayers( + scope => new PromptLayer(scope), + () => { this.ctx.emit('system-prompt/change') }, + ) private readonly toolOrder: string[] | undefined constructor(ctx: Context, config: Config) { @@ -255,34 +285,11 @@ export class SystemPrompt extends Service { if (!Number.isFinite(section.order)) { throw new TypeError(`prompt section "${section.name}" order must be a finite number`) } - const scope = scopeOf(this.ctx) - const dispose = this.ctx.effect(function* (this: SystemPrompt) { - const layer = scope === undefined - ? this.sections - : this.scopedSections.get(scope) ?? (() => { - const created: PromptSection[] = [] - this.scopedSections.set(scope, created) - return created - })() - if (layer.some(existing => existing.name === section.name)) { - throw new Error(scope === undefined - ? `prompt section "${section.name}" is already registered (for a per-agent override, register through that agent's \`agent.ctx\` instead)` - : `prompt section "${section.name}" is already registered in this scope`) - } - layer.push(section) - // Install rollback before notifying listeners that may throw. - yield () => { - const index = layer.indexOf(section) - /* v8 ignore next 3 -- defensive: section was registered, so indexOf is guaranteed >= 0 */ - if (index >= 0) layer.splice(index, 1) - if (scope !== undefined && layer.length === 0) this.scopedSections.delete(scope) - this.ctx.emit('system-prompt/change') - } - this.ctx.emit('system-prompt/change') - }.bind(this), 'systemPrompt.section()') - // Return the exact disposer so composite effects preserve teardown order. - // eslint-disable-next-line @typescript-eslint/no-misused-promises -- synchronous cleanup; direct return preserves disposer identity - return dispose + return this.layers.effect( + this.ctx, + layer => layer.sections.insert(section.name, section), + { label: 'systemPrompt.section()' }, + ) } /** @@ -293,29 +300,11 @@ export class SystemPrompt extends Service { * @returns the exact Cordis effect disposer. */ tools(provider: (context: AssembleContext) => ToolProviderResult): () => void { - const scope = scopeOf(this.ctx) - const dispose = this.ctx.effect(function* (this: SystemPrompt) { - const layer = scope === undefined - ? this.toolProviders - : this.scopedToolProviders.get(scope) ?? (() => { - const created: ((context: AssembleContext) => ToolProviderResult)[] = [] - this.scopedToolProviders.set(scope, created) - return created - })() - layer.push(provider) - // Install rollback before notifying listeners that may throw. - yield () => { - const index = layer.indexOf(provider) - /* v8 ignore next 3 -- defensive: provider was registered, so indexOf is guaranteed >= 0 */ - if (index >= 0) layer.splice(index, 1) - if (scope !== undefined && layer.length === 0) this.scopedToolProviders.delete(scope) - this.ctx.emit('system-prompt/change') - } - this.ctx.emit('system-prompt/change') - }.bind(this), 'systemPrompt.tools()') - // Return the exact disposer so composite effects preserve teardown order. - // eslint-disable-next-line @typescript-eslint/no-misused-promises -- synchronous cleanup; direct return preserves disposer identity - return dispose + return this.layers.effect( + this.ctx, + layer => layer.toolProviders.append(provider), + { label: 'systemPrompt.tools()' }, + ) } /** @@ -330,32 +319,11 @@ export class SystemPrompt extends Service { if (!VARIABLE_NAME.test(name)) { throw new Error(`invalid prompt variable name "${name}" (must match ${String(VARIABLE_NAME)})`) } - const scope = scopeOf(this.ctx) - const dispose = this.ctx.effect(function* (this: SystemPrompt) { - const layer = scope === undefined - ? this.variableProviders - : this.scopedVariableProviders.get(scope) ?? (() => { - const created = new Map string | undefined>() - this.scopedVariableProviders.set(scope, created) - return created - })() - if (layer.has(name)) { - throw new Error(scope === undefined - ? `prompt variable "${name}" is already registered (for a per-agent value, register through that agent's \`agent.ctx\` instead)` - : `prompt variable "${name}" is already registered in this scope`) - } - layer.set(name, provider) - // Install rollback before notifying listeners that may throw. - yield () => { - layer.delete(name) - if (scope !== undefined && layer.size === 0) this.scopedVariableProviders.delete(scope) - this.ctx.emit('system-prompt/change') - } - this.ctx.emit('system-prompt/change') - }.bind(this), 'systemPrompt.variable()') - // Return the exact disposer so composite effects preserve teardown order. - // eslint-disable-next-line @typescript-eslint/no-misused-promises -- synchronous cleanup; direct return preserves disposer identity - return dispose + return this.layers.effect( + this.ctx, + layer => layer.variables.insert(name, provider), + { label: 'systemPrompt.variable()' }, + ) } /** @@ -370,23 +338,19 @@ export class SystemPrompt extends Service { const scope = context.scope // Scoped variables shadow globals. const variables: Record = {} - for (const [name, provider] of this.variableProviders) { + for (const [name, provider] of this.layers.global.variables.entries()) { variables[name] = provider(context) } - const scopedVariables = scope === undefined ? undefined : this.scopedVariableProviders.get(scope) - for (const [name, provider] of scopedVariables ?? []) { + const scopedVariables = this.layers.peek(scope)?.variables + for (const [name, provider] of scopedVariables?.entries() ?? []) { variables[name] = provider(context) } // Scoped sections shadow globals before the stable order sort. - const sectionByName = new Map() - for (const section of this.sections) sectionByName.set(section.name, section) - for (const section of (scope === undefined ? [] : this.scopedSections.get(scope)) ?? []) { - sectionByName.set(section.name, section) - } + const sectionByName = this.layers.merge(scope, layer => layer.sections) // Validate order against pre-restriction names while collecting visible schemas. const providers = [ - ...this.toolProviders, - ...(scope === undefined ? [] : this.scopedToolProviders.get(scope)) ?? [], + ...this.layers.global.toolProviders.values(), + ...(this.layers.peek(scope)?.toolProviders.values() ?? []), ] const collected: ToolSchema[] = [] const knownNames = new Set() diff --git a/packages/core/system-prompt/tests/scoped.spec.ts b/packages/core/system-prompt/tests/scoped.spec.ts index aac3f79e88..a3c85b7886 100644 --- a/packages/core/system-prompt/tests/scoped.spec.ts +++ b/packages/core/system-prompt/tests/scoped.spec.ts @@ -1,4 +1,4 @@ -import { describe, expect, it } from 'vitest' +import { describe, expect, it, vi } from 'vitest' import { Context } from 'cordis' import { createScope, scopeOf } from '@deepseek-ai/dsh-scope' import type { Scope, ScopeKey } from '@deepseek-ai/dsh-scope' @@ -63,6 +63,21 @@ describe('scoped sections', () => { expect(() => scope.ctx.systemPrompt.section({ name: 'y', order: 1, text: 'b' })).toThrow(/already registered in this scope/) }) + it('shadows a global section before evaluating either text provider', async () => { + const ctx = await mount() + const scope = await mintScope(ctx, 'child') + const globalText = vi.fn(() => 'global text') + const scopedText = vi.fn(() => 'scoped text') + ctx.systemPrompt.section({ name: 'shared', order: 1, text: globalText }) + scope.ctx.systemPrompt.section({ name: 'shared', order: 1, text: scopedText }) + + const assembly = await ctx.systemPrompt.assemble({ scope: scopeKeyOf(scope) }) + + expect(assembly.sections.find(section => section.name === 'shared')?.text).toBe('scoped text') + expect(globalText).not.toHaveBeenCalled() + expect(scopedText).toHaveBeenCalledOnce() + }) + }) describe('scoped variables', () => { diff --git a/packages/core/system-prompt/tests/system-prompt.spec.ts b/packages/core/system-prompt/tests/system-prompt.spec.ts index b084104d22..d4fdbdd684 100644 --- a/packages/core/system-prompt/tests/system-prompt.spec.ts +++ b/packages/core/system-prompt/tests/system-prompt.spec.ts @@ -157,6 +157,24 @@ describe('SystemPrompt', () => { expect((await ctx.systemPrompt.assemble()).tools.map(t => t.name)).toEqual(['t']) }) + it('snapshots tool-provider membership before evaluating an assembly', async () => { + const ctx = new Context() + await ctx.plugin(SystemPrompt) + let added = false + ctx.systemPrompt.tools(() => { + if (!added) { + added = true + ctx.systemPrompt.tools(() => ({ + schemas: [{ name: 'late', description: '', parameters: {} }], + })) + } + return { schemas: [{ name: 'first', description: '', parameters: {} }] } + }) + + expect((await ctx.systemPrompt.assemble()).tools.map(tool => tool.name)).toEqual(['first']) + expect((await ctx.systemPrompt.assemble()).tools.map(tool => tool.name)).toEqual(['first', 'late']) + }) + it('rolls back a variable when a system-prompt/change listener throws (P1-1)', async () => { const ctx = new Context() await ctx.plugin(SystemPrompt) @@ -314,6 +332,24 @@ describe('SystemPrompt', () => { expect((await ctx.systemPrompt.assemble()).variables).toEqual({}) }) + it('live-iterates variables registered by an earlier provider', async () => { + const ctx = new Context() + await ctx.plugin(SystemPrompt) + let added = false + ctx.systemPrompt.variable('first', () => { + if (!added) { + added = true + ctx.systemPrompt.variable('late', () => 'second value') + } + return 'first value' + }) + + expect((await ctx.systemPrompt.assemble()).variables).toEqual({ + first: 'first value', + late: 'second value', + }) + }) + it('rejects a duplicate variable name and an unreferenceable name', async () => { const ctx = new Context() await ctx.plugin(SystemPrompt) diff --git a/packages/core/tools/src/index.ts b/packages/core/tools/src/index.ts index 086d6a3271..11f57995e9 100644 --- a/packages/core/tools/src/index.ts +++ b/packages/core/tools/src/index.ts @@ -6,8 +6,8 @@ import { Context, Service } from 'cordis' import z from 'schemastery' -import { scopeOf, scopeTarget } from '@deepseek-ai/dsh-scope' -import type { ScopeKey, Scoped } from '@deepseek-ai/dsh-scope' +import { AnonymousEntries, NamedEntries, ScopedLayers, scopeOf, scopeTarget } from '@deepseek-ai/dsh-scope' +import type { ScopeKey, ScopeLayer, Scoped } from '@deepseek-ai/dsh-scope' import type { CallId, ContentBlock, ToolSchema } from '@deepseek-ai/dsh-llm' import { assertNever, deepFreeze, HarnessError } from '@deepseek-ai/dsh-llm' import type { Agent, HookContext } from '@deepseek-ai/dsh-agent' @@ -463,9 +463,40 @@ interface ToolView { */ export type ToolGuard = (execution: Readonly) => string | undefined -/** One guard registration; the wrapper preserves independent duplicate registrations. */ -interface ToolGuardRegistration { - guard: ToolGuard +/** One scope's complete tool-registry contribution. */ +class ToolLayer implements ScopeLayer { + readonly tools: NamedEntries + readonly restrictions = new AnonymousEntries() + readonly guards = new AnonymousEntries() + + constructor(scope: ScopeKey | undefined) { + this.tools = new NamedEntries(name => new Error(scope === undefined + ? `tool "${name}" is already registered (for a per-agent variant, register through that agent's \`agent.ctx\` instead)` + : `tool "${name}" is already registered in this scope`)) + } + + /** Whether every contribution table in this aggregate layer is empty. */ + isEmpty(): boolean { + return this.tools.isEmpty() && this.restrictions.isEmpty() && this.guards.isEmpty() + } + + /** Whether every compiled restriction in this layer admits a global tool name. */ + admits(name: string): boolean { + for (const filter of this.restrictions.values()) { + if ((filter.allow !== undefined && !filter.allow.has(name)) + || (filter.deny !== undefined && filter.deny.has(name))) return false + } + return true + } + + /** First monotonic denial from this layer's live guard registrations. */ + guardReason(exec: ToolExecution): string | undefined { + for (const guard of this.guards.values()) { + const reason = guard(exec) + if (reason !== undefined) return reason + } + return undefined + } } /** Approval decision plus whether the approval channel reported cancellation. */ @@ -509,13 +540,10 @@ export class ToolRegistry extends Service { private deferredContexts = new WeakMap() /** Original caller cancellation, kept outside the wrapper-mutable execution object. */ private cancellationStates = new WeakMap() - private global = new Map() - private scoped = new Map>() - /** Compiled restriction filters, per scope (see {@link restrict}). */ - private restrictions = new Map() - /** Monotonic post-policy guards, split into global and per-agent layers. */ - private globalGuards = new Set() - private scopedGuards = new Map>() + private readonly layers = new ScopedLayers( + scope => new ToolLayer(scope), + () => { this.ctx.emit('tools/change') }, + ) private readonly mode: ToolPresentationMode /** Reserved presentation transport, kept outside the filterable registration layers. */ private readonly codeTransport: ToolDefinition | undefined @@ -593,7 +621,6 @@ export class ToolRegistry extends Service { * @returns the exact disposer that unregisters the tool. */ register(definition: ToolDefinition): () => void { - const scope = scopeOf(this.ctx) const name = definition.name const timeoutMs = definition.timeoutMs if (timeoutMs !== undefined @@ -603,26 +630,11 @@ export class ToolRegistry extends Service { if (this.codeTransport !== undefined && name === RUN_CODE_NAME) { throw new Error(`tool name "${RUN_CODE_NAME}" is reserved for the Code Mode presentation transport and cannot be registered or shadowed`) } - const dispose = this.ctx.effect(function* (this: ToolRegistry) { - const layer = scope === undefined ? this.global : this.layerFor(scope) - if (layer.has(name)) { - throw new Error(scope === undefined - ? `tool "${name}" is already registered (for a per-agent variant, register through that agent's \`agent.ctx\` instead)` - : `tool "${name}" is already registered in this scope`) - } - layer.set(name, definition) - // Install rollback before notifying listeners. - yield () => { - layer.delete(name) - // Drop empty scope layers. - if (scope !== undefined && layer.size === 0) this.scoped.delete(scope) - this.ctx.emit('tools/change') - } - this.ctx.emit('tools/change') - }.bind(this), 'tools.register()') - // Return the exact disposer so composite effects preserve teardown order. - // eslint-disable-next-line @typescript-eslint/no-misused-promises -- synchronous cleanup; direct return preserves disposer identity - return dispose + return this.layers.effect( + this.ctx, + layer => layer.tools.insert(name, definition), + { label: 'tools.register()' }, + ) } /** @@ -655,22 +667,11 @@ export class ToolRegistry extends Service { if (unknown.length > 0) { throw new Error(`tools.restrict() names unknown global tool${unknown.length > 1 ? 's' : ''} ${unknown.map(n => `"${n}"`).join(', ')}; known global tools: ${[...known].sort().join(', ') || '(none)'}`) } - const dispose = this.ctx.effect(function* (this: ToolRegistry) { - const list = this.restrictions.get(scope) ?? [] - this.restrictions.set(scope, list) - list.push(compiled) - yield () => { - const index = list.indexOf(compiled) - /* v8 ignore next 3 -- defensive: the compiled restriction was pushed, so indexOf is guaranteed >= 0 */ - if (index >= 0) list.splice(index, 1) - if (list.length === 0) this.restrictions.delete(scope) - this.ctx.emit('tools/change') - } - this.ctx.emit('tools/change') - }.bind(this), 'tools.restrict()') - // Return the exact disposer so composite effects preserve teardown order. - // eslint-disable-next-line @typescript-eslint/no-misused-promises -- synchronous cleanup; direct return preserves disposer identity - return dispose + return this.layers.effect( + this.ctx, + layer => layer.restrictions.append(compiled), + { label: 'tools.restrict()' }, + ) } /** @@ -684,63 +685,18 @@ export class ToolRegistry extends Service { * @returns the exact disposer that unregisters the guard. */ guard(guard: ToolGuard): () => void { - const scope = scopeOf(this.ctx) - const registration = { guard } - const dispose = this.ctx.effect(function* (this: ToolRegistry) { - const layer = scope === undefined ? this.globalGuards : this.guardLayerFor(scope) - layer.add(registration) - yield () => { - layer.delete(registration) - if (scope !== undefined && layer.size === 0) this.scopedGuards.delete(scope) - } - }.bind(this), 'tools.guard()') - // eslint-disable-next-line @typescript-eslint/no-misused-promises -- synchronous cleanup; direct return preserves disposer identity - return dispose - } - - /** The (created-on-demand) scoped layer for `scope`. */ - private layerFor(scope: ScopeKey): Map { - let layer = this.scoped.get(scope) - if (!layer) { - layer = new Map() - this.scoped.set(scope, layer) - } - return layer - } - - /** Get or create the guard layer for one agent scope. */ - private guardLayerFor(scope: ScopeKey): Set { - let layer = this.scopedGuards.get(scope) - if (layer === undefined) { - layer = new Set() - this.scopedGuards.set(scope, layer) - } - return layer + return this.layers.effect( + this.ctx, + layer => layer.guards.append(guard), + { label: 'tools.guard()', notify: false }, + ) } /** First monotonic denial from the global then matching scoped guard layers. */ private guardReason(exec: ToolExecution): string | undefined { - for (const { guard } of this.globalGuards) { - const reason = guard(exec) - if (reason !== undefined) return reason - } - if (exec.agent !== undefined) { - for (const { guard } of this.scopedGuards.get(exec.agent) ?? []) { - const reason = guard(exec) - if (reason !== undefined) return reason - } - } - return undefined - } - - /** Whether every restriction registered for `scope` admits the global tool `name` (intersection semantics). */ - private admits(scope: ScopeKey | undefined, name: string): boolean { - if (scope === undefined) return true - const filters = this.restrictions.get(scope) - if (!filters) return true - return filters.every(filter => - (filter.allow === undefined || filter.allow.has(name)) - && (filter.deny === undefined || !filter.deny.has(name))) + const globalReason = this.layers.global.guardReason(exec) + if (globalReason !== undefined) return globalReason + return exec.agent === undefined ? undefined : this.layers.peek(exec.agent)?.guardReason(exec) } /** @@ -752,18 +708,18 @@ export class ToolRegistry extends Service { * @returns the complete derived view for that scope. */ private view(scope?: ScopeKey): ToolView { - const layer = scope === undefined ? undefined : this.scoped.get(scope) + const layer = this.layers.peek(scope) const visible = new Map() const knownNames = new Set() const restrictableNames = new Set() - for (const [name, definition] of this.global) { + for (const [name, definition] of this.layers.global.tools.entries()) { knownNames.add(name) restrictableNames.add(name) - if (this.admits(scope, name)) visible.set(name, definition) + if (layer?.admits(name) ?? true) visible.set(name, definition) } // Scoped layer second: same-name entries REPLACE (shadow) the global ones, // and scope-local registrations are never part of the global filter above. - for (const [name, definition] of layer ?? []) { + for (const [name, definition] of layer?.tools.entries() ?? []) { knownNames.add(name) visible.set(name, definition) } diff --git a/packages/core/tools/tests/scoped.spec.ts b/packages/core/tools/tests/scoped.spec.ts index ab3b869953..cab0ee7155 100644 --- a/packages/core/tools/tests/scoped.spec.ts +++ b/packages/core/tools/tests/scoped.spec.ts @@ -266,6 +266,27 @@ describe('scoped execution dispatch', () => { expect(bodyCalls).toBe(0) }) + it('live-iterates a guard registered by an earlier guard', async () => { + const ctx = await mount() + const calls: string[] = [] + let added = false + ctx.tools.register(tool('t')) + ctx.tools.guard(() => { + calls.push('first') + if (!added) { + added = true + ctx.tools.guard(() => { + calls.push('late') + return 'late denial' + }) + } + return undefined + }) + + expect(await run(ctx, 't')).toBe('Error: late denial') + expect(calls).toEqual(['first', 'late']) + }) + it('shares one token and materialized argument value across the pipeline', async () => { const ctx = await mount() const { scope, key } = await mintAgentScope(ctx, 'a')