agent-orchestrator/backend/internal/adapters/runtime/conpty/attach.go

125 lines
3.6 KiB
Go

// attach.go - conpty Attach: a loopback Stream over the B3 pty-host. Unlike
// tmux, conpty does not spawn an attach CLI; it dials the session's
// loopback host and speaks the B1 framing protocol directly. The host replays
// the scrollback Snapshot as the first MsgTerminalData on connect, so a fresh
// Read naturally yields the repaint first.
package conpty
import (
"context"
"encoding/json"
"fmt"
"io"
"sync"
"github.com/aoagents/agent-orchestrator/backend/internal/ports"
)
var _ ports.Attacher = (*Runtime)(nil)
// Attach opens a fresh attach Stream for the session by dialing its loopback
// pty-host. rows/cols size the host's PTY from birth when known (a MsgResize is
// sent right after connect). ctx cancellation closes the Stream.
func (r *Runtime) Attach(ctx context.Context, handle ports.RuntimeHandle, rows, cols uint16) (ports.Stream, error) {
sess := r.resolve(handle.ID)
if sess == nil {
return nil, fmt.Errorf("conpty: session %q not found", handle.ID)
}
conn, err := dialHost(sess.addr, dialTimeout)
if err != nil {
return nil, fmt.Errorf("conpty: dial host for %q: %w", handle.ID, err)
}
pr, pw := io.Pipe()
s := &loopbackStream{conn: conn, pr: pr, pw: pw}
// Pump host frames: MsgTerminalData payloads go into the pipe that Read
// drains. The first such frame is the scrollback snapshot, so the replay
// arrives before any live output.
go s.pump()
// ctx cancellation must terminate the stream (mirrors the unix/windows
// spawn paths closing the PTY on ctx.Done).
go func() {
<-ctx.Done()
_ = s.Close()
}()
if rows > 0 && cols > 0 {
if err := s.Resize(rows, cols); err != nil {
_ = s.Close()
return nil, err
}
}
return s, nil
}
// loopbackStream is a ports.Stream backed by a single loopback connection to the
// pty-host. The pump goroutine reframes host output into an io.Pipe so Read
// presents a plain byte stream; Write/Resize encode client frames onto the conn.
type loopbackStream struct {
conn io.ReadWriteCloser
pr *io.PipeReader
pw *io.PipeWriter
closeOnce sync.Once
}
// pump reads framed host messages and writes MsgTerminalData payloads into the
// pipe. It closes the pipe when the connection ends so Read returns EOF.
func (s *loopbackStream) pump() {
parser := NewMessageParser(func(msgType byte, payload []byte) {
if msgType == MsgTerminalData {
// Write blocks until Read drains, preserving back-pressure and order.
_, _ = s.pw.Write(payload)
}
})
buf := make([]byte, 4096)
for {
n, err := s.conn.Read(buf)
if n > 0 {
parser.Feed(buf[:n])
}
if err != nil {
_ = s.pw.CloseWithError(err)
return
}
}
}
func (s *loopbackStream) Read(p []byte) (int, error) { return s.pr.Read(p) }
func (s *loopbackStream) Write(p []byte) (int, error) {
frame, err := EncodeMessage(MsgTerminalInput, p)
if err != nil {
return 0, err
}
if _, err := s.conn.Write(frame); err != nil {
return 0, err
}
return len(p), nil
}
func (s *loopbackStream) Resize(rows, cols uint16) error {
payload, _ := json.Marshal(ResizePayload{Cols: int(cols), Rows: int(rows)})
frame, err := EncodeMessage(MsgResize, payload) // small JSON payload, never overflows uint32
if err != nil {
return err
}
_, err = s.conn.Write(frame)
return err
}
// Close closes the conn and the pipe. Idempotent. Closing the conn unblocks
// pump's Read, which then closes the pipe-writer too; closing both here makes
// Close safe to call directly (e.g. on ctx cancel) without waiting for pump.
func (s *loopbackStream) Close() error {
var err error
s.closeOnce.Do(func() {
err = s.conn.Close()
_ = s.pw.Close()
_ = s.pr.Close()
})
return err
}