diff --git a/packages/agent-core-v2/src/_base/di/collection.ts b/packages/agent-core-v2/src/_base/di/collection.ts index 4f69dedbba..75a96e6509 100644 --- a/packages/agent-core-v2/src/_base/di/collection.ts +++ b/packages/agent-core-v2/src/_base/di/collection.ts @@ -36,8 +36,15 @@ export interface CollectionToken { const _collectionTokens = new Map>(); const _collectionTokenSet = new WeakSet(); +const _collectionValidators = new WeakMap< + object, + (value: unknown, existing: readonly unknown[]) => void +>(); -export function collection(name: string): CollectionToken { +export function collection( + name: string, + options: { readonly validate?: (value: T, existing: readonly T[]) => void } = {}, +): CollectionToken { const existing = _collectionTokens.get(name); if (existing !== undefined) { return existing as CollectionToken; @@ -61,6 +68,12 @@ export function collection(name: string): CollectionToken { Object.defineProperty(token, 'name', { value: name, enumerable: false, configurable: true }); _collectionTokens.set(name, token as CollectionToken); _collectionTokenSet.add(token); + if (options.validate !== undefined) { + _collectionValidators.set( + token, + options.validate as (value: unknown, existing: readonly unknown[]) => void, + ); + } return token; } @@ -119,6 +132,10 @@ export class CollectionStore { records = new Map(); this._records.set(token as CollectionToken, records); } + _collectionValidators.get(token)?.( + value, + [...records.values()].map((entry) => entry.value), + ); const record: StoredRecord = { id: ++this._nextId, value, diff --git a/packages/agent-core-v2/src/app/feature/featureManager.ts b/packages/agent-core-v2/src/app/feature/featureManager.ts index 6694d6bb3b..be2314df85 100644 --- a/packages/agent-core-v2/src/app/feature/featureManager.ts +++ b/packages/agent-core-v2/src/app/feature/featureManager.ts @@ -26,6 +26,7 @@ import type { ServiceRecipe, } from '#/_base/di/fiber'; import { createDecorator, type ServiceIdentifier } from '#/_base/di/instantiation'; +import type { ContributedFeatureService } from './featureServiceContribution'; export interface ManagedUnitInfo { readonly name: string; @@ -46,6 +47,7 @@ export interface IFeatureManager { updateUnit(name: string, config?: unknown): Promise; units(): readonly ManagedUnitInfo[]; + contributedServices(): readonly ContributedFeatureService[]; readonly onDidChangeUnits: Event; } diff --git a/packages/agent-core-v2/src/app/feature/featureManagerService.ts b/packages/agent-core-v2/src/app/feature/featureManagerService.ts index 7e521c7201..04e8e3e2d0 100644 --- a/packages/agent-core-v2/src/app/feature/featureManagerService.ts +++ b/packages/agent-core-v2/src/app/feature/featureManagerService.ts @@ -8,6 +8,7 @@ * the previous handle (retract-then-assemble is the caller's cascade). */ +import type { CollectionView } from '#/_base/di/collection'; import { Emitter, type Event } from '#/_base/event'; import type { FiberHandle, @@ -23,6 +24,10 @@ import { IFeatureManager, type ManagedUnitInfo, } from './featureManager'; +import { + FeatureServiceContribution, + type ContributedFeatureService, +} from './featureServiceContribution'; export class FeatureManagerService extends Service implements IFeatureManager { declare readonly _serviceBrand: undefined; @@ -31,7 +36,10 @@ export class FeatureManagerService extends Service implements IFeatureManager { private readonly _onDidChangeUnits = new Emitter(); readonly onDidChangeUnits: Event = this._onDidChangeUnits.event; - constructor() { + constructor( + @FeatureServiceContribution + private readonly _contributedServices: CollectionView, + ) { super(); this._register(this._onDidChangeUnits); } @@ -97,6 +105,10 @@ export class FeatureManagerService extends Service implements IFeatureManager { } return infos; } + + contributedServices(): readonly ContributedFeatureService[] { + return this._contributedServices.items; + } } registerScopedService( diff --git a/packages/agent-core-v2/src/app/feature/featureServiceContribution.ts b/packages/agent-core-v2/src/app/feature/featureServiceContribution.ts new file mode 100644 index 0000000000..635b152a4e --- /dev/null +++ b/packages/agent-core-v2/src/app/feature/featureServiceContribution.ts @@ -0,0 +1,21 @@ +import { collection } from '#/_base/di/collection'; +import type { ServiceIdentifier } from '#/_base/di/instantiation'; +import type { LifecycleScope } from '#/app/scopes'; + +export interface ContributedFeatureService { + readonly scope: LifecycleScope; + readonly id: ServiceIdentifier; +} + +export const FeatureServiceContribution = collection( + 'feature-service', + { + validate(value, existing) { + if (existing.some((entry) => entry.scope === value.scope && entry.id === value.id)) { + throw new Error( + `Service ${String(value.id)} is already contributed at scope ${value.scope}`, + ); + } + }, + }, +); diff --git a/packages/agent-core-v2/src/features/feature.ts b/packages/agent-core-v2/src/features/feature.ts index b62a6666b0..0c5e7b91fb 100644 --- a/packages/agent-core-v2/src/features/feature.ts +++ b/packages/agent-core-v2/src/features/feature.ts @@ -28,6 +28,7 @@ import { AgentProfileContribution, AGENT_PROFILE_SOURCE_PRIORITY, } from '#/app/agentProfileCatalog/agentProfileContribution'; +import { FeatureServiceContribution } from '#/app/feature/featureServiceContribution'; import type { AgentProfile } from '#/app/agentProfileCatalog/agentProfileCatalog'; import type { ConfigSchema, RegisterSectionOptions } from '#/app/config/config'; import { ConfigSectionContribution } from '#/app/config/configSectionContributions'; @@ -43,8 +44,6 @@ import { type AnyAgentTool, } from '#/agent/toolRegistry/toolContribution'; -import { recordContributedService } from './featureRegistry'; - export abstract class Feature extends Service { contribute(token: CollectionToken, value: T): FiberHandle { return this.provide(token, value); @@ -68,7 +67,7 @@ export abstract class Feature extends Service { ctor: ServiceClassRecipe, opts?: FiberProvideOptions, ): FiberHandle { - recordContributedService(scope, id); + this.provide(FeatureServiceContribution, { scope, id }); return this.provide(ScopeUnits(scope), { name: `${this.name}:${String(id)}`, apply(fiber: Fiber): void { diff --git a/packages/agent-core-v2/src/features/featureRegistry.ts b/packages/agent-core-v2/src/features/featureRegistry.ts index b268a473b8..b08b26412c 100644 --- a/packages/agent-core-v2/src/features/featureRegistry.ts +++ b/packages/agent-core-v2/src/features/featureRegistry.ts @@ -11,7 +11,6 @@ */ import type { ServiceClassRecipe } from '#/_base/di/fiber'; -import type { ServiceIdentifier } from '#/_base/di/instantiation'; const _featureRecipes: ServiceClassRecipe[] = []; @@ -26,23 +25,3 @@ export function getFeatureRecipes(): readonly ServiceClassRecipe[] { export function _clearFeatureRecipesForTests(): void { _featureRecipes.length = 0; } - -const _contributedServices: { scope: string; id: ServiceIdentifier }[] = []; - -export function recordContributedService(scope: string, id: ServiceIdentifier): void { - if (_contributedServices.some((entry) => entry.scope === scope && entry.id === id)) { - return; - } - _contributedServices.push({ scope, id }); -} - -export function getContributedServices(): ReadonlyArray<{ - scope: string; - id: ServiceIdentifier; -}> { - return _contributedServices; -} - -export function _clearContributedServicesForTests(): void { - _contributedServices.length = 0; -} diff --git a/packages/agent-core-v2/src/index.ts b/packages/agent-core-v2/src/index.ts index 33e343ba40..886462ba45 100644 --- a/packages/agent-core-v2/src/index.ts +++ b/packages/agent-core-v2/src/index.ts @@ -215,6 +215,7 @@ export * from '#/app/capability/capabilityService'; export * from '#/app/capability/errors'; export * from '#/app/capability/types'; export * from '#/app/feature/featureManager'; +export * from '#/app/feature/featureServiceContribution'; import '#/app/feature/featureManagerService'; export * from '#/features/feature'; export * from '#/features/featureAssembly'; diff --git a/packages/agent-core-v2/test/features/debugEvents/debugEvents.test.ts b/packages/agent-core-v2/test/features/debugEvents/debugEvents.test.ts index 7d6457a3d6..38bee5b9f6 100644 --- a/packages/agent-core-v2/test/features/debugEvents/debugEvents.test.ts +++ b/packages/agent-core-v2/test/features/debugEvents/debugEvents.test.ts @@ -14,7 +14,6 @@ import { IFeatureAssemblyService } from '#/features/featureAssembly'; import { FeatureAssemblyService } from '#/features/featureAssemblyService'; import { _clearFeatureRecipesForTests, - getContributedServices, registerFeature, } from '#/features/featureRegistry'; @@ -53,9 +52,9 @@ describe('DebugEventsFeature — App-scope introspection service', () => { const manager = host.app.accessor.get(IFeatureManager); expect(manager.units().map((unit) => unit.name)).toContain('debugEvents'); expect( - getContributedServices().some( - (entry) => entry.scope === LifecycleScope.App && entry.id === IDebugEventsService, - ), + manager + .contributedServices() + .some((entry) => entry.scope === LifecycleScope.App && entry.id === IDebugEventsService), ).toBe(true); const result = host.app.accessor.get(IDebugEventsService).subscriptions(); @@ -68,6 +67,11 @@ describe('DebugEventsFeature — App-scope introspection service', () => { await host.app.instantiation.cascade.whenIdle(); await new Promise((resolve) => setTimeout(resolve, 0)); expect(() => host.app.accessor.get(IDebugEventsService)).toThrow(); + expect( + manager + .contributedServices() + .some((entry) => entry.scope === LifecycleScope.App && entry.id === IDebugEventsService), + ).toBe(false); host.dispose(); }); }); diff --git a/packages/agent-core-v2/test/features/feature.test.ts b/packages/agent-core-v2/test/features/feature.test.ts index 13e25c0942..8211968d0a 100644 --- a/packages/agent-core-v2/test/features/feature.test.ts +++ b/packages/agent-core-v2/test/features/feature.test.ts @@ -6,6 +6,7 @@ import { createDecorator, ScopeActivation } from '#/_base/di/instantiation'; import { type InstantiationService } from '#/_base/di/instantiationService'; import { _clearScopedRegistryForTests, + getScopedServiceDescriptors, registerScopedService, type Scope, } from '#/_base/di/scope'; @@ -20,7 +21,10 @@ import { AgentToolContribution } from '#/agent/toolRegistry/toolContribution'; import { Feature } from '#/features/feature'; import { IFeatureAssemblyService } from '#/features/featureAssembly'; import { FeatureAssemblyService } from '#/features/featureAssemblyService'; -import { _clearFeatureRecipesForTests, registerFeature } from '#/features/featureRegistry'; +import { + _clearFeatureRecipesForTests, + registerFeature, +} from '#/features/featureRegistry'; import type { AgentTool, ToolExecution } from '#/tool/toolContract'; interface IGreeter { @@ -138,6 +142,82 @@ describe('Feature — built-in capability assembly (src/features)', () => { host.dispose(); }); + it('rejects duplicate service contributions until the provider unloads', async () => { + class FirstFeature extends Feature { + static override readonly name = 'first-feature'; + + constructor() { + super(); + this.contributeAgentService(IGreeter, GreeterService); + } + } + class SecondFeature extends Feature { + static override readonly name = 'second-feature'; + + constructor() { + super(); + this.contributeAgentService(IGreeter, GreeterService); + } + } + registerFeature(FirstFeature); + + const host = createScopedTestHost(); + const manager = host.app.accessor.get(IFeatureManager); + const agent = host.child(LifecycleScope.Agent, 'agent-1'); + const original = agent.accessor.get(IGreeter); + expect(() => manager.provideUnit(SecondFeature)).toThrow( + /Service test-feature-greeter is already contributed at scope agent/, + ); + expect(manager.units().map((unit) => unit.name)).toEqual(['first-feature']); + expect( + manager + .contributedServices() + .filter((entry) => entry.scope === LifecycleScope.Agent && entry.id === IGreeter), + ).toHaveLength(1); + expect(collectionViewOf(host.app, ScopeUnits(LifecycleScope.Agent)).items).toHaveLength(1); + expect(agent.accessor.get(IGreeter)).toBe(original); + + await manager.unprovideUnit('first-feature'); + await host.app.instantiation.cascade.whenIdle(); + expect(() => manager.provideUnit(SecondFeature)).not.toThrow(); + expect(manager.units().map((unit) => unit.name)).toEqual(['second-feature']); + expect(agent.accessor.get(IGreeter).greet()).toBe('hi'); + + host.dispose(); + }); + + it('isolates equal service contributions between App roots', async () => { + class SharedFeature extends Feature { + static override readonly name = 'shared-feature'; + + constructor() { + super(); + this.contributeAgentService(IGreeter, GreeterService); + } + } + registerFeature(SharedFeature); + + const first = createScopedTestHost(); + const second = createScopedTestHost(); + const firstManager = first.app.accessor.get(IFeatureManager); + const secondManager = second.app.accessor.get(IFeatureManager); + const firstAgent = first.child(LifecycleScope.Agent, 'agent-1'); + const secondAgent = second.child(LifecycleScope.Agent, 'agent-1'); + expect(firstManager.contributedServices()).toHaveLength(1); + expect(secondManager.contributedServices()).toHaveLength(1); + expect(firstAgent.accessor.get(IGreeter)).not.toBe(secondAgent.accessor.get(IGreeter)); + + await firstManager.unprovideUnit('shared-feature'); + await first.app.instantiation.cascade.whenIdle(); + expect(firstManager.contributedServices()).toHaveLength(0); + expect(secondManager.contributedServices()).toHaveLength(1); + expect(() => firstAgent.accessor.get(IGreeter)).toThrow(); + expect(secondAgent.accessor.get(IGreeter).greet()).toBe('hi'); + + first.dispose(); + second.dispose(); + }); + it('materializes a per-scope class recipe contributed through contribute()', () => { class SoloAgentUnit extends Service { static override readonly name = 'solo-feature/agent'; diff --git a/packages/agent-core-v2/test/harness/agent.ts b/packages/agent-core-v2/test/harness/agent.ts index 4de0ba782e..c75f46ffe0 100644 --- a/packages/agent-core-v2/test/harness/agent.ts +++ b/packages/agent-core-v2/test/harness/agent.ts @@ -8,7 +8,7 @@ import { expect, vi } from 'vitest'; import { toDisposable } from '#/_base/di/lifecycle'; import type { IInstantiationService } from '#/_base/di/instantiation'; import type { IAgentScopeHandle } from '#/_base/di/scope'; -import { getContributedServices } from '#/features/featureRegistry'; +import { IFeatureManager } from '#/app/feature/featureManager'; import { Emitter, Event } from '#/_base/event'; import { IAgentLifecycleService } from '#/session/agentLifecycle/agentLifecycle'; import type { Promisable, PromisifyMethods } from '#/_base/utils/types'; @@ -910,7 +910,9 @@ function reassertServiceOverrides( instantiation: IInstantiationService, ): void { const contributed = new Set( - getContributedServices() + instantiation + .invokeFunction((accessor) => accessor.get(IFeatureManager)) + .contributedServices() .filter((entry) => entry.scope === scope) .map((entry) => entry.id), ); diff --git a/packages/kap-server/src/transport/channelRegistry.ts b/packages/kap-server/src/transport/channelRegistry.ts index 58cd41e76d..56c2268b51 100644 --- a/packages/kap-server/src/transport/channelRegistry.ts +++ b/packages/kap-server/src/transport/channelRegistry.ts @@ -15,12 +15,12 @@ import { Disposable, - getContributedServices, getScopedServiceDescriptors, + IFeatureManager, LifecycleScope, } from '@moonshot-ai/agent-core-v2'; -import type { ScopedEntry, ServiceIdentifier } from '@moonshot-ai/agent-core-v2'; +import type { Scope, ScopedEntry, ServiceIdentifier } from '@moonshot-ai/agent-core-v2'; export interface ChannelMethodDescriptor { readonly name: string; @@ -87,10 +87,16 @@ function scopedServiceNameIndex(): Map> { } /** Resolve a wire name to its `ServiceIdentifier` anywhere in the DI registry. */ -export function resolveAnyScopedServiceId(name: string): ServiceIdentifier | undefined { +export function resolveAnyScopedServiceId( + core: Scope, + name: string, +): ServiceIdentifier | undefined { return ( scopedServiceNameIndex().get(name) ?? - getContributedServices().find((entry) => entry.id.toString() === name)?.id + core.accessor + .get(IFeatureManager) + .contributedServices() + .find((entry) => entry.id.toString() === name)?.id ); } diff --git a/packages/kap-server/src/transport/dispatcher.ts b/packages/kap-server/src/transport/dispatcher.ts index e3b75dba49..94b6e3f804 100644 --- a/packages/kap-server/src/transport/dispatcher.ts +++ b/packages/kap-server/src/transport/dispatcher.ts @@ -87,7 +87,7 @@ export async function resolveService( scopeKind: ScopeKind, params: Record, serviceName: string, - lookup: ChannelLookup = resolveAnyScopedServiceId, + lookup: ChannelLookup = (name) => resolveAnyScopedServiceId(core, name), ): Promise { const scope = await resolveScope(core, scopeKind, params); if (scope === undefined) { @@ -128,7 +128,7 @@ export async function dispatch( serviceName: string, method: string, arg: unknown, - lookup: ChannelLookup = resolveAnyScopedServiceId, + lookup: ChannelLookup = (name) => resolveAnyScopedServiceId(core, name), ): Promise { const service = await resolveService(core, scopeKind, params, serviceName, lookup); const member = (service as Record)[method]; diff --git a/packages/kap-server/src/transport/registerDebugRoutes.ts b/packages/kap-server/src/transport/registerDebugRoutes.ts index 11ec8e57ca..6cdb3b9af2 100644 --- a/packages/kap-server/src/transport/registerDebugRoutes.ts +++ b/packages/kap-server/src/transport/registerDebugRoutes.ts @@ -21,7 +21,7 @@ import { type RouteHost, registerServiceDispatcherRoutes } from './serviceDispat export function registerDebugRoutes(app: RouteHost, core: Scope): void { registerServiceDispatcherRoutes(app, core, '/debug', { - lookup: resolveAnyScopedServiceId, + lookup: (name) => resolveAnyScopedServiceId(core, name), describe: describeAllChannels, }); } diff --git a/packages/kap-server/src/transport/serviceDispatcherRoutes.ts b/packages/kap-server/src/transport/serviceDispatcherRoutes.ts index 59327f1c15..c325dec9c2 100644 --- a/packages/kap-server/src/transport/serviceDispatcherRoutes.ts +++ b/packages/kap-server/src/transport/serviceDispatcherRoutes.ts @@ -74,7 +74,7 @@ export function registerServiceDispatcherRoutes( basePath: string, opts: ServiceDispatcherRouteOptions = {}, ): void { - const lookup = opts.lookup ?? resolveAnyScopedServiceId; + const lookup = opts.lookup ?? ((name) => resolveAnyScopedServiceId(core, name)); const scopeRoutes: { path: string; scopeKind: ScopeKind }[] = [ { path: `${basePath}/:service/:method`, scopeKind: 'core' }, { path: `${basePath}/workspace/:workspace_id/:service/:method`, scopeKind: 'workspace' }, diff --git a/packages/kap-server/test/channelRegistry.test.ts b/packages/kap-server/test/channelRegistry.test.ts new file mode 100644 index 0000000000..d81126a375 --- /dev/null +++ b/packages/kap-server/test/channelRegistry.test.ts @@ -0,0 +1,42 @@ +import { + createDecorator, + Feature, + IFeatureManager, + LifecycleScope, + Service, + createAppScope, +} from '@moonshot-ai/agent-core-v2'; +import { describe, expect, it } from 'vitest'; + +import { resolveAnyScopedServiceId } from '../src/transport/channelRegistry'; + +describe('channelRegistry', () => { + it('resolves contributed services from the current core only', async () => { + const id = createDecorator('test-contributed-service'); + class TestService extends Service {} + class TestFeature extends Feature { + constructor() { + super(); + this.contributeService(LifecycleScope.Agent, id, TestService); + } + } + const first = createAppScope(); + const second = createAppScope(); + const firstManager = first.accessor.get(IFeatureManager); + const secondManager = second.accessor.get(IFeatureManager); + firstManager.provideUnit(TestFeature); + secondManager.provideUnit(TestFeature); + + expect(resolveAnyScopedServiceId(first, String(id))).toBe(id); + expect(resolveAnyScopedServiceId(second, String(id))).toBe(id); + + await firstManager.unprovideUnit('TestFeature'); + expect(resolveAnyScopedServiceId(first, String(id))).toBeUndefined(); + expect(resolveAnyScopedServiceId(second, String(id))).toBe(id); + + await secondManager.unprovideUnit('TestFeature'); + expect(resolveAnyScopedServiceId(second, String(id))).toBeUndefined(); + first.dispose(); + second.dispose(); + }); +});