File size: 4,679 Bytes
68d7816 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | import { type CollectionView } from '#/_base/di/collection';
import { IInstantiationService } from '#/_base/di/instantiation';
import { type IDisposable } from '#/_base/di/lifecycle';
import { Service } from '#/_base/di/service';
import { LifecycleScope } from '#/app/scopes';
import { ScopeActivation, registerScopedService } from '#/_base/di/scope';
import { IEventBus } from '#/app/event/eventBus';
import { IAgentProfileService } from '#/agent/profile/profile';
import { AgentStatusUpdated } from '#/agent/usage/usageEvents';
import { isToolActive } from '#/agent/toolPolicy/evaluate';
import { SELECT_TOOLS_TOOL_NAME } from '#/agent/toolSelect/toolSelect';
import { IAgentToolRegistryService } from '#/agent/toolRegistry/toolRegistry';
import { AgentToolContribution } from '#/agent/toolRegistry/toolContribution';
import { ISessionToolPolicyGate } from '#/session/sessionToolPolicyGate/sessionToolPolicyGate';
import { IAgentRuntimeService } from '#/agent/runtimeBinding/agentRuntime';
import { IAgentToolActivationService } from './toolActivation';
export class AgentToolActivationService extends Service implements IAgentToolActivationService {
declare readonly _serviceBrand: undefined;
private readonly registrations = new Map<AgentToolContribution, IDisposable>();
constructor(
@IInstantiationService private readonly instantiationService: IInstantiationService,
@IAgentToolRegistryService private readonly toolRegistry: IAgentToolRegistryService,
@IAgentProfileService private readonly profile: IAgentProfileService,
@ISessionToolPolicyGate private readonly toolPolicyGate: ISessionToolPolicyGate,
@IAgentRuntimeService private readonly runtime: IAgentRuntimeService,
@IEventBus eventBus: IEventBus,
@AgentToolContribution private readonly contributions: CollectionView<AgentToolContribution>,
) {
super();
this._register(
eventBus.subscribe(AgentStatusUpdated, () => {
void this.activate();
}),
);
this._register(this.runtime.onDidChange(() => this.refreshRuntimeRecords()));
this._register(
this.contributions.onDidChange((change) => {
this.activateRecords(change.added);
for (const record of change.removed) {
this.deactivateRecord(record);
}
}),
);
}
activate(): Promise<void> {
this.activateRecords(this.contributions.items);
return Promise.resolve();
}
private activateRecords(records: readonly AgentToolContribution[]): void {
if (records.length === 0) return;
const data = this.profile.data();
const policy = { tools: data.activeToolNames, disallowedTools: data.disallowedTools };
const disclosurePolicy = { disallowedTools: data.disallowedTools };
const workspaceVeto = { disallowedTools: this.toolPolicyGate.disabledTools };
this.instantiationService.invokeFunction((accessor) => {
for (const record of records) {
const { id, options } = record;
const source = options.source ?? 'builtin';
if (this.toolRegistry.resolve(options.name) !== undefined) continue;
if (!this.runtimeAllows(record)) continue;
if (!isToolActive(workspaceVeto, options.name, source)) continue;
const activeByProfile =
options.name === SELECT_TOOLS_TOOL_NAME
? isToolActive(disclosurePolicy, options.name, source)
: isToolActive(policy, options.name, source);
if (!activeByProfile) continue;
if (options.when !== undefined && !options.when(accessor)) continue;
const tool = accessor.get(id);
const registration = this.toolRegistry.register(tool, {
source: options.source,
disclosure: options.disclosure,
});
this.registrations.set(record, registration);
this._register(registration);
}
});
}
private refreshRuntimeRecords(): void {
for (const record of this.contributions.items) {
if (!this.runtimeAllows(record)) this.deactivateRecord(record);
}
this.activateRecords(this.contributions.items);
}
private runtimeAllows(record: AgentToolContribution): boolean {
const required = record.options.requiredRuntimeCapabilities;
return required === undefined || this.runtime.isAvailable(required);
}
private deactivateRecord(record: AgentToolContribution): void {
const registration = this.registrations.get(record);
if (registration === undefined) return;
this.registrations.delete(record);
registration.dispose();
}
}
registerScopedService(
LifecycleScope.Agent,
IAgentToolActivationService,
AgentToolActivationService,
ScopeActivation.OnScopeCreated,
'toolActivation',
);
|