fix(pty): reset persistent shell on cancellation

This commit is contained in:
Tianyi Cui
2026-07-29 21:33:39 +08:00
parent fbad903dd0
commit 205702adaf
2 changed files with 31 additions and 22 deletions

View File

@@ -309,6 +309,10 @@ async function executeCommand(
SHELL_RESET_MESSAGE, SHELL_RESET_MESSAGE,
].join('\n') ].join('\n')
} }
if (commandDeadline.signal.aborted) {
await shells.reset(owner, 'persistent bash command aborted')
commandDeadline.signal.throwIfAborted()
}
if (latest.text.includes(marker.end)) { if (latest.text.includes(marker.end)) {
const complete = commandOutput(retainedScrollback(ctx, owner, id, latest), marker) const complete = commandOutput(retainedScrollback(ctx, owner, id, latest), marker)
return renderCaptured(complete, config.maxOutputChars) return renderCaptured(complete, config.maxOutputChars)
@@ -325,10 +329,6 @@ async function executeCommand(
SHELL_RESET_MESSAGE, SHELL_RESET_MESSAGE,
].filter(part => part.length > 0).join('\n') ].filter(part => part.length > 0).join('\n')
} }
if (commandDeadline.signal.aborted) {
await shells.reset(owner, 'persistent bash command aborted')
commandDeadline.signal.throwIfAborted()
}
if (promptCompleted(result)) { if (promptCompleted(result)) {
const snapshot = retainedScrollback(ctx, owner, id, latest) const snapshot = retainedScrollback(ctx, owner, id, latest)
return renderCaptured( return renderCaptured(

View File

@@ -80,6 +80,7 @@ type StubMode =
| 'exit' | 'exit'
| 'signal-exit' | 'signal-exit'
| 'wait-for-abort' | 'wait-for-abort'
| 'end-on-abort'
| 'idle-then-normal' | 'idle-then-normal'
| 'large' | 'large'
| 'nonzero' | 'nonzero'
@@ -119,11 +120,16 @@ class StubPtySession implements PtyBackendSession {
return this.operation(Promise.resolve(this.result(this.motd, 'stdin_read'))) return this.operation(Promise.resolve(this.result(this.motd, 'stdin_read')))
} }
if (this.mode === 'send-error') throw new Error('stub send failed') if (this.mode === 'send-error') throw new Error('stub send failed')
if (this.mode === 'wait-for-abort') { if (this.mode === 'wait-for-abort' || this.mode === 'end-on-abort') {
const done = new Promise<ReturnType<StubPtySession['result']>>((resolve) => { const done = new Promise<ReturnType<StubPtySession['result']>>((resolve) => {
request.signal?.addEventListener('abort', () => { request.signal?.addEventListener('abort', () => {
this.scrollback += 'partial output' const start = /__DSH_PERSISTENT_BASH_START_[^_]+(?:-[^_]+)*__/.exec(request.text)?.[0]
resolve(this.result('partial output', 'stdin_read')) const end = /__DSH_PERSISTENT_BASH_END_[^:]+:/.exec(request.text)?.[0]
const output = this.mode === 'end-on-abort'
? `${start ?? ''}\ninterrupted\n${end ?? ''}130\n${this.motd}`
: 'partial output'
this.scrollback += output
resolve(this.result(output, 'stdin_read'))
}, { once: true }) }, { once: true })
}) })
return this.operation(done) return this.operation(done)
@@ -389,22 +395,25 @@ describe('tool-bash-persistent', () => {
expect(stub.sessions[0]?.closed).toContain('persistent bash command timed out') expect(stub.sessions[0]?.closed).toContain('persistent bash command timed out')
}) })
it('cancels in-flight work, resets the shell, and releases a queued call', async () => { it.each(['wait-for-abort', 'end-on-abort'] as const)(
const { ctx, owner, stub } = await setup({ backendType: 'stub', timeoutMs: 5_000 }) 'cancels %s work, resets the shell, and releases a queued call',
await call(ctx, owner, 'warm up') async (mode) => {
stub.sessions[0]!.mode = 'wait-for-abort' const { ctx, owner, stub } = await setup({ backendType: 'stub', timeoutMs: 5_000 })
const controller = new AbortController() await call(ctx, owner, 'warm up')
const cancelled = call(ctx, owner, 'hang', controller.signal) stub.sessions[0]!.mode = mode
const queued = call(ctx, owner, 'after cancellation') const controller = new AbortController()
setTimeout(() => { const cancelled = call(ctx, owner, 'hang', controller.signal)
controller.abort(new Error('caller stopped')) const queued = call(ctx, owner, 'after cancellation')
}, 5) setTimeout(() => {
controller.abort(new Error('caller stopped'))
}, 5)
expect((await cancelled).isError).toBe(true) expect((await cancelled).isError).toBe(true)
expect(text(await queued)).toBe('hello from stub') expect(text(await queued)).toBe('hello from stub')
expect(stub.sessions[0]?.closed).toContain('persistent bash command aborted') expect(stub.sessions[0]?.closed).toContain('persistent bash command aborted')
expect(stub.sessions).toHaveLength(2) expect(stub.sessions).toHaveLength(2)
}) },
)
it.each(['init-exit', 'init-timeout'] as const)( it.each(['init-exit', 'init-timeout'] as const)(
'fails initialization and closes the unusable shell for %s', 'fails initialization and closes the unusable shell for %s',