add framebuffer state tracking to WASM client (match Dart)

- Transport: expose LastRecvOldNum/LastRecvNewNum for state tracking - WASM: stateTracker with shadow VT emulator and received states map - Each diff applied to base state copy, stored by newNum - Display diffing: only changed cells sent to terminal - Eliminates garbled output from missed diffs and state transitions - Build with -tags osusergo for WASM os/user compatibility
2026-03-30 20:47:18 -07:00
parent 4a40090a5f
commit 8d9bc2eba0
3 changed files with 228 additions and 22 deletions
@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"sync"
 	"syscall/js"
+	"time"

 	mosh "github.com/unixshells/mosh-go"
 )
@@ -16,23 +17,23 @@ func main() {
 	select {} // keep alive
 }

-// moshConnect(url, key) → Promise<MoshSession>
-// url: WebTransport URL (e.g., "https://relay.example.com/mosh/user/device")
-// key: base64-encoded mosh key
+// moshConnect(url, key, cols, rows) → Promise<MoshSession>
 func moshConnect(this js.Value, args []js.Value) interface{} {
-	if len(args) < 2 {
-		return reject("moshConnect requires (url, key)")
+	if len(args) < 4 {
+		return reject("moshConnect requires (url, key, cols, rows)")
 	}

 	url := args[0].String()
 	key := args[1].String()
+	cols := args[2].Int()
+	rows := args[3].Int()

 	handler := js.FuncOf(func(this js.Value, pargs []js.Value) interface{} {
 		resolve := pargs[0]
 		rejectFn := pargs[1]

 		go func() {
-			session, err := newSession(url, key)
+			session, err := newSession(url, key, cols, rows)
 			if err != nil {
 				rejectFn.Invoke(err.Error())
 				return
@@ -47,13 +48,14 @@ func moshConnect(this js.Value, args []js.Value) interface{} {
 }

 type session struct {
-	client *mosh.Client
-	conn   *wtConn
-	mu     sync.Mutex
-	closed bool
+	client  *mosh.Client
+	conn    *wtConn
+	tracker *stateTracker
+	mu      sync.Mutex
+	closed  bool
 }

-func newSession(url, key string) (*session, error) {
+func newSession(url, key string, cols, rows int) (*session, error) {
 	for len(key)%4 != 0 {
 		key += "="
 	}
@@ -71,14 +73,22 @@ func newSession(url, key string) (*session, error) {
 		return nil, err
 	}

-	// Use DialConnManual — no internal sendLoop. JS drives Tick() via setInterval.
 	client, err := mosh.DialConnManual(conn, ocb)
 	if err != nil {
 		conn.Close()
 		return nil, err
 	}

-	return &session{client: client, conn: conn}, nil
+	s := &session{
+		client:  client,
+		conn:    conn,
+		tracker: newStateTracker(cols, rows),
+	}
+
+	// Background goroutine: read from mosh client, feed to state tracker.
+	go s.recvLoop()
+
+	return s, nil
 }

 func (s *session) jsObject() js.Value {
@@ -97,7 +107,10 @@ func (s *session) jsObject() js.Value {
 		if len(args) < 2 {
 			return nil
 		}
-		s.client.Resize(uint16(args[0].Int()), uint16(args[1].Int()))
+		cols := args[0].Int()
+		rows := args[1].Int()
+		s.client.Resize(uint16(cols), uint16(rows))
+		s.tracker.resize(cols, rows)
 		s.client.Tick()
 		return nil
 	}))
@@ -108,11 +121,9 @@ func (s *session) jsObject() js.Value {
 		return nil
 	}))

-	// poll() — check for output. Called by JS setInterval.
-	// Returns string if output available, null otherwise.
-	// Using JS-driven polling avoids Go goroutine scheduler starvation.
+	// poll() — returns pre-diffed ANSI output from the state tracker.
 	obj.Set("poll", js.FuncOf(func(this js.Value, args []js.Value) interface{} {
-		out := s.client.Recv(0)
+		out := s.tracker.poll()
 		if out == nil {
 			return js.Null()
 		}
@@ -130,6 +141,27 @@ func (s *session) jsObject() js.Value {
 	return obj
 }

+// recvLoop reads raw diffs from the mosh client and feeds them
+// to the state tracker for framebuffer-based processing.
+func (s *session) recvLoop() {
+	for {
+		out := s.client.Recv(60 * time.Second)
+		if out == nil {
+			s.mu.Lock()
+			closed := s.closed
+			s.mu.Unlock()
+			if closed {
+				return
+			}
+			continue
+		}
+		t := s.client.Transport()
+		oldNum := t.LastRecvOldNum()
+		newNum := t.LastRecvNewNum()
+		s.tracker.applyDiff(out, oldNum, newNum)
+	}
+}
+
 func reject(msg string) js.Value {
 	return js.Global().Get("Promise").Call("reject", msg)
 }
@@ -0,0 +1,154 @@
+//go:build js && wasm
+
+package main
+
+import (
+	"sync"
+
+	mosh "github.com/unixshells/mosh-go"
+	vt "github.com/unixshells/vt-go"
+)
+
+// stateTracker implements Dart-style received state management.
+// Each incoming diff is applied to a copy of the base state (looked up
+// by oldNum), producing a new state stored by newNum. The display
+// always shows the latest received state, diffed against what's
+// currently on screen. This eliminates garbled output from missed
+// diffs and overlapping state transitions.
+type stateTracker struct {
+	mu sync.Mutex
+
+	shadow      *vt.Emulator
+	shadowState uint64
+	cols, rows  int
+
+	// Received states: stateNum → framebuffer snapshot.
+	states       map[uint64]*mosh.Framebuffer
+	latestState  uint64
+	displayedFB  *mosh.Framebuffer
+
+	// Output buffer: pre-diffed ANSI ready for the terminal.
+	output []byte
+}
+
+func newStateTracker(cols, rows int) *stateTracker {
+	return &stateTracker{
+		shadow: vt.NewEmulator(cols, rows),
+		cols:   cols,
+		rows:   rows,
+		states: make(map[uint64]*mosh.Framebuffer),
+	}
+}
+
+// applyDiff processes an incoming server diff using state tracking.
+func (st *stateTracker) applyDiff(diff []byte, oldNum, newNum uint64) {
+	st.mu.Lock()
+	defer st.mu.Unlock()
+
+	// Already have this state — skip.
+	if _, ok := st.states[newNum]; ok {
+		return
+	}
+
+	// Snapshot current shadow as oldNum if we don't have it.
+	if oldNum == st.shadowState {
+		if _, ok := st.states[oldNum]; !ok {
+			st.states[oldNum] = mosh.SnapshotEmulator(st.shadow, true)
+		}
+	}
+
+	// Need the base state to apply the diff.
+	base, ok := st.states[oldNum]
+	if !ok {
+		return
+	}
+
+	// Restore shadow to base state if needed.
+	if st.shadowState != oldNum {
+		st.shadow = vt.NewEmulator(st.cols, st.rows)
+		// Write the base framebuffer as ANSI to restore state.
+		if base != nil {
+			st.shadow.Write(base.Diff(nil))
+		}
+	}
+
+	// Apply diff: feed hoststrings to shadow.
+	instrs, err := mosh.UnmarshalHostMessage(diff)
+	if err != nil {
+		return
+	}
+	for _, hi := range instrs {
+		if len(hi.Hoststring) > 0 {
+			st.shadow.Write(hi.Hoststring)
+		}
+	}
+	st.shadowState = newNum
+
+	// Snapshot result.
+	snap := mosh.SnapshotEmulator(st.shadow, true)
+	st.states[newNum] = snap
+	if newNum > st.latestState {
+		st.latestState = newNum
+	}
+
+	// Bound states list.
+	if len(st.states) > 128 {
+		st.pruneStates()
+	}
+
+	// Display diff: compute ANSI difference between displayed and latest.
+	latest := st.states[st.latestState]
+	if latest != nil {
+		ansi := latest.Diff(st.displayedFB)
+		if len(ansi) > 0 {
+			st.output = append(st.output, ansi...)
+		}
+		st.displayedFB = latest
+	}
+}
+
+// poll returns accumulated ANSI output, or nil.
+func (st *stateTracker) poll() []byte {
+	st.mu.Lock()
+	defer st.mu.Unlock()
+	if len(st.output) == 0 {
+		return nil
+	}
+	out := st.output
+	st.output = nil
+	return out
+}
+
+// resize resets the state tracker for new dimensions.
+func (st *stateTracker) resize(cols, rows int) {
+	st.mu.Lock()
+	defer st.mu.Unlock()
+	st.cols = cols
+	st.rows = rows
+	st.shadow = vt.NewEmulator(cols, rows)
+	st.shadowState = 0
+	st.states = make(map[uint64]*mosh.Framebuffer)
+	st.displayedFB = nil
+}
+
+func (st *stateTracker) pruneStates() {
+	if len(st.states) <= 64 {
+		return
+	}
+	// Keep the 64 most recent states.
+	var nums []uint64
+	for n := range st.states {
+		nums = append(nums, n)
+	}
+	// Simple sort for small lists.
+	for i := range nums {
+		for j := i + 1; j < len(nums); j++ {
+			if nums[j] < nums[i] {
+				nums[i], nums[j] = nums[j], nums[i]
+			}
+		}
+	}
+	for _, n := range nums[:len(nums)-64] {
+		delete(st.states, n)
+	}
+}
@@ -34,10 +34,12 @@ type Transport struct {
 	pendingDataAck bool   // true = send ack ASAP (received data, not just ack)

 	// Incoming state — list of received state nums for old_num validation.
-	receivedNums []uint64 // ordered list of state nums we have
-	ackNum       uint64   // latest received state num
-	sentAckNum   uint64   // last ackNum we actually sent on wire
-	throwawayNum uint64   // oldest state we still hold
+	receivedNums   []uint64 // ordered list of state nums we have
+	ackNum         uint64   // latest received state num
+	sentAckNum     uint64   // last ackNum we actually sent on wire
+	throwawayNum   uint64   // oldest state we still hold
+	lastRecvOldNum uint64   // oldNum from most recently received diff
+	lastRecvNewNum uint64   // newNum from most recently received diff

 	// Sequence counter for the crypto layer (independent of SSP state numbering).
 	seqOut      uint64
@@ -326,6 +328,10 @@ func (t *Transport) Recv(wire []byte) []byte {
 		t.receivedNums = filtered
 	}

+	// Track oldNum/newNum for state management.
+	t.lastRecvOldNum = ti.OldNum
+	t.lastRecvNewNum = ti.NewNum
+
 	// Add new state.
 	t.receivedNums = append(t.receivedNums, ti.NewNum)

@@ -359,6 +365,20 @@ func (t *Transport) SentNum() uint64 {
 	return t.sentNum
 }

+// LastRecvOldNum returns the oldNum from the most recently received diff.
+func (t *Transport) LastRecvOldNum() uint64 {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	return t.lastRecvOldNum
+}
+
+// LastRecvNewNum returns the newNum from the most recently received diff.
+func (t *Transport) LastRecvNewNum() uint64 {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	return t.lastRecvNewNum
+}
+
 // LastRecv returns the time of the last received datagram.
 func (t *Transport) LastRecv() time.Time {
 	t.mu.Lock()