pi-agent-rb 0.2.2 → 0.3.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: db7831746ed3cbbd46e215293bb6a6e34dff2e31915b47a733737aa91ae15e4b
-  data.tar.gz: 61af774b3c9125f05aa4cfcafe11e5af31459a7db8a796487ed8857f5b677c1c
+  metadata.gz: ec081bcb7d1b2bac67301f37239e002e9b32d0eeb065d80a36b4d217896e1485
+  data.tar.gz: f19a72e0adb62fb42a49e3a21f2bcdffe3d46b3b546752c66e36427469c04c6e
 SHA512:
-  metadata.gz: cb4424e7698e4481dbda464cd5df382a3a555316ee923d5ca5d55faf9692701b606932d753e13ebf5dacf94c878868c7455e401786e1b5667f0d8daaccbf5c15
-  data.tar.gz: ccc703e5e81c2bf3ff3b47fa2605a391eeea66e9f1cee7603fba83db2e74f948a7c4b14b90db79ade99a5d8024513619c8764595a9d9456a8af6ede8b34815d1
+  metadata.gz: 1493e49939b44a45ff13fa7652b855ff48e25c611ad59bb6ac00c725bec549e76297156a455b20668bc5641a7fb0d8bcadd3eb535227e4e988c06c004a34dd92
+  data.tar.gz: 9cbd030aa9ae76199aa5730aab067bfd447a15416e9845d17fa3632f8baffd341586d0ca11b423eb4b66da6baa837c92511ec2ad7a6703a85c40e8d5996396d4

data/CHANGELOG.md

@@ -7,6 +7,35 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [0.3.0] - 2026-07-31
+
+### Added
+- Transport death notification. When pi dies unexpectedly (OOM kill,
+  missing binary after spawn, sandbox teardown), the client now learns
+  immediately instead of waiting out the 30s ack / 300s event timeouts
+  with a generic `TimeoutError`:
+  - Transports may accept an `on_close:` callable alongside
+    `on_message:`/`on_stderr:` and invoke it exactly once, with a short
+    human-readable reason, when they reach a terminal state on their
+    own — never for a caller-initiated `#close`. `Transport::Subprocess`
+    implements this (reporting exit status or signal, and watching the
+    child process itself so a descendant holding the stdout pipe open
+    can't defer the notification). The contract is documented in
+    `Transport` and the README's "Custom transports" section.
+  - On notification, `Client` rejects in-flight request futures with the
+    new `PiAgent::TransportClosedError` (a `ProtocolError` subclass with
+    the reason on `#reason`) and wakes subscribers with a synthetic
+    `Client::TRANSPORT_CLOSED_TYPE` message, which `Session` event
+    streams re-raise as `TransportClosedError`. Later `request`/`notify`
+    calls fail fast with the same error, and subscribers registered
+    after the death get the notification replayed once.
+  - Backward compatible: `Client#start` passes `on_close:` only when the
+    factory accepts the keyword (or `**kwargs`); existing
+    `(on_message:, on_stderr:)` factories keep their previous
+    timeout-backstop behavior unchanged.
+  - Subscriber callbacks are isolated during fanout: one raising
+    subscriber no longer prevents the rest from receiving a message.
+
 ## [0.2.2] - 2026-07-30
 
 ### Added

data/README.md

@@ -220,10 +220,60 @@ end
 a pi extension vetoes the operation — that is an expected outcome, not
 an error.
 
+## Custom transports
+
+By default the client spawns `pi --mode rpc` as a local subprocess. Pass
+`transport_factory:` to run pi somewhere else — a container, a remote
+sandbox:
+
+```ruby
+factory = lambda do |on_message:, on_stderr:, on_close:|
+  MySandboxTransport.new(
+    sandbox: sandbox,
+    on_message: on_message, on_stderr: on_stderr, on_close: on_close
+  )
+end
+
+PiAgent.session(transport_factory: factory) do |session|
+  # pi runs inside the sandbox; the protocol flows through your transport
+end
+```
+
+The factory receives the client's handlers and returns an object
+implementing `#start`, `#write(Hash)`, `#close(timeout:)`, and `#alive?`.
+[`Transport`](lib/pi_agent/transport.rb) documents the full contract;
+[`Transport::Subprocess`](lib/pi_agent/transport/subprocess.rb) is the
+reference implementation.
+
+`on_close:` is the transport's death notification. Invoke it exactly
+once, with a short human-readable reason (e.g.
+`"process terminated by signal 9"`), when the transport reaches a
+terminal state *on its own* — process exit, read-stream EOF, fatal stream
+error — and only after delivering any stdout messages already read. Do
+not invoke it for a shutdown initiated through `#close`. When it fires,
+the client fails in-flight requests and live event streams promptly with
+`PiAgent::TransportClosedError` (reason on `#reason`) instead of letting
+them wait out the 30s ack / 300s event timeouts, and later
+`request`/`notify` calls fail fast with the same error. Subscribers that
+register after the death get the notification replayed once, so an event
+stream started late still ends promptly.
+
+`on_close:` is optional and backward compatible: the client inspects the
+factory's parameters and passes the keyword only when the factory accepts
+it (explicitly or via `**kwargs`). An existing
+`(on_message:, on_stderr:)` factory keeps working unchanged — the
+timeouts then remain the only backstop when pi dies.
+
 ## Errors
 
 - A failed RPC command (`success: false`) raises `PiAgent::CommandError`,
   which carries the failing `#command` name.
+- If the transport dies out from under the client (pi OOM-killed, sandbox
+  torn down), in-flight requests and event streams raise
+  `PiAgent::TransportClosedError` promptly, with the death reason on
+  `#reason`. A caller-initiated `close` never raises it. This requires
+  the transport to report death — the bundled subprocess transport does;
+  for custom transports see [Custom transports](#custom-transports).
 - Agent-side errors arrive *in* the event stream, not as exceptions —
   inspect them with `Event#error?`, `#error_message`, and `#error_reason`
   (`"aborted"` vs `"error"`). This covers `extension_error` events and

data/lib/pi_agent/client.rb

@@ -14,7 +14,10 @@ module PiAgent
   # By default the client spawns `pi --mode rpc` as a local subprocess.
   # Pass `transport_factory:` — a callable `(on_message:, on_stderr:) ->
   # transport` — to run pi somewhere else (e.g. inside a remote sandbox).
-  # See Transport for the transport contract.
+  # A factory may additionally accept `on_close:` to report transport
+  # death; the keyword is passed only when the factory accepts it, so
+  # older two-keyword factories keep working. See Transport for the
+  # transport contract.
   #
   # Since pi 0.79.0 project-local inputs (.pi/settings.json, project
   # extensions, resources, packages) are trust-gated, and in RPC mode pi
@@ -25,6 +28,12 @@ module PiAgent
     DEFAULT_BIN = "pi"
     DEFAULT_ARGS = ["--mode", "rpc"].freeze
 
+    # Synthetic message fanned out to subscribers when the transport dies,
+    # so event streams wake promptly instead of waiting out their timeouts.
+    # Never sent by pi (the "pi_agent/" prefix keeps it out of upstream's
+    # event namespace); carries the death reason under "reason".
+    TRANSPORT_CLOSED_TYPE = "pi_agent/transport_closed"
+
     attr_reader :bin
 
     def self.resolve_bin(override = nil)
@@ -64,13 +73,15 @@ module PiAgent
       @subscribers_mutex = Mutex.new
       @transport = nil
       @extension_ui = nil
+      # Close/death bookkeeping, guarded by @pending_mutex except
+      # @close_broadcast (@subscribers_mutex).
+      @caller_closed = false
+      @close_reason = nil
+      @close_broadcast = nil
     end
 
     def start
-      @transport = @transport_factory.call(
-        on_message: method(:handle_message),
-        on_stderr: method(:handle_stderr)
-      )
+      @transport = @transport_factory.call(**transport_callbacks)
       @extension_ui = ExtensionUI.new(
         writer: @transport,
         handler: @extension_ui_handler,
@@ -83,21 +94,44 @@ module PiAgent
     def request(type, params = {})
       id = next_id
       future = Future.new
-      @pending_mutex.synchronize { @pending[id] = future }
+      @pending_mutex.synchronize do
+        raise TransportClosedError, @close_reason if @close_reason
+
+        @pending[id] = future
+      end
       payload = { id: id, type: type }.merge(params)
-      @transport.write(payload)
+      begin
+        @transport.write(payload)
+      rescue StandardError => e
+        raise abandon_request(id, e)
+      end
       future
     end
 
     def notify(type, params = {})
+      @pending_mutex.synchronize do
+        raise TransportClosedError, @close_reason if @close_reason
+      end
       payload = { type: type }.merge(params)
-      @transport.write(payload)
+      begin
+        @transport.write(payload)
+      rescue StandardError => e
+        raise close_error_or(e)
+      end
     end
 
+    # Subscribers registered after a transport death still learn about it:
+    # the synthetic TRANSPORT_CLOSED_TYPE message is replayed to them once,
+    # outside any lock. (The death fanout snapshots subscribers atomically
+    # with recording the broadcast, so nobody sees it twice.)
     def subscribe(&block)
       raise ArgumentError, "subscribe requires a block" unless block
 
-      @subscribers_mutex.synchronize { @subscribers << block }
+      replay = @subscribers_mutex.synchronize do
+        @subscribers << block
+        @close_broadcast
+      end
+      deliver([block], replay) if replay
       block
     end
 
@@ -106,6 +140,10 @@ module PiAgent
     end
 
     def close
+      # A caller-initiated close is a clean shutdown: mark it first so the
+      # transport teardown's own death notification is ignored and never
+      # surfaces as a TransportClosedError.
+      @pending_mutex.synchronize { @caller_closed = true }
       # Drain extension UI handler threads while the transport is still
       # open so their responses can still be written.
       @extension_ui&.shutdown
@@ -125,14 +163,42 @@ module PiAgent
     def build_subprocess_factory(bin, args, env, cwd)
       @bin = self.class.resolve_bin(bin)
       command = [@bin, *Array(args)]
-      lambda do |on_message:, on_stderr:|
+      lambda do |on_message:, on_stderr:, on_close:|
         Transport::Subprocess.new(
           command: command, env: env, cwd: cwd,
-          on_message: on_message, on_stderr: on_stderr
+          on_message: on_message, on_stderr: on_stderr, on_close: on_close
         )
       end
     end
 
+    def transport_callbacks
+      callbacks = {
+        on_message: method(:handle_message),
+        on_stderr: method(:handle_stderr)
+      }
+      callbacks[:on_close] = method(:handle_transport_close) if factory_accepts_on_close?
+      callbacks
+    end
+
+    # External factories predate on_close; pass it only to callables that
+    # declare the keyword (or **kwargs), so factories with the older
+    # `(on_message:, on_stderr:)` shape keep working unchanged.
+    def factory_accepts_on_close?
+      factory_parameters(@transport_factory).any? do |kind, name|
+        kind == :keyrest || (%i[keyreq key].include?(kind) && name == :on_close)
+      end
+    end
+
+    # Proc#parameters / Method#parameters describe the callable itself.
+    # Any other callable object may define an unrelated #parameters method,
+    # so inspect its #call method instead of trusting that name.
+    def factory_parameters(factory)
+      case factory
+      when Proc, Method then factory.parameters
+      else factory.method(:call).parameters
+      end
+    end
+
     def next_id
       @pending_mutex.synchronize do
         @next_id += 1
@@ -163,21 +229,76 @@ module PiAgent
     end
 
     def notify_subscribers(msg)
-      callbacks = @subscribers_mutex.synchronize { @subscribers.dup }
-      callbacks.each { |cb| cb.call(msg) }
+      deliver(@subscribers_mutex.synchronize { @subscribers.dup }, msg)
+    end
+
+    # Fan a message out, isolating each subscriber: one raising callback
+    # must not starve the rest (in particular, a raising logger must not
+    # keep a Session stream from seeing the death notification).
+    def deliver(callbacks, msg)
+      callbacks.each do |cb|
+        cb.call(msg)
+      rescue StandardError => e
+        handle_stderr("[pi-agent-rb] subscriber raised #{e.class}: #{e.message}")
+      end
     end
 
     def handle_stderr(line)
       # no-op by default; future versions may wire a logger
     end
 
-    def reject_pending(error)
+    # Transport-side death (process exit, stream EOF, fatal stream error).
+    # Idempotent, and a no-op after a caller-initiated #close — a clean
+    # shutdown is not an error. Fails all pending futures and wakes
+    # subscribers with a TRANSPORT_CLOSED_TYPE message so event streams
+    # end promptly instead of waiting out their timeouts. The broadcast is
+    # recorded atomically with the subscriber snapshot, so subscribers
+    # registered later get it replayed exactly once (see #subscribe).
+    def handle_transport_close(reason)
       pending = @pending_mutex.synchronize do
-        snapshot = @pending.values
-        @pending.clear
-        snapshot
+        return if @caller_closed || @close_reason
+
+        @close_reason = reason
+        take_pending
+      end
+      message = { "type" => TRANSPORT_CLOSED_TYPE, "reason" => reason }
+      callbacks = @subscribers_mutex.synchronize do
+        @close_broadcast = message
+        @subscribers.dup
       end
+      pending.each { |f| f.reject(TransportClosedError.new(reason)) }
+      deliver(callbacks, message)
+    end
+
+    def reject_pending(error)
+      pending = @pending_mutex.synchronize { take_pending }
       pending.each { |f| f.reject(error) }
     end
+
+    # Must be called holding @pending_mutex.
+    def take_pending
+      snapshot = @pending.values
+      @pending.clear
+      snapshot
+    end
+
+    # A write failed after the future was registered. Unregister it and,
+    # when the transport is dead, surface the death rather than the raw
+    # pipe error. Returns the error to raise; the future (if still ours)
+    # is rejected with the same error so it never dangles.
+    def abandon_request(id, error)
+      future = @pending_mutex.synchronize { @pending.delete(id) }
+      final = close_error_or(error)
+      future&.reject(final)
+      final
+    end
+
+    # Surface the death rather than the raw pipe error when the transport
+    # has already been reported dead. When the write error beats the
+    # notification, the raw error stands — it fails just as promptly.
+    def close_error_or(error)
+      reason = @pending_mutex.synchronize { @close_reason }
+      reason ? TransportClosedError.new(reason) : error
+    end
   end
 end

data/lib/pi_agent/errors.rb

@@ -9,6 +9,21 @@ module PiAgent
   class SessionError < Error; end
   class TimeoutError < Error; end
 
+  # Raised when the transport dies out from under the client — child
+  # process exit, read-stream EOF, fatal stream error — while requests or
+  # event streams are outstanding, and on any `request`/`notify` attempted
+  # after the death. A caller-initiated `Client#close` never raises it.
+  # `#reason` carries the transport's short description of what happened
+  # (e.g. "process terminated by signal 9").
+  class TransportClosedError < ProtocolError
+    attr_reader :reason
+
+    def initialize(reason)
+      @reason = reason
+      super("Transport closed: #{reason}")
+    end
+  end
+
   # Raised when an RPC command returns `success: false`. Carries the
   # failing command name so callers can branch on it.
   class CommandError < Error

data/lib/pi_agent/session.rb

@@ -19,6 +19,10 @@ module PiAgent
   # submit it as a streaming-aware prompt and drain it race-free. `events` is
   # a prompt-less drain for when you have already subscribed before processing
   # starts.
+  #
+  # If the transport dies mid-stream (pi killed, sandbox torn down), event
+  # streams raise TransportClosedError promptly instead of waiting out the
+  # event timeout — provided the transport reports death (see Transport).
   class Session
     # Max time to wait for the next event before assuming the agent stalled.
     DEFAULT_EVENT_TIMEOUT = 300
@@ -322,6 +326,10 @@ module PiAgent
       loop do
         msg = queue.pop(timeout: event_timeout)
         raise TimeoutError, "No event received within #{event_timeout}s" if msg.nil?
+        # Client's synthetic death notification: the transport is gone, no
+        # further events can arrive — end the stream now rather than
+        # timing out.
+        raise TransportClosedError, msg["reason"] if msg["type"] == Client::TRANSPORT_CLOSED_TYPE
 
         event = Event.new(msg)
         yielder << event

data/lib/pi_agent/transport/subprocess.rb

@@ -15,20 +15,29 @@ module PiAgent
     class Subprocess
       DEFAULT_CHUNK_SIZE = 4096
       DEFAULT_CLOSE_TIMEOUT = 5
+      # How long to let the stdout reader drain buffered output after the
+      # child exits before reporting the death (or closing) anyway — a
+      # descendant that inherited the pipe can hold it open indefinitely.
+      EXIT_DRAIN_TIMEOUT = 1
 
       attr_reader :pid
 
       # `cwd` sets the child's working directory — pi's built-in tools
       # (bash/read/edit/...) operate relative to it. nil leaves the
       # child in this process's working directory.
-      def initialize(command:, env: {}, cwd: nil, on_message: nil, on_stderr: nil)
+      #
+      # `on_close`, when given, is invoked exactly once with a short
+      # human-readable reason when the child exits on its own (see
+      # watch_exit). A caller-initiated #close is not reported.
+      def initialize(command:, env: {}, cwd: nil, on_message: nil, on_stderr: nil, on_close: nil)
         @command = Array(command)
         @env = env.transform_keys(&:to_s)
         @cwd = cwd
         @on_message = on_message
         @on_stderr = on_stderr
+        @on_close = on_close
         @write_mutex = Mutex.new
-        @closed = false
+        @owner_closed = false
       end
 
       def start
@@ -39,13 +48,14 @@ module PiAgent
         @stderr.binmode
         @stdout_thread = Thread.new { read_loop(@stdout, :stdout) }
         @stderr_thread = Thread.new { read_loop(@stderr, :stderr) }
+        @exit_watch_thread = Thread.new { watch_exit }
         self
       end
 
       def write(obj)
         payload = "#{JSON.generate(obj)}\n"
         @write_mutex.synchronize do
-          raise ProtocolError, "Transport closed" if @closed
+          raise ProtocolError, "Transport closed" if @owner_closed
 
           @stdin.write(payload)
           @stdin.flush
@@ -55,13 +65,18 @@ module PiAgent
       end
 
       def close(timeout: DEFAULT_CLOSE_TIMEOUT)
-        return if mark_closed!
+        return if mark_owner_closed!
 
         safe_close(@stdin)
         wait_for_exit(timeout)
-        [@stdout_thread, @stderr_thread].compact.each(&:join)
+        readers = [@stdout_thread, @stderr_thread].compact
+        # Bounded drain: a descendant that inherited a pipe can hold it
+        # open past the child's exit; force the readers out by closing
+        # the pipes after the window rather than hanging the close.
+        readers.each { |t| t.join(EXIT_DRAIN_TIMEOUT) }
         safe_close(@stdout)
         safe_close(@stderr)
+        (readers + [@exit_watch_thread]).compact.each(&:join)
       end
 
       def alive?
@@ -80,11 +95,13 @@ module PiAgent
         args
       end
 
-      def mark_closed!
+      # Marks that the owner requested shutdown via #close (as opposed to
+      # the child dying on its own). Returns whether it was already set.
+      def mark_owner_closed!
         @write_mutex.synchronize do
-          return true if @closed
+          return true if @owner_closed
 
-          @closed = true
+          @owner_closed = true
           false
         end
       end
@@ -117,6 +134,21 @@ module PiAgent
         # Pipe closed; reader exits normally. (EOFError descends from IOError.)
       end
 
+      # Sole death notifier. Watches the child process itself rather than
+      # the stdout pipe, whose EOF a descendant that inherited the write
+      # end can defer indefinitely. Gives the stdout reader a bounded
+      # window to drain buffered output first, so responses that raced the
+      # exit are dispatched before the notification. Owner-initiated
+      # #close marks @owner_closed before teardown, so the expected exit
+      # stays silent.
+      def watch_exit
+        status = @wait_thr.value
+        @stdout_thread&.join(EXIT_DRAIN_TIMEOUT)
+        return if @on_close.nil? || owner_closed?
+
+        @on_close.call(close_reason(status))
+      end
+
       def dispatch_stdout(line)
         msg = JSON.parse(line)
         @on_message&.call(msg)
@@ -124,6 +156,18 @@ module PiAgent
         @on_stderr&.call("[pi-agent-rb] invalid JSON on stdout: #{e.message}: #{line.inspect}")
       end
 
+      def owner_closed?
+        @write_mutex.synchronize { @owner_closed }
+      end
+
+      def close_reason(status)
+        if status.signaled?
+          "process terminated by signal #{status.termsig}"
+        else
+          "process exited with status #{status.exitstatus}"
+        end
+      end
+
       def terminate_process
         Process.kill("TERM", @pid)
       rescue Errno::ESRCH, Errno::EPERM

data/lib/pi_agent/transport.rb

@@ -6,10 +6,11 @@ module PiAgent
   # subprocess (Transport::Subprocess), or — via a caller-supplied
   # transport — a process inside a remote sandbox.
   #
-  # Contract. A transport is constructed with two callables:
+  # Contract. A transport is constructed with callables:
   #
   #   on_message: ->(Hash)    # one parsed JSON message from pi's stdout
   #   on_stderr:  ->(String)  # one line from pi's stderr
+  #   on_close:   ->(String)  # death notification (optional; see below)
   #
   # and responds to:
   #
@@ -18,6 +19,23 @@ module PiAgent
   #   #close(timeout:)  -> shut pi down
   #   #alive?           -> Boolean
   #
+  # Death notification (on_close). A transport MUST invoke on_close
+  # exactly once when it reaches a terminal state on its own — child
+  # process exit, read-stream EOF, fatal stream error — with a short
+  # human-readable reason (e.g. "process terminated by signal 9").
+  # Dispatch any already-read stdout messages via on_message *before*
+  # calling on_close, so responses that raced the death are not lost. A
+  # shutdown initiated by the owner via #close is not a death and must
+  # not be reported. On notification, Client rejects in-flight requests
+  # and ends Session event streams promptly with TransportClosedError
+  # instead of letting them wait out their timeouts.
+  #
+  # on_close is optional for compatibility: Client inspects the transport
+  # factory's parameters and passes the keyword only when the factory
+  # accepts it (explicitly or via **kwargs). A factory with the older
+  # `(on_message:, on_stderr:)` shape keeps its previous behavior —
+  # request and event timeouts remain the only death backstop.
+  #
   # Implementations own framing (pi speaks strict-LF JSONL — see Framer)
   # and the thread-safety of #write. Client injects its own handlers via
   # a transport factory, so transports never need settable callbacks.

data/lib/pi_agent/version.rb

@@ -1,7 +1,7 @@
 # frozen_string_literal: true
 
 module PiAgent
-  VERSION = "0.2.2"
+  VERSION = "0.3.0"
 
   # Pinned upstream pi-coding-agent version this gem is verified against.
   # See: https://www.npmjs.com/package/@earendil-works/pi-coding-agent

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: pi-agent-rb
 version: !ruby/object:Gem::Version
-  version: 0.2.2
+  version: 0.3.0
 platform: ruby
 authors:
 - chagel
@@ -85,7 +85,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.6.9
+rubygems_version: 4.0.10
 specification_version: 4
 summary: Ruby client for the pi coding agent (drives `pi --mode rpc`)
 test_files: []