feat: expand frame type to support stdout/stderr chan

agent/internal/sshutil/sessions.go

@@ -1,9 +1,11 @@
 package sshutil
 
 import (
+	"errors"
 	"io"
 	"sync"
 
+	"github.com/tale/headplane/agent/internal/util"
 	"golang.org/x/crypto/ssh"
 )
 
@@ -12,56 +14,70 @@ type SessionContext struct {
 	Session *ssh.Session
 	Stdin   io.WriteCloser
 	Stdout  io.Reader
+	Stderr  io.Reader
 	InputCh chan []byte
 }
 
 var sessions = make(map[string]*SessionContext)
-var sessionsMu sync.RWMutex
+var sessionsLock sync.RWMutex
 
-func addSession(id string, session *ssh.Session) *SessionContext {
-	sessionsMu.Lock()
-	defer sessionsMu.Unlock()
+func registerSessionChans(id string, session *ssh.Session) (*SessionContext, error) {
+	log := util.GetLogger()
+
+	sessionsLock.Lock()
+	defer sessionsLock.Unlock()
 
 	if _, exists := sessions[id]; exists {
-		return nil // Session with this ID already exists
+		return sessions[id], nil
 	}
 
 	stdin, err := session.StdinPipe()
 	if err != nil {
-		return nil // Handle error appropriately in production code
+		return nil, errors.New("failed to create stdin pipe: " + err.Error())
 	}
 
 	stdout, err := session.StdoutPipe()
 	if err != nil {
-		stdin.Close() // Close stdin if stdout pipe creation fails
-		return nil // Handle error appropriately in production code
+		stdin.Close()
+		return nil, errors.New("failed to create stdout pipe: " + err.Error())
 	}
 
-	sessionContext := &SessionContext{
+	stderr, err := session.StderrPipe()
+	if err != nil {
+		stdin.Close()
+		return nil, errors.New("failed to create stderr pipe: " + err.Error())
+	}
+
+	ctx := &SessionContext{
 		ID:      id,
 		Session: session,
 		Stdin:   stdin,
 		Stdout:  stdout,
+		Stderr:  stderr,
+		// Buffered channel to queue input data
+		InputCh: make(chan []byte, 256),
 	}
 
-	sessions[id] = sessionContext
-	return sessionContext
+	sessions[id] = ctx
+	log.Debug("Registered session %s with stdin, stdout, and stderr pipes", id)
+	return ctx, nil
 }
 
-func GetSession(id string) *SessionContext {
-	sessionsMu.RLock()
-	defer sessionsMu.RUnlock()
+func lookupSession(id string) (*SessionContext, bool) {
+	sessionsLock.RLock()
+	defer sessionsLock.RUnlock()
 
-	return sessions[id] // Returns nil if session does not exist
+	sessionContext, exists := sessions[id]
+	return sessionContext, exists
 }
 
 func RemoveSession(id string) {
-	sessionsMu.Lock()
-	defer sessionsMu.Unlock()
+	sessionsLock.Lock()
+	defer sessionsLock.Unlock()
 
 	if sessionContext, exists := sessions[id]; exists {
-		sessionContext.Stdin.Close() // Close the stdin pipe
+		sessionContext.Stdin.Close()   // Close the stdin pipe
 		sessionContext.Session.Close() // Close the SSH session
-		delete(sessions, id) // Remove from the map
+		delete(sessions, id)           // Remove from the map
 	}
 }