ollie/cmd/toolsrv/integration_test.go

// integration_test.go - End-to-end tests for 9P toolsrv with Tauth.
//
// These tests build and spawn a real toolsrv process, then connect via 9P
// to verify the authentication flow works correctly.
package main

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"
	"time"

	toolsrvclient "ollie/toolsrv/client"
	"ollie/toolsrv/metadata"
	"ollie/toolsrv/protocol"
)

var testToolsrvBinary string

// TestMain builds the toolsrv binary once for all integration tests.
func TestMain(m *testing.M) {
	// Build toolsrv binary to a temp location
	tmpDir, err := os.MkdirTemp("", "toolsrv-test-*")
	if err != nil {
		fmt.Fprintf(os.Stderr, "failed to create temp dir: %v\n", err)
		os.Exit(1)
	}
	defer os.RemoveAll(tmpDir)

	testToolsrvBinary = filepath.Join(tmpDir, "toolsrv")

	// Build the binary
	cmd := exec.Command("go", "build", "-o", testToolsrvBinary, "ollie/cmd/toolsrv")
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		fmt.Fprintf(os.Stderr, "failed to build toolsrv: %v\n", err)
		os.Exit(1)
	}

	os.Exit(m.Run())
}

// startTestServer starts a toolsrv process for testing.
// Returns socketPath and a cleanup function.
func startTestServer(t *testing.T) (socketPath string, cleanup func()) {
	t.Helper()

	tmpDir, err := os.MkdirTemp("", "toolsrv-socket-*")
	if err != nil {
		t.Fatalf("failed to create temp dir: %v", err)
	}

	// Create sandbox config in temp dir
	cfgDir := filepath.Join(tmpDir, "ollie")
	if err := os.MkdirAll(cfgDir, 0755); err != nil {
		os.RemoveAll(tmpDir)
		t.Fatalf("failed to create config dir: %v", err)
	}
	sandboxCfg := `filesystem:
  ro:
    - /proc
    - /sys
  rw:
    - /dev
  rox:
    - /usr
    - /lib
    - /lib64
    - /bin
    - /sbin
    - /etc
  rwx:
    - "{CWD}"
    - "{TMPDIR}"
    - "{HOME}"
network:
  unrestricted: true
`
	if err := os.WriteFile(filepath.Join(cfgDir, "sandbox.yaml"), []byte(sandboxCfg), 0644); err != nil {
		os.RemoveAll(tmpDir)
		t.Fatalf("failed to write sandbox.yaml: %v", err)
	}

	// Symlink the real tools directory so tests can find them
	realToolsPath := metadata.ToolsPath()
	if _, err := os.Stat(realToolsPath); err == nil {
		if err := os.Symlink(realToolsPath, filepath.Join(cfgDir, "tools")); err != nil {
			os.RemoveAll(tmpDir)
			t.Fatalf("failed to symlink tools dir: %v", err)
		}
	}

	socketPath = filepath.Join(tmpDir, "toolsrvclient.sock")
	cwd, _ := os.Getwd()

	ctx, cancel := context.WithCancel(context.Background())
	cmd := exec.CommandContext(ctx, testToolsrvBinary, "serve", "--cwd", cwd, "--listen", socketPath)
	cmd.Stderr = os.Stderr
	// Set session ID so tool registry is configured, and XDG_CONFIG_HOME so sandbox.yaml is found
	cmd.Env = append(os.Environ(),
		"OLLIE_SESSION_ID=test-session",
		"XDG_CONFIG_HOME="+tmpDir,
	)

	if err := cmd.Start(); err != nil {
		cancel()
		os.RemoveAll(tmpDir)
		t.Fatalf("failed to start toolsrv: %v", err)
	}

	// Wait for socket to be ready
	deadline := time.Now().Add(5 * time.Second)
	for time.Now().Before(deadline) {
		if _, err := os.Stat(socketPath); err == nil {
			break
		}
		time.Sleep(50 * time.Millisecond)
	}

	if _, err := os.Stat(socketPath); err != nil {
		cancel()
		cmd.Wait()
		os.RemoveAll(tmpDir)
		t.Fatalf("toolsrv socket not ready: %v", err)
	}

	return socketPath, func() {
		cancel()
		cmd.Wait()
		os.RemoveAll(tmpDir)
	}
}

func TestIntegration_FirstConnectionSetsSecret(t *testing.T) {
	socketPath, cleanup := startTestServer(t)
	defer cleanup()

	// First connection with a secret should succeed
	secret := "test-secret-first-connection"
	conn, err := toolsrvclient.Dial(socketPath, secret)
	if err != nil {
		t.Fatalf("first Dial failed: %v", err)
	}
	defer conn.Close()

	// Should have received a token
	if conn.Token() == "" {
		t.Error("expected non-empty token")
	}

	// Secret should be stored
	if conn.Secret() != secret {
		t.Errorf("Secret() = %q, want %q", conn.Secret(), secret)
	}
}

func TestIntegration_SameSecretReconnects(t *testing.T) {
	socketPath, cleanup := startTestServer(t)
	defer cleanup()

	secret := "test-secret-reconnect"

	// First connection establishes the secret
	conn1, err := toolsrvclient.Dial(socketPath, secret)
	if err != nil {
		t.Fatalf("first Dial failed: %v", err)
	}
	token1 := conn1.Token()
	conn1.Close()

	// Second connection with same secret should succeed
	conn2, err := toolsrvclient.Dial(socketPath, secret)
	if err != nil {
		t.Fatalf("second Dial failed: %v", err)
	}
	defer conn2.Close()

	// Should get the same token
	if conn2.Token() != token1 {
		t.Errorf("token changed: %q -> %q", token1, conn2.Token())
	}
}

func TestIntegration_WrongSecretFails(t *testing.T) {
	socketPath, cleanup := startTestServer(t)
	defer cleanup()

	// First connection establishes the secret
	conn1, err := toolsrvclient.Dial(socketPath, "correct-secret")
	if err != nil {
		t.Fatalf("first Dial failed: %v", err)
	}
	conn1.Close()

	// Second connection with wrong secret should fail
	conn2, err := toolsrvclient.Dial(socketPath, "wrong-secret")
	if err == nil {
		conn2.Close()
		t.Fatal("Dial with wrong secret should fail")
	}
}

func TestIntegration_EmptySecretGeneratesOne(t *testing.T) {
	socketPath, cleanup := startTestServer(t)
	defer cleanup()

	// Empty secret should generate a random one
	conn, err := toolsrvclient.Dial(socketPath, "")
	if err != nil {
		t.Fatalf("Dial with empty secret failed: %v", err)
	}
	defer conn.Close()

	// Should have generated and stored a secret
	if conn.Secret() == "" {
		t.Error("expected non-empty generated secret")
	}

	// Token should be set
	if conn.Token() == "" {
		t.Error("expected non-empty token")
	}
}

func TestIntegration_BasicOperations(t *testing.T) {
	socketPath, cleanup := startTestServer(t)
	defer cleanup()

	conn, err := toolsrvclient.Dial(socketPath, "test-secret-ops")
	if err != nil {
		t.Fatalf("Dial failed: %v", err)
	}
	defer conn.Close()
	conn.SetAgentID("test-agent")

	// Test Ping
	if err := conn.Ping(); err != nil {
		t.Errorf("Ping failed: %v", err)
	}

	// Test HostInfo
	info, err := conn.HostInfo()
	if err != nil {
		t.Errorf("HostInfo failed: %v", err)
	}
	if info.Platform == "" {
		t.Error("HostInfo.Platform is empty")
	}

	// Test ListTools (should be empty initially, but shouldn't error)
	tools, err := conn.ListTools()
	if err != nil {
		t.Errorf("ListTools failed: %v", err)
	}
	t.Logf("ListTools returned %d tools", len(tools))
}

func TestIntegration_ToolExecution(t *testing.T) {
	// Skip if tools directory doesn't exist (CI environment)
	if _, err := os.ReadDir(metadata.ToolsPath()); err != nil {
		t.Skipf("skipping: tools directory not available: %v", err)
	}

	socketPath, cleanup := startTestServer(t)
	defer cleanup()

	conn, err := toolsrvclient.Dial(socketPath, "test-secret-exec")
	if err != nil {
		t.Fatalf("Dial failed: %v", err)
	}
	defer conn.Close()
	conn.SetAgentID("test-agent")

	// Load the shell tool
	if err := conn.LoadTool("shell"); err != nil {
		t.Skipf("skipping: LoadTool failed (shell tool not available): %v", err)
	}

	// Verify it's loaded
	tools, err := conn.ListTools()
	if err != nil {
		t.Fatalf("ListTools failed: %v", err)
	}
	found := false
	for _, tool := range tools {
		if tool.Name == "shell" {
			found = true
			break
		}
	}
	if !found {
		t.Fatal("shell tool not found after loading")
	}

	// Execute a tool - this mimics what agent/runtime.go does
	ctx := context.Background()
	args := json.RawMessage(`{"cmd": "echo hello from test"}`)
	result, err := conn.CallTool(ctx, "shell", args)
	if err != nil {
		t.Fatalf("CallTool failed: %v", err)
	}

	// Parse the result the same way agent/runtime.go does
	var toolResult protocol.ToolResult
	if err := json.Unmarshal(result, &toolResult); err != nil {
		t.Fatalf("failed to parse protocol.ToolResult: %v (raw: %s)", err, string(result))
	}

	// Verify structure
	if len(toolResult.Content) == 0 {
		t.Fatal("protocol.ToolResult.Content is empty")
	}
	if toolResult.Content[0].Type != "text" {
		t.Errorf("expected type 'text', got %q", toolResult.Content[0].Type)
	}
	if !strings.Contains(toolResult.Content[0].Text, "hello from test") {
		t.Errorf("output doesn't contain expected text: %q", toolResult.Content[0].Text)
	}
	if toolResult.IsError {
		t.Error("protocol.ToolResult.IsError should be false")
	}

	t.Logf("Tool result: %+v", toolResult)
}

func TestIntegration_MultilineContent(t *testing.T) {
	// Skip if tools directory doesn't exist (CI environment)
	if _, err := os.ReadDir(metadata.ToolsPath()); err != nil {
		t.Skipf("skipping: tools directory not available: %v", err)
	}

	socketPath, cleanup := startTestServer(t)
	defer cleanup()

	conn, err := toolsrvclient.Dial(socketPath, "test-secret-multiline")
	if err != nil {
		t.Fatalf("Dial failed: %v", err)
	}
	defer conn.Close()
	conn.SetAgentID("test-agent")

	// Load the file_write tool
	if err := conn.LoadTool("file_write"); err != nil {
		t.Skipf("skipping: LoadTool failed (file_write tool not available): %v", err)
	}

	// Create a temp file path
	tmpFile := filepath.Join(t.TempDir(), "multiline_test.txt")

	// Test with multiline content including special characters
	multilineContent := "line one\nline two\nline three with \"quotes\"\nline four with 'single quotes'\nline five with \\backslash\n"

	ctx := context.Background()
	args := json.RawMessage(fmt.Sprintf(`{"path": %q, "content": %q}`, tmpFile, multilineContent))
	result, err := conn.CallTool(ctx, "file_write", args)
	if err != nil {
		t.Fatalf("CallTool failed: %v", err)
	}

	var toolResult protocol.ToolResult
	if err := json.Unmarshal(result, &toolResult); err != nil {
		t.Fatalf("failed to parse protocol.ToolResult: %v (raw: %s)", err, string(result))
	}

	if toolResult.IsError {
		t.Fatalf("file_write returned error: %s", toolResult.Content[0].Text)
	}

	// Read the file back and verify
	written, err := os.ReadFile(tmpFile)
	if err != nil {
		t.Fatalf("failed to read written file: %v", err)
	}

	if string(written) != multilineContent {
		t.Errorf("content mismatch:\nexpected: %q\ngot: %q", multilineContent, string(written))
	}
}

func TestIntegration_ToolExecutionError(t *testing.T) {
	// Skip if tools directory doesn't exist (CI environment)
	if _, err := os.ReadDir(metadata.ToolsPath()); err != nil {
		t.Skipf("skipping: tools directory not available: %v", err)
	}

	socketPath, cleanup := startTestServer(t)
	defer cleanup()

	conn, err := toolsrvclient.Dial(socketPath, "test-secret-exec-err")
	if err != nil {
		t.Fatalf("Dial failed: %v", err)
	}
	defer conn.Close()
	conn.SetAgentID("test-agent")

	// Load the shell tool
	if err := conn.LoadTool("shell"); err != nil {
		t.Skipf("skipping: LoadTool failed (shell tool not available): %v", err)
	}

	// Execute a command that fails
	ctx := context.Background()
	args := json.RawMessage(`{"cmd": "exit 42"}`)
	result, err := conn.CallTool(ctx, "shell", args)
	if err != nil {
		t.Fatalf("CallTool failed: %v", err)
	}

	// Parse the result
	var toolResult protocol.ToolResult
	if err := json.Unmarshal(result, &toolResult); err != nil {
		t.Fatalf("failed to parse protocol.ToolResult: %v (raw: %s)", err, string(result))
	}

	// Should be an error result with exit code info
	if len(toolResult.Content) == 0 {
		t.Fatal("protocol.ToolResult.Content is empty")
	}
	if !toolResult.IsError {
		t.Error("expected IsError=true for non-zero exit")
	}
	if !strings.Contains(toolResult.Content[0].Text, "exit status 42") && !strings.Contains(toolResult.Content[0].Text, "exit: 42") {
		t.Errorf("output doesn't contain exit code: %q", toolResult.Content[0].Text)
	}

	t.Logf("Error tool result: %+v", toolResult)
}

func TestIntegration_ConcurrentConnections(t *testing.T) {
	socketPath, cleanup := startTestServer(t)
	defer cleanup()

	secret := "test-secret-concurrent"

	// Establish the secret
	conn0, err := toolsrvclient.Dial(socketPath, secret)
	if err != nil {
		t.Fatalf("initial Dial failed: %v", err)
	}
	expectedToken := conn0.Token()
	conn0.Close()

	// Open multiple concurrent connections
	const numConns = 5
	conns := make([]*toolsrvclient.Conn, numConns)
	errors := make([]error, numConns)

	for i := 0; i < numConns; i++ {
		conns[i], errors[i] = toolsrvclient.Dial(socketPath, secret)
	}

	// Check all succeeded with same token
	for i := 0; i < numConns; i++ {
		if errors[i] != nil {
			t.Errorf("connection %d failed: %v", i, errors[i])
			continue
		}
		if conns[i].Token() != expectedToken {
			t.Errorf("connection %d has different token", i)
		}
		conns[i].Close()
	}
}

func TestIntegration_ProcessKeeperRespawn(t *testing.T) {
	// This test verifies that when a toolsrv process dies,
	// the ProcessKeeper respawns it with a new secret.

	cwd, _ := os.Getwd()

	// We need to use the real Spawn which requires the binary in PATH
	// For now, manually simulate what ProcessKeeper does

	// Start first server
	socketPath1, cleanup1 := startTestServer(t)

	secret1 := "secret-for-first-server"
	conn1, err := toolsrvclient.Dial(socketPath1, secret1)
	if err != nil {
		cleanup1()
		t.Fatalf("first Dial failed: %v", err)
	}
	token1 := conn1.Token()
	conn1.Close()

	// Kill the server
	cleanup1()

	// Start second server (simulates respawn)
	socketPath2, cleanup2 := startTestServer(t)
	defer cleanup2()

	// New server should accept a new secret (any secret, since it's fresh)
	secret2 := "secret-for-second-server"
	conn2, err := toolsrvclient.Dial(socketPath2, secret2)
	if err != nil {
		t.Fatalf("second Dial failed: %v", err)
	}
	token2 := conn2.Token()
	conn2.Close()

	// Tokens should be different (different servers)
	if token1 == token2 {
		t.Error("tokens should be different after respawn")
	}

	// Old secret should NOT work on new server
	conn3, err := toolsrvclient.Dial(socketPath2, secret1)
	if err == nil {
		conn3.Close()
		t.Error("old secret should not work on new server")
	}

	t.Logf("Successfully verified: old_token=%s new_token=%s", token1[:8], token2[:8])
	_ = cwd // silence unused warning
}

// TestIntegration_ToolCancellation tests that canceling the context during a tool
// call causes the call to return quickly (not wait for the full tool execution).
func TestIntegration_ToolCancellation(t *testing.T) {
	// Skip if tools directory doesn't exist (CI environment)
	if _, err := os.ReadDir(metadata.ToolsPath()); err != nil {
		t.Skipf("skipping: tools directory not available: %v", err)
	}

	socketPath, cleanup := startTestServer(t)
	defer cleanup()

	conn, err := toolsrvclient.Dial(socketPath, "test-secret-cancel")
	if err != nil {
		t.Fatalf("Dial failed: %v", err)
	}
	defer conn.Close()
	conn.SetAgentID("test-agent-cancel")

	// Load the shell tool
	if err := conn.LoadTool("shell"); err != nil {
		t.Skipf("skipping: LoadTool failed (shell tool not available): %v", err)
	}

	// Create a context that we'll cancel after the tool starts
	ctx, cancel := context.WithCancel(context.Background())

	// Run a long-running command
	args := json.RawMessage(`{"cmd": "sleep 30"}`)

	done := make(chan struct{})
	var callErr error
	var result json.RawMessage
	start := time.Now()

	go func() {
		defer close(done)
		result, callErr = conn.CallTool(ctx, "shell", args)
	}()

	// Wait a bit for the tool to start, then cancel
	time.Sleep(200 * time.Millisecond)
	cancel()

	// Wait for CallTool to return
	select {
	case <-done:
		// Expected
	case <-time.After(5 * time.Second):
		t.Fatal("CallTool did not return after context cancellation")
	}

	elapsed := time.Since(start)

	// The call should return quickly (not wait 30 seconds)
	// This is the key assertion - cancellation caused early return
	if elapsed > 3*time.Second {
		t.Errorf("CallTool took %v, expected quick return after cancellation", elapsed)
	}

	// Note: The error may be nil if the fid was closed before the read completed
	// The important thing is that it returned quickly
	t.Logf("Context cancellation worked: elapsed=%v, err=%v, result_len=%d", elapsed, callErr, len(result))
}

// TestIntegration_ToolCancellation_ContextPropagation tests that context cancellation
// properly propagates through the entire chain: client -> toolsrv -> tool execution.
// This verifies that the client returns quickly when canceled, even if the tool would
// take a long time to complete.
func TestIntegration_ToolCancellation_ContextPropagation(t *testing.T) {
	// Skip if tools directory doesn't exist (CI environment)
	if _, err := os.ReadDir(metadata.ToolsPath()); err != nil {
		t.Skipf("skipping: tools directory not available: %v", err)
	}

	socketPath, cleanup := startTestServer(t)
	defer cleanup()

	conn, err := toolsrvclient.Dial(socketPath, "test-secret-cancel-prop")
	if err != nil {
		t.Fatalf("Dial failed: %v", err)
	}
	defer conn.Close()
	conn.SetAgentID("test-agent-cancel-prop")

	// Load the shell tool
	if err := conn.LoadTool("shell"); err != nil {
		t.Skipf("skipping: LoadTool failed (shell tool not available): %v", err)
	}

	// Command: echo immediately, then sleep for a long time
	// We verify that cancellation works by checking the call returns quickly
	args := json.RawMessage(`{"cmd": "echo started && sleep 30 && echo finished"}`)

	ctx, cancel := context.WithCancel(context.Background())

	done := make(chan struct{})
	var callErr error
	var result json.RawMessage
	start := time.Now()

	go func() {
		defer close(done)
		result, callErr = conn.CallTool(ctx, "shell", args)
	}()

	// Wait a bit for the tool to start, then cancel
	time.Sleep(300 * time.Millisecond)
	cancel()

	// Wait for CallTool to return
	select {
	case <-done:
		t.Logf("CallTool returned: err=%v result=%s", callErr, string(result))
	case <-time.After(5 * time.Second):
		t.Fatal("CallTool did not return after cancellation")
	}

	elapsed := time.Since(start)

	// The call should return quickly (not wait 30 seconds)
	if elapsed > 3*time.Second {
		t.Errorf("CallTool took %v, expected quick return after cancellation", elapsed)
	}

	// If the result contains "finished", the process wasn't killed
	if result != nil && strings.Contains(string(result), "finished") {
		t.Error("result contains 'finished' - process was not killed on cancellation")
	}

	t.Logf("Context propagation test passed: elapsed=%v", elapsed)
}