Skip to content

Commit 8a54c78

Browse files
V3RONV3RON
committed
feat(leasing): proxy CLI names each new lease, daemon keeps a caller-chosen lease id
1 parent 7da2934 commit 8a54c78

14 files changed

Lines changed: 284 additions & 4 deletions

File tree

‎packages/contracts/src/__tests__/lease-scope.test.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -130,6 +130,7 @@ test('leaseScopeToLeaseRpcParams projects canonical provider and command-specifi
130130
leaseProvider: 'proxy',
131131
clientId: 'client-a',
132132
deviceKey: 'android:emulator-5554',
133+
leaseId: 'lease-1',
133134
ttlMs: 60_000,
134135
backend: 'android-instance',
135136
},

‎packages/contracts/src/client-lease.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,12 @@ export type LeaseOptions = AgentDeviceRequestOverrides &
2929
export type LeaseAllocateOptions = LeaseOptions & {
3030
tenant: string;
3131
runId: string;
32+
/**
33+
* The id for the new lease (16-128 hex characters); the client's configured `leaseId` is never
34+
* sent here. Use the returned lease's `leaseId`: a daemon that predates this field mints its own,
35+
* and a run that already holds a lease gets that lease back.
36+
*/
37+
leaseId?: string;
3238
leaseBackend?: LeaseBackend;
3339
leaseProvider?: string;
3440
provider?: string;

‎packages/contracts/src/lease-scope.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ export type LeaseRpcCommand = 'lease_allocate' | 'lease_heartbeat' | 'lease_rele
3232
export type LeaseAllocateRequestScope = {
3333
tenantId: string;
3434
runId: string;
35+
leaseId?: string;
3536
leaseBackend?: LeaseBackend;
3637
leaseProvider?: string;
3738
deviceKey?: string;
@@ -141,6 +142,7 @@ export function leaseScopeToAllocateRequest(scope: LeaseScope): LeaseAllocateReq
141142
return stripUndefined({
142143
tenantId: scope.tenantId ?? '',
143144
runId: scope.runId ?? '',
145+
leaseId: scope.leaseId,
144146
leaseBackend: scope.leaseBackend,
145147
leaseProvider: scope.leaseProvider,
146148
deviceKey: scope.deviceKey,
@@ -184,6 +186,7 @@ export function leaseScopeToLeaseRpcParams(
184186
return {
185187
...common,
186188
...stripUndefined({
189+
leaseId: scope.leaseId,
187190
ttlMs: scope.leaseTtlMs,
188191
retainOnClose: scope.leaseRetainOnClose,
189192
backend: scope.leaseBackend,

‎src/__tests__/remote-connection-proxy-lease-id.test.ts‎

Lines changed: 140 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,140 @@
1+
import { test, vi } from 'vitest';
2+
import assert from 'node:assert/strict';
3+
import fs from 'node:fs';
4+
import { AppError } from '@agent-device/kernel/errors';
5+
import {
6+
connectionWorkspace,
7+
createTestClient,
8+
type LeaseAllocateRequest,
9+
seedConnectionState,
10+
} from './remote-connection.fixtures.ts';
11+
12+
vi.mock('../metro/client-metro-companion.ts', () => ({
13+
stopMetroCompanion: vi.fn(),
14+
}));
15+
16+
vi.mock('../client/client-react-devtools-companion.ts', () => ({
17+
stopReactDevtoolsCompanion: vi.fn(),
18+
}));
19+
20+
import { materializeRemoteConnectionForCommand } from '../cli/commands/connection-runtime.ts';
21+
import { readRemoteConnectionState } from '../remote/remote-connection-state.ts';
22+
23+
const DAEMON_LEASE_ID = 'abc123abc123abc1';
24+
25+
function recordingAllocate(requests: LeaseAllocateRequest[]) {
26+
return async (request: LeaseAllocateRequest) => {
27+
requests.push(request);
28+
// The daemon's own id, as an older daemon that ignores the requested one answers.
29+
return {
30+
leaseId: DAEMON_LEASE_ID,
31+
tenantId: request.tenant,
32+
runId: request.runId,
33+
backend: request.leaseBackend ?? 'android-instance',
34+
leaseProvider: request.leaseProvider,
35+
clientId: request.clientId,
36+
deviceKey: request.deviceKey,
37+
};
38+
};
39+
}
40+
41+
test('proxy allocation sends a fresh lease id on every attempt and keeps the id the daemon returns', async () => {
42+
const { tempRoot, stateDir, remoteConfigPath } = connectionWorkspace(
43+
'agent-device-proxy-lease-id-',
44+
);
45+
fs.writeFileSync(remoteConfigPath, JSON.stringify({ daemonBaseUrl: 'https://daemon.example' }));
46+
seedConnectionState({
47+
stateDir,
48+
state: {
49+
session: 'adc-proxy',
50+
remoteConfigPath,
51+
daemon: { baseUrl: 'https://daemon.example' },
52+
tenant: 'proxy',
53+
runId: 'proxy-client-1',
54+
leaseProvider: 'proxy',
55+
clientId: 'client-1',
56+
},
57+
});
58+
const requests: LeaseAllocateRequest[] = [];
59+
const client = createTestClient({
60+
allocate: recordingAllocate(requests),
61+
heartbeat: async () => {
62+
throw new AppError('UNAUTHORIZED', 'Lease expired', { reason: 'LEASE_EXPIRED' });
63+
},
64+
});
65+
const flags = {
66+
json: true,
67+
help: false,
68+
version: false,
69+
stateDir,
70+
remoteConfig: remoteConfigPath,
71+
daemonBaseUrl: 'https://daemon.example',
72+
tenant: 'proxy',
73+
runId: 'proxy-client-1',
74+
session: 'adc-proxy',
75+
platform: 'android',
76+
} as const;
77+
78+
const first = await materializeRemoteConnectionForCommand({
79+
command: 'install',
80+
flags: { ...flags },
81+
client,
82+
});
83+
const second = await materializeRemoteConnectionForCommand({
84+
command: 'install',
85+
flags: { ...flags },
86+
client,
87+
});
88+
89+
assert.equal(requests.length, 2);
90+
assert.match(requests[0]?.leaseId ?? '', /^[0-9a-f]{32}$/);
91+
assert.match(requests[1]?.leaseId ?? '', /^[0-9a-f]{32}$/);
92+
assert.notEqual(requests[0]?.leaseId, requests[1]?.leaseId);
93+
assert.equal(first.flags.leaseId, DAEMON_LEASE_ID);
94+
assert.equal(second.flags.leaseId, DAEMON_LEASE_ID);
95+
assert.equal(
96+
readRemoteConnectionState({ stateDir, session: 'adc-proxy' })?.leaseId,
97+
DAEMON_LEASE_ID,
98+
);
99+
fs.rmSync(tempRoot, { recursive: true, force: true });
100+
});
101+
102+
test('a non-proxy remote connection leaves the lease id to the daemon', async () => {
103+
const { tempRoot, stateDir, remoteConfigPath } = connectionWorkspace(
104+
'agent-device-cloud-lease-id-',
105+
);
106+
fs.writeFileSync(remoteConfigPath, JSON.stringify({}));
107+
seedConnectionState({
108+
stateDir,
109+
state: {
110+
session: 'adc-cloud',
111+
remoteConfigPath,
112+
tenant: 'acme',
113+
runId: 'run-123',
114+
leaseProvider: 'aws-device-farm',
115+
platform: 'ios',
116+
},
117+
});
118+
const requests: LeaseAllocateRequest[] = [];
119+
120+
const materialized = await materializeRemoteConnectionForCommand({
121+
command: 'screenshot',
122+
flags: {
123+
json: true,
124+
help: false,
125+
version: false,
126+
stateDir,
127+
remoteConfig: remoteConfigPath,
128+
tenant: 'acme',
129+
runId: 'run-123',
130+
session: 'adc-cloud',
131+
platform: 'ios',
132+
},
133+
client: createTestClient({ allocate: recordingAllocate(requests) }),
134+
});
135+
136+
assert.equal(requests.length, 1);
137+
assert.equal(requests[0]?.leaseId, undefined);
138+
assert.equal(materialized.flags.leaseId, DAEMON_LEASE_ID);
139+
fs.rmSync(tempRoot, { recursive: true, force: true });
140+
});

‎src/cli/commands/connection-runtime.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import crypto from 'node:crypto';
12
import type { MetroBridgeScope } from '@agent-device/contracts/remote';
23
import { resolveDaemonPaths } from '@agent-device/daemon-contracts/daemon-resolution';
34
import { stopReactDevtoolsCompanion } from '../../client/client-react-devtools-companion.ts';
@@ -424,6 +425,8 @@ function buildMaterializedLeaseState(
424425
type ConnectionLeasePolicy = {
425426
shouldAllocate(command: string): boolean;
426427
ttlMs(state: RemoteConnectionState): number | undefined;
428+
/** The id this client names a new lease with; without it the daemon mints one. */
429+
newLeaseId?(): string;
427430
resolveLeaseState(options: {
428431
command: string;
429432
client: AgentDeviceClient;
@@ -464,6 +467,8 @@ const DEFAULT_CONNECTION_LEASE_POLICY: ConnectionLeasePolicy = {
464467
const PROXY_CONNECTION_LEASE_POLICY: ConnectionLeasePolicy = {
465468
shouldAllocate: (command) => command !== 'devices' && !leaseDeferredCommands.has(command),
466469
ttlMs: () => PROXY_REMOTE_LEASE_TTL_MS,
470+
// Known before the daemon allocates, so a host in front of it can name the lease with the same id.
471+
newLeaseId: () => crypto.randomBytes(16).toString('hex'),
467472
resolveLeaseState: resolveProxyLeaseState,
468473
};
469474

@@ -904,6 +909,7 @@ async function allocateOrReuseLease(
904909
const lease = await client.leases.allocate({
905910
tenant: state.tenant,
906911
runId: state.runId,
912+
leaseId: policy.newLeaseId?.(),
907913
leaseBackend,
908914
...connection,
909915
ttlMs: policy.ttlMs(state),

‎src/client/lease-client.test.ts‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,3 +97,37 @@ test('allocate reports whether the daemon kept retainOnClose on the lease', asyn
9797
assert.equal((await allocate(false)).retainOnClose, undefined);
9898
assert.equal(requests[0]?.meta?.leaseRetainOnClose, true);
9999
});
100+
101+
test('allocate sends the lease id named on the call, never the configured one', async () => {
102+
const requests: Array<Omit<DaemonRequest, 'token'>> = [];
103+
const client = createAgentDeviceClient(
104+
{ leaseId: 'feedfacefeedface' },
105+
{
106+
transport: async (request) => {
107+
requests.push(request);
108+
return {
109+
ok: true,
110+
data: {
111+
lease: {
112+
leaseId: 'abcdef0123456789',
113+
tenantId: 't',
114+
runId: 'r',
115+
backend: 'ios-instance',
116+
},
117+
},
118+
};
119+
},
120+
},
121+
);
122+
123+
await client.leases.allocate({ tenant: 't', runId: 'r' });
124+
const lease = await client.leases.allocate({
125+
tenant: 't',
126+
runId: 'r',
127+
leaseId: '0123456789abcdef',
128+
});
129+
130+
assert.equal(requests[0]?.meta?.leaseId, undefined);
131+
assert.equal(requests[1]?.meta?.leaseId, '0123456789abcdef');
132+
assert.equal(lease.leaseId, 'abcdef0123456789');
133+
});

‎src/client/lease-client.ts‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,10 @@ export function createLeaseClient(
1818
return {
1919
allocate: async (options) =>
2020
normalizeLease(
21-
await execute(INTERNAL_COMMANDS.leaseAllocate, [], { ...options, leaseId: undefined }),
21+
await execute(INTERNAL_COMMANDS.leaseAllocate, [], {
22+
...options,
23+
leaseId: options.leaseId,
24+
}),
2225
),
2326
heartbeat: async (options) =>
2427
normalizeLease(await execute(INTERNAL_COMMANDS.leaseHeartbeat, [], options)),

‎src/daemon-client/daemon-client-rpc.test.ts‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,31 @@ test('lease allocation transports an optional initial provider app', () => {
3333
});
3434
});
3535

36+
test('lease allocation transports a caller-chosen lease id', () => {
37+
const payload = buildHttpRpcPayload(
38+
{
39+
token: 'daemon-token',
40+
session: 'default',
41+
command: 'lease_allocate',
42+
positionals: [],
43+
meta: {
44+
requestId: 'lease-req',
45+
tenantId: 'proxy',
46+
runId: 'proxy-client-1',
47+
leaseId: 'abcdef0123456789abcdef0123456789',
48+
},
49+
},
50+
{ includeTokenParam: false },
51+
);
52+
53+
assert.deepEqual(payload.params, {
54+
session: 'default',
55+
tenantId: 'proxy',
56+
runId: 'proxy-client-1',
57+
leaseId: 'abcdef0123456789abcdef0123456789',
58+
});
59+
});
60+
3661
test('lease allocation transports the provider configuration the session needs (#2494)', () => {
3762
const payload = buildHttpRpcPayload(
3863
{

‎src/daemon/__tests__/lease-registry.test.ts‎

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import { test } from 'vitest';
22
import assert from 'node:assert/strict';
3-
import { createRequestCanceledError } from '@agent-device/kernel/errors';
3+
import { AppError, createRequestCanceledError } from '@agent-device/kernel/errors';
44
import { LeaseRegistry } from '../lease-registry.ts';
55
import {
66
HUMAN_CONTROL_LEASE_REQUEST,
@@ -30,6 +30,34 @@ test('allocateLease creates lease and enforces tenant/run validation', () => {
3030
);
3131
});
3232

33+
test('allocateLease keeps a caller-chosen lease id and refuses one an active lease holds', () => {
34+
const registry = new LeaseRegistry();
35+
const chosen = 'ABCDEF0123456789abcdef0123456789';
36+
const lease = registry.allocateLease({ tenantId: 'tenant-a', runId: 'run-1', leaseId: chosen });
37+
assert.equal(lease.leaseId, chosen.toLowerCase());
38+
assert.equal(registry.heartbeatLease({ leaseId: lease.leaseId }).leaseId, lease.leaseId);
39+
40+
assert.throws(
41+
() => registry.allocateLease({ tenantId: 'tenant-b', runId: 'run-2', leaseId: chosen }),
42+
(error: unknown) =>
43+
error instanceof AppError &&
44+
error.code === 'INVALID_ARGS' &&
45+
error.details?.reason === 'LEASE_ID_IN_USE',
46+
);
47+
assert.throws(
48+
() => registry.allocateLease({ tenantId: 'tenant-b', runId: 'run-2', leaseId: 'not-hex' }),
49+
(error: unknown) => error instanceof AppError && error.code === 'INVALID_ARGS',
50+
);
51+
52+
assert.equal(registry.releaseLease({ leaseId: lease.leaseId }).released, true);
53+
const reused = registry.allocateLease({ tenantId: 'tenant-b', runId: 'run-2', leaseId: chosen });
54+
assert.equal(reused.leaseId, lease.leaseId);
55+
assert.notEqual(
56+
registry.allocateLease({ tenantId: 'tenant-c', runId: 'run-3' }).leaseId,
57+
lease.leaseId,
58+
);
59+
});
60+
3361
test('allocateLease is idempotent per tenant/run/backend and refreshes expiry', () => {
3462
let now = 1_000;
3563
const registry = new LeaseRegistry({

‎src/daemon/handlers/__tests__/lease.test.ts‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,26 @@ for (const operation of ['allocate', 'release'] as const) {
6161
});
6262
}
6363

64+
test('lease_allocate answers with the lease id the request named', async () => {
65+
const registry = new LeaseRegistry();
66+
const leaseId = 'abcdef0123456789abcdef0123456789';
67+
const response = await handleLeaseCommands({
68+
req: {
69+
token: 'test-token',
70+
session: 'default',
71+
command: 'lease_allocate',
72+
positionals: [],
73+
meta: { tenantId: 'tenant-a', runId: 'run-1', leaseId },
74+
},
75+
sessionName: 'default',
76+
sessionStore: makeSessionStore('agent-device-chosen-lease-id-'),
77+
leaseRegistry: registry,
78+
});
79+
assert.ok(response?.ok);
80+
assert.equal((response.data?.lease as DeviceLease | undefined)?.leaseId, leaseId);
81+
assert.equal(registry.listActiveLeases()[0]?.leaseId, leaseId);
82+
});
83+
6484
test('lease_release still releases a retainOnClose lease through the provider', async () => {
6585
const registry = new LeaseRegistry();
6686
const lease = registry.allocateLease({

0 commit comments

Comments
 (0)