@zgeoff/atc 2.15.0 → 2.15.2

package/src/daemon/sessions.ts

@@ -1,21 +1,25 @@
 import { writeFileSync } from 'node:fs';
 import type {
-  AdapterEvent,
   AgentAdapter,
-  AgentID,
   SpawnOptions,
   SpawnOverrides,
   SpawnPlan,
 } from '../agents/agent-adapter';
-import { truncateDetail } from '../agents/truncate-detail';
+import type { AdapterEvent } from '../protocol/adapter-event';
+import { countSessionStates } from '../protocol/count-session-states';
 import { DaemonError } from '../protocol/daemon-error';
+import type { HookEvent } from '../protocol/hook-event';
 import type { ErrorCode } from '../protocol/protocol';
+import type { SessionState } from '../protocol/session-state';
+import { sortSessionViews } from '../protocol/sort-session-views';
+import type { AgentID } from '../shared/agent-id';
 import type { AgentSessionID } from '../shared/agent-session-id';
 import type { TargetConfigError } from '../shared/collect-targets';
 import { socketPath, statusFile } from '../shared/config';
 import type { DaemonID } from '../shared/daemon-id';
 import { resolveRepoRoot } from '../shared/resolve-repo-root';
 import type { SessionID } from '../shared/session-id';
+import { truncateDetail } from '../shared/truncate-detail';
 import { truncateToBytes } from '../shared/truncate-to-bytes';
 import type { FleetEntry, FleetEntryUpdate, FleetStore } from '../store/fleet-entry';
 import type { SessionWorkspace } from '../store/workspace-materialization';
@@ -31,14 +35,11 @@ import type {
   HarnessRelay,
 } from './execution-provider';
 import { findExecutionRefusal } from './find-execution-refusal';
-import type { HookEvent } from './hooks';
 import type { BridgeBinding } from './is-binding-current';
 import { LocalPTYProvider } from './local-pty-provider';
 import { mintSessionID } from './mint-session-id';
 import { pickSessionState } from './pick-session-state';
 
-export type SessionState = 'running' | 'needs_you' | 'done' | 'exited';
-
 export type SessionEventKind = 'added' | 'state' | 'renamed' | 'removed';
 
 // The over-the-wire view of a session: everything but the harness handle,
@@ -181,6 +182,10 @@ interface MaterializedSpawn {
 // stored identity ran on.
 const LOCAL_TARGET_IDENTITY = buildTargetIdentity('local-pty', {});
 
+// How long a failed spawn's rollback waits for the killed process to exit,
+// once after its kill and once more after a forced kill.
+const FAILED_SPAWN_EXIT_WAIT_MS = 2000;
+
 // The target a session runs on and the identity it is bound to there; null
 // binds a new session to the target as it stands now.
 interface TargetBinding {
@@ -236,6 +241,13 @@ export class SessionManager {
   // write deletes; each stays here until a write carrying it lands.
   private readonly removedIDs = new Set<SessionID>();
 
+  // Failed spawns whose rollback is running, and those whose killed process
+  // it could not confirm gone, each with the harness it ran in, so no revive
+  // starts a second harness while the first may still run.
+  private readonly rollingBack = new Set<SessionID>();
+
+  private readonly unconfirmedKills = new Map<SessionID, HarnessHandle>();
+
   // The epoch the next harness start or attach takes, unique across every
   // session this manager holds.
   private nextBridgeEpoch = 1;
@@ -467,14 +479,27 @@ export class SessionManager {
   // Adopts a headless session back into a terminal: a fresh PTY resumes the
   // same agent session id. On a remote host the host wakes first, and a
   // harness still running inside it is attached rather than started again;
-  // every other session left asleep on that host comes back with it.
-  async adoptTerminal(id: SessionID, cols: number, rows: number): Promise<Session | null> {
+  // every other session left asleep on that host comes back with it. The
+  // caller's check runs again after each await, and a failed check leaves
+  // the session as it was.
+  async adoptTerminal(
+    id: SessionID,
+    cols: number,
+    rows: number,
+    canProceed: () => boolean = () => true,
+  ): Promise<Session | null> {
     const s = this.sessions.find((x) => x.id === id);
 
     if (!s || s.pty !== null || s.agentSessionID === undefined || this.adopting.has(id)) {
       return null;
     }
 
+    const settled = await this.isRollbackSettled(id);
+
+    if (!settled || !canProceed()) {
+      return null;
+    }
+
     const adapter = this.findAdapter(s.agent);
 
     if (adapter === null) {
@@ -498,8 +523,16 @@ export class SessionManager {
       this.adopting.delete(id);
     }
 
-    // A kill or a second adopt can land while the host wakes.
-    if (s.pty !== null || !this.sessions.includes(s)) {
+    // A kill, a second adopt, a failed spawn's rollback, or a change that
+    // takes the session out of the caller's reach can land while the host
+    // wakes.
+    if (
+      !canProceed() ||
+      s.pty !== null ||
+      !this.sessions.includes(s) ||
+      this.rollingBack.has(id) ||
+      this.unconfirmedKills.has(id)
+    ) {
       return null;
     }
 
@@ -705,6 +738,7 @@ export class SessionManager {
     id: SessionID = mintSessionID(),
     target = 'local',
     materialized: MaterializedSpawn | null = null,
+    requireInReach: () => void = () => {},
   ): Promise<Session> {
     const adapter = this.findAdapter(agent);
 
@@ -727,6 +761,10 @@ export class SessionManager {
       ...overrides,
     });
 
+    // The caller's check runs again once the host is ready, before the
+    // harness starts.
+    requireInReach();
+
     const binding = this.mintBridgeBinding(id, target, execution.identity, hostKey);
 
     const pty = provider.spawnHarness({
@@ -893,8 +931,11 @@ export class SessionManager {
 
   // Takes back a spawn that failed after its process started: the process
   // dies and the session leaves the list and the fleet, so the spawn leaves
-  // nothing behind. It resolves only once the fleet without the session is
-  // durable, and throws when the kill or that write fails. A session that
+  // nothing behind. It resolves only once the process has exited and the
+  // fleet without the session is durable, and throws when the kill, the
+  // exit, or that write cannot be confirmed. A session whose exit is not
+  // confirmed stays listed, so the spawn's caller can still find it, and
+  // no revive starts it until its process is found gone. A session that
   // never registered is left alone.
   async removeFailedSpawn(id: SessionID): Promise<void> {
     const s = this.sessions.find((x) => x.id === id);
@@ -903,13 +944,87 @@ export class SessionManager {
       return;
     }
 
-    s.pty?.kill();
-    s.pty = null;
+    const pty = s.pty;
 
-    this.remove(s);
-    this.emitChange();
+    this.rollingBack.add(id);
 
-    await this.writeFleet();
+    try {
+      if (pty !== null) {
+        await this.killFailedSpawn(id, pty);
+      }
+
+      // A harness other than the killed one would outlive the rollback.
+      if (s.pty !== null && s.pty !== pty) {
+        throw new Error(`session ${id} started another harness while its spawn rolled back`);
+      }
+
+      s.pty = null;
+
+      this.remove(s);
+      this.emitChange();
+
+      await this.writeFleet();
+    } finally {
+      this.rollingBack.delete(id);
+    }
+  }
+
+  // Kills a failed spawn's harness and waits for its exit. A kill that
+  // throws or an exit it cannot confirm leaves the harness recorded as
+  // possibly running, which blocks revives of the session.
+  private async killFailedSpawn(id: SessionID, pty: HarnessHandle): Promise<void> {
+    this.unconfirmedKills.set(id, pty);
+    pty.kill();
+
+    const exited = await this.waitForKilledExit(pty);
+
+    if (!exited) {
+      throw new Error(`session ${id} did not exit after its kill`);
+    }
+
+    this.unconfirmedKills.delete(id);
+  }
+
+  // Waits for a killed harness to exit, ending it with a signal it cannot
+  // ignore when the first wait runs out and its provider can send one.
+  private async waitForKilledExit(pty: HarnessHandle): Promise<boolean> {
+    const exited = await pty.waitForExit(FAILED_SPAWN_EXIT_WAIT_MS);
+
+    if (exited) {
+      return true;
+    }
+
+    if (pty.killForced === undefined) {
+      return false;
+    }
+
+    pty.killForced();
+
+    return pty.waitForExit(FAILED_SPAWN_EXIT_WAIT_MS);
+  }
+
+  // Whether a failed spawn's rollback leaves the session free to revive:
+  // none is running, and any process it killed is gone now. A session whose
+  // killed process may still run stays blocked until a later check finds it
+  // gone.
+  private async isRollbackSettled(id: SessionID): Promise<boolean> {
+    if (this.rollingBack.has(id)) {
+      return false;
+    }
+
+    const killed = this.unconfirmedKills.get(id);
+
+    if (killed !== undefined) {
+      const gone = await killed.waitForExit(0);
+
+      if (!gone) {
+        return false;
+      }
+
+      this.unconfirmedKills.delete(id);
+    }
+
+    return true;
   }
 
   collectDescriptors(): SessionDescriptor[] {
@@ -1147,7 +1262,8 @@ export class SessionManager {
   // every harness on it kept inside, and fails whole when the host stays
   // awake. A dead session on a target that can destroy its host is not
   // forgotten by a kill: forgetting it destroys the host, which takes a
-  // confirmed forget.
+  // confirmed forget. The set is taken before the first await, so a
+  // sub-session spawned while the kill waits on a host is not part of it.
   async kill(id: SessionID): Promise<void> {
     const s = this.sessions.find((x) => x.id === id);
 
@@ -1155,11 +1271,15 @@ export class SessionManager {
       return;
     }
 
+    const children = this.collectChildren(id);
+
     if (s.pty) {
       await this.stopHarness(s);
 
-      for (const child of this.collectChildren(id)) {
-        await this.tryStopHarness(child);
+      for (const child of children) {
+        if (this.sessions.includes(child) && child.parent === id) {
+          await this.tryStopHarness(child);
+        }
       }
     } else {
       if (this.findProvider(s)?.capabilities.destroy === true) {
@@ -1170,8 +1290,7 @@ export class SessionManager {
         );
       }
 
-      this.updateForgottenChildren(id);
-      this.remove(s);
+      this.removeWithChildren(s, children);
     }
 
     await this.writeFleet();
@@ -1187,7 +1306,9 @@ export class SessionManager {
    * kept asleep in that host is refused until the host wakes. Its dead
    * sub-sessions on other hosts go with it, unless their own target can
    * destroy their host, and its live ones become top-level. A failed destroy
-   * throws before anything is forgotten.
+   * throws before anything is forgotten. The sub-sessions it may forget are
+   * taken before the first await; one spawned while the forget waits on a
+   * host becomes top-level.
    */
   async forget(id: SessionID): Promise<boolean> {
     const s = this.sessions.find((x) => x.id === id);
@@ -1196,6 +1317,7 @@ export class SessionManager {
       return false;
     }
 
+    const children = this.collectChildren(id);
     const provider = this.findProvider(s);
     const destroys = provider !== null && provider.capabilities.destroy && s.hostKey === s.id;
 
@@ -1229,8 +1351,7 @@ export class SessionManager {
       this.killTerminal(s);
     }
 
-    this.updateForgottenChildren(id);
-    this.remove(s);
+    this.removeWithChildren(s, children);
 
     await this.writeFleet();
 
@@ -1239,22 +1360,42 @@ export class SessionManager {
     return destroys;
   }
 
-  // A forgotten parent's dead sub-sessions go with it, except one whose own
-  // target can destroy its host: forgetting that one destroys the host, which
-  // takes its own confirmed forget. Every sub-session that stays becomes
-  // top-level.
-  private updateForgottenChildren(id: SessionID): void {
-    for (const child of this.collectChildren(id)) {
+  // Removes a forgotten session and its dead sub-sessions, except one whose
+  // own target can destroy its host: forgetting that one destroys the host,
+  // which takes its own confirmed forget. Only a sub-session among the given
+  // ones may go. Every sub-session that stays becomes top-level before the
+  // session's removal is announced, and is announced after it, so no
+  // announcement ever finds a sub-session whose parent is gone or a parent
+  // standing without the sub-sessions that leave with it.
+  private removeWithChildren(s: Session, forgettable: readonly Session[]): void {
+    const kept: Session[] = [];
+    const dropped: Session[] = [];
+
+    for (const child of this.collectChildren(s.id)) {
       const live = child.pty !== null || (child.kind === 'headless' && child.state !== 'exited');
 
-      if (live || this.findProvider(child)?.capabilities.destroy === true) {
+      if (
+        live ||
+        !forgettable.includes(child) ||
+        this.findProvider(child)?.capabilities.destroy === true
+      ) {
         child.parent = null;
 
-        this.onEvent('state', child);
+        kept.push(child);
       } else {
-        this.remove(child);
+        dropped.push(child);
       }
     }
+
+    this.remove(s);
+
+    for (const child of kept) {
+      this.onEvent('state', child);
+    }
+
+    for (const child of dropped) {
+      this.remove(child);
+    }
   }
 
   /**
@@ -1482,95 +1623,6 @@ export class SessionManager {
   }
 }
 
-export function countSessionStates(
-  list: readonly { readonly state: SessionState }[],
-): Record<SessionState, number> {
-  const c = { needs_you: 0, running: 0, done: 0, exited: 0 };
-
-  for (const s of list) {
-    c[s.state]++;
-  }
-
-  return c;
-}
-
-interface SortableSessionView {
-  readonly id: string;
-  readonly parent: string | null;
-  readonly state: SessionState;
-  readonly pinned: boolean;
-  readonly lastAttachedAt: number;
-  readonly createdAt: number;
-}
-
-// Overlay order: pinned sessions first in most-recently-attached order, then
-// everyone else by urgency — who needs you, finished turns, busy, dead —
-// with most-recently-attached breaking ties inside each state. A
-// sub-session sits directly under its parent, ranked among its siblings
-// alone, so its attention never moves the parent's row; a sub-session whose
-// parent is not listed ranks as a top-level row.
-export function sortSessionViews<T extends SortableSessionView>(list: readonly T[]): T[] {
-  const rank: Record<SessionState, number> = {
-    needs_you: 0,
-    done: 1,
-    running: 2,
-    exited: 3,
-  };
-
-  const ranked = [...list].toSorted((a, b) => {
-    if (a.pinned !== b.pinned) {
-      return a.pinned ? -1 : 1;
-    }
-
-    const recency = b.lastAttachedAt - a.lastAttachedAt || b.createdAt - a.createdAt;
-
-    return a.pinned ? recency : rank[a.state] - rank[b.state] || recency;
-  });
-
-  const listed = new Set(ranked.map((s) => s.id));
-
-  const sorted: T[] = [];
-
-  for (const s of ranked) {
-    if (s.parent !== null && listed.has(s.parent)) {
-      continue;
-    }
-
-    sorted.push(s, ...ranked.filter((child) => child.parent === s.id));
-  }
-
-  return sorted;
-}
-
-// Never a filesystem path, so a repository can't collide with it.
-export const PINNED_GROUP_KEY = ' pinned';
-
-// Overlay display order for the grouped view: the flat sort with each
-// repository's sessions pulled together at the position of its best-ranked
-// member, so the renderer's adjacency-based headers appear once per group.
-// Pinned sessions form their own leading group. A sub-session keys by its
-// parent, so a set never splits across groups.
-export function sortGroupedSessionViews<
-  T extends SortableSessionView & { readonly repoRoot: string },
->(list: readonly T[]): T[] {
-  const byID = new Map(list.map((s) => [s.id, s]));
-  const buckets = new Map<string, T[]>();
-
-  for (const s of sortSessionViews(list)) {
-    const owner = (s.parent === null ? undefined : byID.get(s.parent)) ?? s;
-    const key = owner.pinned ? PINNED_GROUP_KEY : owner.repoRoot;
-    const bucket = buckets.get(key);
-
-    if (bucket === undefined) {
-      buckets.set(key, [s]);
-    } else {
-      bucket.push(s);
-    }
-  }
-
-  return [...buckets.values()].flat();
-}
-
 function buildLifecycle(s: Session): SessionLifecycle {
   return buildSessionLifecycle({
     desired: s.desired,

package/src/daemon/start-headless-turn.ts

@@ -1,9 +1,9 @@
-import type { AdapterEvent } from '../agents/agent-adapter';
-import { truncateDetail } from '../agents/truncate-detail';
+import type { AdapterEvent } from '../protocol/adapter-event';
 import type { SessionID } from '../shared/session-id';
+import { truncateDetail } from '../shared/truncate-detail';
+import { truncateSummary } from '../shared/truncate-summary';
 import type { SessionRuntime } from './session-runtime';
 import type { SessionManager } from './sessions';
-import { truncateSummary } from './truncate-summary';
 
 /**
  * Starts one headless turn for a session: finds its adapter's headless

package/src/daemon/start-session-bridge.ts

@@ -1,12 +1,12 @@
 import { z } from 'zod';
+import type { HookEvent } from '../protocol/hook-event';
 import type { EventMsg } from '../protocol/protocol';
 import type { MessageID } from '../shared/message-id';
 import type { SessionID } from '../shared/session-id';
 import { toMessageID } from '../shared/to-message-id';
 import type { MessageRecord } from '../store/message-record';
-import type { TapClient } from './daemon-connection';
+import type { TapClient } from './daemon-context';
 import type { HarnessRelay } from './execution-provider';
-import type { HookEvent } from './hooks';
 import { isBindingCurrent } from './is-binding-current';
 import type { BridgeBinding } from './is-binding-current';
 import { parseHookLine } from './parse-hook-line';

package/src/daemon/tap-registry.ts

@@ -4,6 +4,23 @@ import type { SessionID } from '../shared/session-id';
 interface TapEntry<TClient> {
   readonly client: TClient;
   readonly claimed: Set<MessageID>;
+
+  // Whether the tap also takes messages held under the session's agent
+  // session id, as opposed to its atc id alone.
+  readonly linked: boolean;
+
+  // Which attach made this tap; every attach makes a new one.
+  readonly generation: number;
+}
+
+/**
+ * A session's tap as a drain finds it before it waits: the attach that made
+ * it and whether it is linked. A delivery made for it lands only while the
+ * same attach holds the tap.
+ */
+export interface TapGeneration {
+  readonly generation: number;
+  readonly linked: boolean;
 }
 
 /**
@@ -13,14 +30,24 @@ interface TapEntry<TClient> {
 export class TapRegistry<TClient> {
   private readonly bySession = new Map<SessionID, TapEntry<TClient>>();
 
+  private nextGeneration = 0;
+
   /**
    * Makes the client the session's tap and returns the client it displaced,
-   * or null when there was none or it is the same connection.
+   * or null when there was none or it is the same connection. A linked tap
+   * also takes the messages held under the session's agent session id.
    */
-  attach(sessionID: SessionID, client: TClient): TClient | null {
+  attach(sessionID: SessionID, client: TClient, linked = true): TClient | null {
     const previous = this.bySession.get(sessionID)?.client ?? null;
 
-    this.bySession.set(sessionID, { client, claimed: new Set() });
+    this.nextGeneration += 1;
+
+    this.bySession.set(sessionID, {
+      client,
+      claimed: new Set(),
+      linked,
+      generation: this.nextGeneration,
+    });
 
     return previous === client ? null : previous;
   }
@@ -45,6 +72,12 @@ export class TapRegistry<TClient> {
     return previous;
   }
 
+  findTap(sessionID: SessionID): TapGeneration | null {
+    const entry = this.bySession.get(sessionID);
+
+    return entry === undefined ? null : { generation: entry.generation, linked: entry.linked };
+  }
+
   hasTap(sessionID: SessionID): boolean {
     return this.bySession.has(sessionID);
   }
@@ -53,10 +86,15 @@ export class TapRegistry<TClient> {
     return this.bySession.get(sessionID)?.client === client;
   }
 
-  claimDelivery(sessionID: SessionID, messageID: MessageID): TClient | null {
+  /**
+   * Claims the message for the session's tap and returns its client, or
+   * null when the tap already took it or another attach replaced the tap
+   * the delivery was made for.
+   */
+  claimDelivery(sessionID: SessionID, messageID: MessageID, generation: number): TClient | null {
     const entry = this.bySession.get(sessionID);
 
-    if (entry === undefined || entry.claimed.has(messageID)) {
+    if (entry === undefined || entry.generation !== generation || entry.claimed.has(messageID)) {
       return null;
     }
 

package/src/mcp/collect-clients.ts

@@ -1,6 +1,6 @@
 import type { Kysely } from 'kysely';
 import { collectJSONStrings } from './collect-json-strings';
-import type { MCPAuthSchema } from './types';
+import type { MCPAuthSchema } from './mcp-auth-schema';
 
 /**
  * A client added with `atc clients add`.

package/src/mcp/collect-grants.ts

@@ -2,7 +2,7 @@ import type { Kysely } from 'kysely';
 import { GRANT_SCOPES } from '../shared/grant-scope';
 import type { GrantScope } from '../shared/grant-scope';
 import { collectJSONStrings } from './collect-json-strings';
-import type { MCPAuthSchema } from './types';
+import type { MCPAuthSchema } from './mcp-auth-schema';
 
 /**
  * One authorization a client holds: every refresh token rotated from one

package/src/mcp/find-client-name.ts

@@ -1,5 +1,5 @@
 import type { Kysely } from 'kysely';
-import type { MCPAuthSchema } from './types';
+import type { MCPAuthSchema } from './mcp-auth-schema';
 
 /**
  * The name a client was added under, or null for an unknown client id or a

package/src/mcp/mcp-auth-schema.ts

@@ -0,0 +1,37 @@
+/**
+ * The columns of the authorization server's database that atc reads and
+ * writes itself. better-auth owns the tables and their migrations; dates are
+ * ISO 8601 strings and string lists are JSON arrays.
+ */
+export interface MCPAuthSchema {
+  readonly oauthClient: {
+    readonly id: string;
+    readonly clientId: string;
+    readonly name: string | null;
+    readonly redirectUris: string;
+    readonly createdAt: string | null;
+  };
+  readonly oauthAccessToken: {
+    readonly id: string;
+    readonly token: string;
+    readonly clientId: string;
+    readonly sessionId: string | null;
+    readonly authorizationCodeId: string | null;
+  };
+  readonly oauthRefreshToken: {
+    readonly id: string;
+    readonly clientId: string;
+    readonly sessionId: string | null;
+    readonly authorizationCodeId: string | null;
+    readonly scopes: string;
+    readonly createdAt: string;
+    readonly expiresAt: string;
+    readonly revoked: string | null;
+  };
+  readonly oauthConsent: { readonly id: string; readonly clientId: string };
+  readonly oauthResource: { readonly identifier: string };
+  readonly session: { readonly id: string; readonly expiresAt: string };
+
+  // When each grant last reached /mcp, keyed by the grant's authorization code id.
+  readonly atc_grant_use: { readonly grant_id: string; readonly last_used_at: string };
+}

package/src/mcp/open-mcp-auth.ts

@@ -8,11 +8,11 @@ import { betterAuth } from 'better-auth';
 import type { BetterAuthPlugin } from 'better-auth';
 import { getMigrations } from 'better-auth/db/migration';
 import { Kysely, SqliteAdapter, SqliteIntrospector, SqliteQueryCompiler, sql } from 'kysely';
+import { BunSqliteDriver } from '../shared/bun-sqlite-driver';
 import { GRANT_SCOPES } from '../shared/grant-scope';
-import { BunSqliteDriver } from '../store/bun-sqlite-driver';
 import { buildOwnerPlugin } from './build-owner-plugin';
 import { deriveTokenHash } from './derive-token-hash';
-import type { MCPAuthSchema } from './types';
+import type { MCPAuthSchema } from './mcp-auth-schema';
 
 interface MCPAuthOptions {
   readonly dbPath: string;