File size: 4,006 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 | import { ScopeActivation, registerScopedService } from '#/_base/di/scope';
import { defineState } from '#/state/state';
import type { IDisposable } from '#/_base/di/lifecycle';
import { Emitter } from '#/_base/event';
import { LifecycleScope } from '#/app/scopes';
import { IAgentScopeContext } from '#/agent/scopeContext/scopeContext';
import { IAgentStateService } from '#/agent/state/agentState';
import type { RuntimeBinding } from '#/runtime/runtime';
import { RuntimeError } from '#/runtime/runtimeRegistry';
import { ISessionContext } from '#/session/sessionContext/sessionContext';
import { IEventDispatcher } from '#/state/eventDispatcher';
import { IRuntimeResolver } from '#/workspace/workspaceInstance/workspaceInstanceManager';
import { IAgentRuntimeBindingSeed, IAgentRuntimeBindingService } from './runtimeBinding';
import { RuntimeSetBinding, runtimeBindingKey } from './runtimeBindingOps';
export const agentRuntimeBindingKey = defineState<RuntimeBinding>('runtime.binding', () => ({ workspaceId: '', runtimeId: 'local' }));
export class AgentRuntimeBindingService implements IAgentRuntimeBindingService {
declare readonly _serviceBrand: undefined;
private readonly changeEmitter = new Emitter<RuntimeBinding>();
readonly onDidChange = this.changeEmitter.event;
private readonly restoreHook: IDisposable;
constructor(
@IAgentScopeContext private readonly scopeContext: IAgentScopeContext,
@IAgentStateService private readonly state: IAgentStateService,
@IAgentRuntimeBindingSeed seed: IAgentRuntimeBindingSeed,
@ISessionContext private readonly session: ISessionContext,
@IRuntimeResolver private readonly resolver: IRuntimeResolver,
@IEventDispatcher private readonly dispatcher: IEventDispatcher,
) {
this.state.contributeState(agentRuntimeBindingKey);
this.state.contributeState(runtimeBindingKey);
const initial = this.state.get(runtimeBindingKey) ?? seed.binding;
this.assertSessionWorkspace(initial);
this.state.set(agentRuntimeBindingKey, initial);
this.restoreHook = dispatcher.hooks.onDidRestore.register('agent-runtime-binding', async (_ctx, next) => {
const replayed = this.state.get(runtimeBindingKey);
if (replayed === undefined) {
void this.dispatcher.dispatch(
new RuntimeSetBinding({ ...this.current, agentId: this.scopeContext.agentId }),
);
} else {
this.assertSessionWorkspace(replayed);
this.state.set(agentRuntimeBindingKey, replayed);
}
await next();
});
}
private assertSessionWorkspace(binding: RuntimeBinding): void {
if (binding.workspaceId !== this.session.workspaceId) {
throw new RuntimeError(
'runtime.not_found',
`runtime binding workspace ${binding.workspaceId} does not match session workspace ${this.session.workspaceId}`,
);
}
}
get current(): RuntimeBinding {
return this.state.get(agentRuntimeBindingKey);
}
get(): RuntimeBinding {
return this.current;
}
set(binding: RuntimeBinding): RuntimeBinding {
this.assertSessionWorkspace(binding);
const lease = this.resolver.acquire(binding, []);
lease.dispose();
if (binding.workspaceId === this.current.workspaceId && binding.runtimeId === this.current.runtimeId) {
return this.current;
}
const next = { workspaceId: binding.workspaceId, runtimeId: binding.runtimeId };
void this.dispatcher.dispatch(
new RuntimeSetBinding({ ...next, agentId: this.scopeContext.agentId }),
);
this.state.set(agentRuntimeBindingKey, next);
this.changeEmitter.fire(next);
return next;
}
switch(runtimeId: string): RuntimeBinding {
return this.set({ workspaceId: this.session.workspaceId, runtimeId });
}
dispose(): void {
this.restoreHook.dispose();
this.changeEmitter.dispose();
}
}
registerScopedService(LifecycleScope.Agent, IAgentRuntimeBindingService, AgentRuntimeBindingService, ScopeActivation.OnScopeCreated, 'agentRuntimeBinding');
|