diff --git a/backend/internal/adapters/systemexec/remote_script_test.go b/backend/internal/adapters/systemexec/remote_script_test.go index 4dba4927e6b..0603309aebf 100644 --- a/backend/internal/adapters/systemexec/remote_script_test.go +++ b/backend/internal/adapters/systemexec/remote_script_test.go @@ -78,9 +78,15 @@ func TestRunInstallScriptCancellationCleansUp(t *testing.T) { if err == nil { t.Fatal("expected cancellation") } + // Cancellation can land while the download is still in flight, before the + // scratch directory is ever created, so an absent directory is also a clean + // outcome: nothing leaked. Only surviving entries indicate a leftover script. entries, readErr := os.ReadDir(filepath.Join(dataDir, "installers", "tmp")) - if readErr != nil || len(entries) != 0 { - t.Fatalf("temporary scripts remain after cancellation: %v, %v", entries, readErr) + if readErr != nil && !os.IsNotExist(readErr) { + t.Fatalf("read installer scratch dir: %v", readErr) + } + if len(entries) != 0 { + t.Fatalf("temporary scripts remain after cancellation: %v", entries) } } diff --git a/cloud/internal/postgres/worker_transport_store.go b/cloud/internal/postgres/worker_transport_store.go index 4cb24984321..191fe071e9a 100644 --- a/cloud/internal/postgres/worker_transport_store.go +++ b/cloud/internal/postgres/worker_transport_store.go @@ -355,12 +355,27 @@ func (s *Store) IssueTerminalTicket( var exited bool err := s.withSessionAccess(ctx, principal, orgID, sessionID, func(tx pgx.Tx, _ sessionAccess) error { return tx.QueryRow(ctx, - `SELECT session.is_terminated OR session.activity_state = 'exited' OR EXISTS ( + `SELECT session.is_terminated OR session.activity_state = 'exited' OR ( + EXISTS ( SELECT 1 FROM ao_terminal_sessions terminal WHERE terminal.org_id = session.org_id AND terminal.session_id = session.id AND terminal.kind = 'agent' AND terminal.state IN ('closed', 'failed') + ) + AND NOT EXISTS ( + SELECT 1 FROM ao_terminal_sessions terminal + JOIN ao_worker_connections worker + ON worker.org_id = terminal.org_id + AND worker.session_id = terminal.session_id + AND worker.epoch = terminal.worker_epoch + AND worker.disconnected_at IS NULL + WHERE terminal.org_id = session.org_id + AND terminal.session_id = session.id + AND terminal.kind = 'agent' + AND terminal.state IN ('opening', 'open') + AND terminal.expires_at > now() + ) ) FROM ao_sessions session WHERE session.org_id = $1 AND session.id = $2`, @@ -472,12 +487,27 @@ func (s *Store) IssueTerminalTicket( if kind == "agent" { var exited bool lookupErr := tx.QueryRow(ctx, - `SELECT session.is_terminated OR session.activity_state = 'exited' OR EXISTS ( + `SELECT session.is_terminated OR session.activity_state = 'exited' OR ( + EXISTS ( SELECT 1 FROM ao_terminal_sessions terminal WHERE terminal.org_id = session.org_id AND terminal.session_id = session.id AND terminal.kind = 'agent' AND terminal.state IN ('closed', 'failed') + ) + AND NOT EXISTS ( + SELECT 1 FROM ao_terminal_sessions terminal + JOIN ao_worker_connections worker + ON worker.org_id = terminal.org_id + AND worker.session_id = terminal.session_id + AND worker.epoch = terminal.worker_epoch + AND worker.disconnected_at IS NULL + WHERE terminal.org_id = session.org_id + AND terminal.session_id = session.id + AND terminal.kind = 'agent' + AND terminal.state IN ('opening', 'open') + AND terminal.expires_at > now() + ) ) FROM ao_sessions session WHERE session.org_id = $1 AND session.id = $2`,