Skip to content

Commit 9281bc9

Browse files
committed
fix(eval): require a server admission token before preserve-detach
Transport in-flight is not Host admission. preserve_environment now calls hosted.execution.admit and detaches only after that token returns. A frame-written interrupt without the token settles the Host instead of claiming execution continues. Cancelled relay cleanup funnels every finalize exception through settle-or-destroy before rethrowing, including execution and persist failures during host abort. Fixes #3150
1 parent 2a907a5 commit 9281bc9

10 files changed

Lines changed: 320 additions & 53 deletions

File tree

packages/cli/src/__tests__/runtime-host-operator-command.test.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,7 @@ describe('Runtime Host operator commands', () => {
138138
'access.credential.issue',
139139
'access.credential.revoke',
140140
'host.upgrade.prepare',
141+
'hosted.execution.admit',
141142
'hosted.execution.cancel',
142143
'hosted.execution.start',
143144
],

packages/eval/harbor/relay_agent.py

Lines changed: 40 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -172,24 +172,16 @@ async def run(self, instruction: str, environment: Any, context: Any) -> None:
172172
current = asyncio.current_task()
173173
if current is not None and hasattr(current, "uncancel"):
174174
current.uncancel()
175-
try:
176-
if request is not None and execution is not None:
177-
await self._finalize_cancelled_execution(
178-
environment,
179-
cwd,
180-
scope_path,
181-
execution,
182-
request,
183-
writer,
184-
execution_reported,
185-
)
186-
except asyncio.CancelledError:
187-
if request is not None and execution is not None:
188-
with contextlib.suppress(Exception):
189-
await _settle_or_destroy(
190-
environment, cwd, scope_path, execution, self._teardown_timeout
191-
)
192-
raise
175+
if request is not None and execution is not None:
176+
await self._cleanup_cancelled_execution(
177+
environment,
178+
cwd,
179+
scope_path,
180+
execution,
181+
request,
182+
writer,
183+
execution_reported,
184+
)
193185
raise
194186
except RelayTransportClosed:
195187
if request is not None and execution is not None:
@@ -227,6 +219,36 @@ async def run(self, instruction: str, environment: Any, context: Any) -> None:
227219
writer.close()
228220
await asyncio.wait_for(writer.wait_closed(), timeout=1)
229221

222+
async def _cleanup_cancelled_execution(
223+
self,
224+
environment: Any,
225+
cwd: str,
226+
scope_path: str,
227+
execution: asyncio.Task[Any],
228+
request: dict[str, Any],
229+
writer: Any,
230+
execution_reported: bool,
231+
) -> None:
232+
try:
233+
await self._finalize_cancelled_execution(
234+
environment,
235+
cwd,
236+
scope_path,
237+
execution,
238+
request,
239+
writer,
240+
execution_reported,
241+
)
242+
except BaseException:
243+
with contextlib.suppress(Exception):
244+
await asyncio.wait_for(
245+
_settle_or_destroy(
246+
environment, cwd, scope_path, execution, self._teardown_timeout
247+
),
248+
timeout=self._teardown_timeout,
249+
)
250+
raise
251+
230252
async def _finalize_cancelled_execution(
231253
self,
232254
environment: Any,

packages/eval/harbor/test_relay_lifecycle.py

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,17 @@ async def exec(self, command, cwd=None, timeout_sec=None):
8383
return SimpleNamespace(return_code=0, stdout="", stderr="")
8484

8585

86+
class DummyWriter:
87+
def is_closing(self):
88+
return False
89+
90+
def write(self, _value):
91+
return None
92+
93+
async def drain(self):
94+
return None
95+
96+
8697
class ClosedWriter:
8798
def is_closing(self):
8899
return False
@@ -188,6 +199,20 @@ async def exec(self, command, cwd=None, timeout_sec=None):
188199
return await super().exec(command, cwd=cwd, timeout_sec=timeout_sec)
189200

190201

202+
class ExplodingSubjectEnvironment:
203+
def __init__(self):
204+
self.stopped = False
205+
206+
async def exec(self, command, cwd=None, timeout_sec=None):
207+
return SimpleNamespace(return_code=0, stdout="", stderr="")
208+
209+
async def stop(self, delete=False):
210+
self.stopped = delete
211+
212+
async def upload_file(self, source, target):
213+
return None
214+
215+
191216
class SlowLeaderStopEnvironment(IgnoringLeaderEnvironment):
192217
def __init__(self):
193218
super().__init__()
@@ -218,6 +243,19 @@ async def exec(self, command, cwd=None, timeout_sec=None):
218243
return await super().exec(command, cwd=cwd, timeout_sec=timeout_sec)
219244

220245

246+
class PersistFailureEnvironment(LiveScopeEnvironment):
247+
def __init__(self):
248+
super().__init__()
249+
self.stopped = False
250+
251+
async def stop(self, delete=False):
252+
self.stopped = delete
253+
254+
async def upload_file(self, source, target):
255+
if str(target).endswith("maka-subject.stdout.txt"):
256+
raise RuntimeError("persist failed")
257+
258+
221259
class TransportLossEnvironment(SimultaneousEnvironment):
222260
async def exec(self, command, cwd=None, timeout_sec=None):
223261
if "kill -TERM" in command:
@@ -734,6 +772,87 @@ async def accept(reader, writer):
734772
server.close()
735773
await server.wait_closed()
736774

775+
async def test_execution_exception_during_host_abort_still_destroys(self):
776+
relay = load_relay()
777+
environment = ExplodingSubjectEnvironment()
778+
779+
async def boom():
780+
raise RuntimeError("subject execution exploded")
781+
782+
execution = asyncio.create_task(boom())
783+
await asyncio.sleep(0)
784+
relay.request_host_teardown()
785+
agent = relay.RelayAgent(
786+
logs_dir=Path(tempfile.gettempdir()),
787+
relay_host="127.0.0.1",
788+
relay_port=1,
789+
relay_token="token",
790+
teardown_timeout_ms=1_000,
791+
)
792+
with self.assertRaisesRegex(RuntimeError, "subject execution exploded"):
793+
await agent._cleanup_cancelled_execution(
794+
environment,
795+
"/",
796+
"/tmp/missing",
797+
execution,
798+
{"resultToken": "0" * 32, "captureStdout": True},
799+
DummyWriter(),
800+
False,
801+
)
802+
self.assertTrue(environment.stopped)
803+
804+
async def test_persistence_failure_during_host_abort_still_destroys(self):
805+
relay = load_relay()
806+
environment = PersistFailureEnvironment()
807+
token = f"host-persist-{os.getpid()}"
808+
connected = asyncio.get_running_loop().create_future()
809+
810+
async def accept(reader, writer):
811+
connected.set_result((reader, writer))
812+
813+
server = await asyncio.start_server(accept, "127.0.0.1", 0)
814+
port = server.sockets[0].getsockname()[1]
815+
agent = relay.RelayAgent(
816+
logs_dir=Path(tempfile.gettempdir()),
817+
relay_host="127.0.0.1",
818+
relay_port=port,
819+
relay_token=token,
820+
teardown_timeout_ms=1_000,
821+
)
822+
running = asyncio.create_task(agent.run("solve", environment, None))
823+
reader, writer = await connected
824+
try:
825+
await reader.readline()
826+
writer.write(
827+
(
828+
__import__("json").dumps(
829+
{
830+
"token": token,
831+
"kind": "execute",
832+
"command": "/bin/true",
833+
"args": [],
834+
"credentials": {},
835+
"resultToken": "0" * 32,
836+
}
837+
)
838+
+ "\n"
839+
).encode()
840+
)
841+
await writer.drain()
842+
environment.release.set()
843+
await environment.finished.wait()
844+
relay.request_host_teardown()
845+
running.cancel()
846+
with self.assertRaisesRegex(RuntimeError, "persist failed"):
847+
await asyncio.wait_for(running, timeout=1)
848+
self.assertTrue(
849+
any(_is_teardown(command) for command in environment.commands),
850+
)
851+
finally:
852+
writer.close()
853+
server.close()
854+
await server.wait_closed()
855+
737856
async def test_framework_timeout_stops_the_subject_without_the_process_group(self):
738857
relay = load_relay()
739858
environment = TimeoutScopeEnvironment()

packages/runtime-host/src/__tests__/hosted-execution-client.test.ts

Lines changed: 20 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -94,18 +94,21 @@ test('post-connect cancellation reports the owned Host settlement outcome', asyn
9494

9595
test('environment-preserving abort detaches without cancelling or settling the Host', async () => {
9696
const abort = new AbortController();
97-
const started = deferred();
97+
const admitted = deferred();
9898
const closed = deferred();
9999
const events: string[] = [];
100100
const connected = ownedHost({
101101
request: async (operation: string) => {
102102
events.push(operation);
103+
if (operation === 'hosted.execution.admit') {
104+
admitted.resolve();
105+
return { executionId: ID, admissionToken: ADMISSION_TOKEN };
106+
}
103107
if (operation !== 'hosted.execution.start') {
104108
throw new Error(`Unexpected operation ${operation}`);
105109
}
106-
started.resolve();
107110
await closed.promise;
108-
throw admittedInterrupt();
111+
throw dispatchedInterrupt();
109112
},
110113
});
111114
connected.connection.close = async () => {
@@ -124,29 +127,34 @@ test('environment-preserving abort detaches without cancelling or settling the H
124127
{ ...input(abort.signal), abortPolicy: 'preserve_environment' },
125128
{ connectOwnedRuntimeHost: async () => connected as never },
126129
);
127-
await started.promise;
130+
await admitted.promise;
128131
abort.abort();
129132
const result = await execution;
130133

131134
assert.equal(result.kind, 'indeterminate');
132135
assert.equal(result.failureReason, 'Hosted execution continues for environment verification');
133-
assert.deepEqual(events, ['hosted.execution.start', 'close', 'release']);
136+
assert.deepEqual(events, [
137+
'hosted.execution.admit',
138+
'close',
139+
'hosted.execution.start',
140+
'release',
141+
]);
134142
});
135143

136-
test('pre-admission preserve abort does not claim execution continues', async () => {
144+
test('frame-written admit interruption is not a server admission', async () => {
137145
const abort = new AbortController();
138146
const started = deferred();
139147
const closed = deferred();
140148
const events: string[] = [];
141149
const connected = ownedHost({
142150
request: async (operation: string) => {
143151
events.push(operation);
144-
if (operation !== 'hosted.execution.start') {
152+
if (operation !== 'hosted.execution.admit') {
145153
throw new Error(`Unexpected operation ${operation}`);
146154
}
147155
started.resolve();
148156
await closed.promise;
149-
throw queuedInterrupt();
157+
throw dispatchedInterrupt();
150158
},
151159
});
152160
connected.connection.close = async () => {
@@ -172,7 +180,7 @@ test('pre-admission preserve abort does not claim execution continues', async ()
172180

173181
assert.equal(result.kind, 'indeterminate');
174182
assert.equal(result.failureReason, 'Hosted execution was cancelled');
175-
assert.deepEqual(events, ['hosted.execution.start', 'close', 'settle']);
183+
assert.deepEqual(events, ['hosted.execution.admit', 'close', 'settle']);
176184
});
177185

178186
test('environment-preserving abort before start does not claim execution continues', async () => {
@@ -322,6 +330,7 @@ test('explicit target reconnect cancellation preserves the cancelled outcome', a
322330

323331
const ID = '00000000-0000-4000-8000-000000000001';
324332
const CONNECTION_ID = '00000000-0000-4000-8000-000000000002';
333+
const ADMISSION_TOKEN = '00000000-0000-4000-8000-0000000000ad';
325334

326335
function input(signal?: AbortSignal) {
327336
return {
@@ -458,24 +467,15 @@ function ownedHost(connection: Record<string, unknown>, clean = false) {
458467
};
459468
}
460469

461-
function admittedInterrupt() {
470+
function dispatchedInterrupt() {
462471
return new RuntimeHostRequestInterruptedError(
463-
'hosted.execution.start',
472+
'hosted.execution.admit',
464473
'command',
465474
'dispatched',
466475
'connection_lost',
467476
);
468477
}
469478

470-
function queuedInterrupt() {
471-
return new RuntimeHostRequestInterruptedError(
472-
'hosted.execution.start',
473-
'command',
474-
'not_dispatched',
475-
'connection_lost',
476-
);
477-
}
478-
479479
function deferred() {
480480
let resolve!: () => void;
481481
const promise = new Promise<void>((accept) => {

packages/runtime-host/src/__tests__/hosted-execution-coordinator.test.ts

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,32 @@ test('cancel before start prevents the hosted execution from running', async ()
2424
assert.equal(runs, 0);
2525
});
2626

27+
test('admit returns a server-owned token before start waits for settlement', async () => {
28+
const release = deferred();
29+
const coordinator = new HostHostedExecutionCoordinator(
30+
async (input) => {
31+
await release.promise;
32+
return settled(input.executionId, 'completed');
33+
},
34+
() => {},
35+
);
36+
37+
const admitted = await coordinator.handlers['hosted.execution.admit'](input(), context());
38+
assert.equal(admitted.ok, true);
39+
if (!admitted.ok) return;
40+
assert.equal(admitted.result.executionId, ID);
41+
assert.match(admitted.result.admissionToken, /^[0-9a-f-]{36}$/i);
42+
43+
const waiting = coordinator.handlers['hosted.execution.start'](input(), context());
44+
const again = await coordinator.handlers['hosted.execution.admit'](input(), context());
45+
assert.equal(again.ok, true);
46+
if (again.ok) assert.equal(again.result.admissionToken, admitted.result.admissionToken);
47+
release.resolve();
48+
const started = await waiting;
49+
assert.equal(started.ok, true);
50+
if (started.ok) assert.equal(started.result.kind, 'settled');
51+
});
52+
2753
test('cancelling a settled subject reclaims its verification environment', async () => {
2854
let drains = 0;
2955
const coordinator = new HostHostedExecutionCoordinator(
@@ -79,3 +105,11 @@ function context() {
79105
acquireResidency: () => ({ release() {} }),
80106
};
81107
}
108+
109+
function deferred() {
110+
let resolve!: () => void;
111+
const promise = new Promise<void>((accept) => {
112+
resolve = accept;
113+
});
114+
return { promise, resolve };
115+
}

packages/runtime-host/src/__tests__/protocol.test.ts

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,12 @@ describe('Runtime Host bootstrap protocol', () => {
4646
assert.ok(RUNTIME_HOST_COMPATIBILITY_EPOCH > 22);
4747
});
4848

49+
test('publishes a new compatibility epoch for hosted.execution.admit', () => {
50+
// Epoch 25 has start and cancel only. Mixed-version peers must fail
51+
// during handshake instead of sending an unknown admit command.
52+
assert.ok(RUNTIME_HOST_COMPATIBILITY_EPOCH > 25);
53+
});
54+
4955
test('rejects the legacy connection update result in the current compatibility epoch', () => {
5056
assert.throws(
5157
() =>

0 commit comments

Comments
 (0)