agentic-coding-workflow/internal/terminal/session_test.go

package terminal

import (
	"context"
	"errors"
	"runtime"
	"strings"
	"testing"
	"time"
)

// testShell 返回各平台一个可交互读 stdin 的 shell。
func testShell() string {
	if runtime.GOOS == "windows" {
		return "cmd"
	}
	return "/bin/sh"
}

func newTestSup(t *testing.T, max int) *Supervisor {
	t.Helper()
	sup := NewSupervisor(SupervisorConfig{
		Shell:       testShell(),
		MaxSessions: max,
		KillGrace:   2 * time.Second,
		EnvWhitelist: []string{
			"PATH", "HOME", "TERM",
			"SystemRoot", "TEMP", "TMP", "USERPROFILE", "ComSpec", "PATHEXT",
		},
	}, nil)
	t.Cleanup(func() {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		sup.CloseAll(ctx)
	})
	return sup
}

// readUntil 在独立 goroutine 中持续读会话输出，直到累计内容包含 marker 或超时。
// 阻塞的 Read goroutine 会在会话被关闭（defer/Cleanup）后随 Read 报错而退出。
func readUntil(sess *Session, marker string, d time.Duration) string {
	out := make(chan []byte, 16)
	go func() {
		buf := make([]byte, 4096)
		for {
			n, err := sess.Read(buf)
			if n > 0 {
				b := make([]byte, n)
				copy(b, buf[:n])
				out <- b
			}
			if err != nil {
				close(out)
				return
			}
		}
	}()
	deadline := time.After(d)
	var acc strings.Builder
	for {
		select {
		case b, ok := <-out:
			if !ok {
				return acc.String()
			}
			acc.Write(b)
			if strings.Contains(acc.String(), marker) {
				return acc.String()
			}
		case <-deadline:
			return acc.String()
		}
	}
}

// TestOpen_ReadsShellOutput 验证开会话后向 shell 写命令能读回其输出。
// Unix 走真 PTY；Windows 走降级管道实现。
func TestOpen_ReadsShellOutput(t *testing.T) {
	sup := newTestSup(t, 5)
	sess, err := sup.Open(OpenParams{})
	if err != nil {
		t.Fatalf("open: %v", err)
	}
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	defer sup.Close(ctx, sess.ID)

	const marker = "ACWMARK42"
	cmd := "echo " + marker + "\n"
	if runtime.GOOS == "windows" {
		cmd = "echo " + marker + "\r\n"
	}
	if _, err := sess.Write([]byte(cmd)); err != nil {
		t.Fatalf("write: %v", err)
	}

	got := readUntil(sess, marker, 5*time.Second)
	if !strings.Contains(got, marker) {
		t.Fatalf("marker %q not found in shell output; got %q", marker, got)
	}
}

// TestOpen_MaxSessions 验证达到上限后再开会话返回 ErrTooManySessions。
func TestOpen_MaxSessions(t *testing.T) {
	sup := newTestSup(t, 1)
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	s1, err := sup.Open(OpenParams{})
	if err != nil {
		t.Fatalf("first open: %v", err)
	}
	defer sup.Close(ctx, s1.ID)

	if _, err := sup.Open(OpenParams{}); !errors.Is(err, ErrTooManySessions) {
		t.Fatalf("expected ErrTooManySessions, got %v", err)
	}
	if got := sup.Count(); got != 1 {
		t.Fatalf("expected 1 active session, got %d", got)
	}
}

// TestClose_TerminatesSession 验证 Close 终止 shell（done 关闭）并从映射移除。
func TestClose_TerminatesSession(t *testing.T) {
	sup := newTestSup(t, 5)
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	sess, err := sup.Open(OpenParams{})
	if err != nil {
		t.Fatalf("open: %v", err)
	}
	if got := sup.Count(); got != 1 {
		t.Fatalf("expected 1 active session, got %d", got)
	}

	sup.Close(ctx, sess.ID)

	select {
	case <-sess.Done():
	case <-time.After(8 * time.Second):
		t.Fatal("session not terminated within 8s after Close")
	}

	// monitor 在 done 关闭后异步移除映射，轮询直至归零。
	deadline := time.Now().Add(3 * time.Second)
	for sup.Count() != 0 {
		if time.Now().After(deadline) {
			t.Fatalf("session not removed from supervisor; count=%d", sup.Count())
		}
		time.Sleep(20 * time.Millisecond)
	}
}

// TestCloseAll_ClosesEverySession 验证 CloseAll 终止全部会话。
func TestCloseAll_ClosesEverySession(t *testing.T) {
	sup := newTestSup(t, 5)
	for i := 0; i < 3; i++ {
		if _, err := sup.Open(OpenParams{}); err != nil {
			t.Fatalf("open %d: %v", i, err)
		}
	}
	if got := sup.Count(); got != 3 {
		t.Fatalf("expected 3 active sessions, got %d", got)
	}

	ctx, cancel := context.WithTimeout(context.Background(), 8*time.Second)
	defer cancel()
	sup.CloseAll(ctx)

	deadline := time.Now().Add(3 * time.Second)
	for sup.Count() != 0 {
		if time.Now().After(deadline) {
			t.Fatalf("sessions not all closed; count=%d", sup.Count())
		}
		time.Sleep(20 * time.Millisecond)
	}
}