125 lines
3.6 KiB
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
|
|
}
|