coder/agent/agentsocket/context_test.go

package agentsocket_test

import (
	"context"
	"testing"

	"github.com/stretchr/testify/require"

	"cdr.dev/slog/v3"
	"github.com/coder/coder/v2/agent/agentcontext"
	"github.com/coder/coder/v2/agent/agentsocket"
	"github.com/coder/coder/v2/testutil"
)

// fakeContextManager is an in-memory agentsocket.ContextManager for tests.
type fakeContextManager struct {
	sources   []agentcontext.Source
	snapshot  agentcontext.Snapshot
	resyncErr error
	resynced  bool
}

func (f *fakeContextManager) Sources() []agentcontext.Source { return f.sources }

func (f *fakeContextManager) HasSource(path string) (string, bool) {
	for _, s := range f.sources {
		if s.Path == path {
			return s.Path, true
		}
	}
	return "", false
}

func (f *fakeContextManager) AddSource(s agentcontext.Source) (agentcontext.Source, error) {
	for _, existing := range f.sources {
		if existing.Path == s.Path {
			return existing, nil
		}
	}
	f.sources = append(f.sources, s)
	return s, nil
}

func (f *fakeContextManager) RemoveSource(path string) error {
	for i, s := range f.sources {
		if s.Path == path {
			f.sources = append(f.sources[:i], f.sources[i+1:]...)
			return nil
		}
	}
	return agentcontext.ErrSourceNotFound
}

func (f *fakeContextManager) Snapshot() agentcontext.Snapshot { return f.snapshot }

func (f *fakeContextManager) Resync(_ context.Context) (agentcontext.Snapshot, error) {
	if f.resyncErr != nil {
		return agentcontext.Snapshot{}, f.resyncErr
	}
	f.resynced = true
	return f.snapshot, nil
}

func TestDRPCAgentSocketService_Context(t *testing.T) {
	t.Parallel()

	t.Run("SourceCRUDAndSnapshot", func(t *testing.T) {
		t.Parallel()

		const sourcePath = "/home/coder/project"
		cm := &fakeContextManager{
			snapshot: agentcontext.Snapshot{
				Version: 7,
				Resources: []agentcontext.Resource{{
					ID:          "instruction_file:" + sourcePath + "/AGENTS.md",
					Kind:        agentcontext.KindInstructionFile,
					Source:      sourcePath + "/AGENTS.md",
					SourcePath:  sourcePath,
					SizeBytes:   42,
					Status:      agentcontext.StatusOK,
					Description: "be concise",
				}, {
					// A built-in resource (no source path) the show filter must skip.
					ID:     "instruction_file:/home/coder/.coder/AGENTS.md",
					Kind:   agentcontext.KindInstructionFile,
					Source: "/home/coder/.coder/AGENTS.md",
					Status: agentcontext.StatusOK,
				}},
			},
		}

		socketPath := testutil.AgentSocketPath(t)
		ctx := testutil.Context(t, testutil.WaitShort)
		server, err := agentsocket.NewServer(
			slog.Make().Leveled(slog.LevelDebug),
			agentsocket.WithPath(socketPath),
			agentsocket.WithContextManager(cm),
		)
		require.NoError(t, err)
		defer server.Close()

		client := newSocketClient(ctx, t, socketPath)

		// Add a source.
		src, err := client.AddContextSource(ctx, sourcePath)
		require.NoError(t, err)
		require.Equal(t, sourcePath, src.Path)

		// It shows up in the list.
		sources, err := client.ContextSources(ctx)
		require.NoError(t, err)
		require.Len(t, sources, 1)
		require.Equal(t, sourcePath, sources[0].Path)

		// Get the registered source.
		got, err := client.GetContextSource(ctx, sourcePath)
		require.NoError(t, err)
		require.Equal(t, sourcePath, got.Path)

		// Getting an unregistered source errors.
		_, err = client.GetContextSource(ctx, "/nope")
		require.Error(t, err)

		// Snapshot carries resources with their source path stamped.
		snap, err := client.GetContextSnapshot(ctx)
		require.NoError(t, err)
		require.EqualValues(t, 7, snap.Version)
		require.Len(t, snap.Resources, 2)
		require.Equal(t, agentcontext.KindInstructionFile.String(), snap.Resources[0].Kind)
		require.Equal(t, sourcePath, snap.Resources[0].SourcePath)
		require.EqualValues(t, 42, snap.Resources[0].SizeBytes)

		// Remove the source; removing again reports not found.
		require.NoError(t, client.RemoveContextSource(ctx, sourcePath))
		err = client.RemoveContextSource(ctx, sourcePath)
		require.Error(t, err)
		require.Contains(t, err.Error(), "not found")
	})

	t.Run("Resync", func(t *testing.T) {
		t.Parallel()

		cm := &fakeContextManager{snapshot: agentcontext.Snapshot{Version: 3}}
		socketPath := testutil.AgentSocketPath(t)
		ctx := testutil.Context(t, testutil.WaitShort)
		server, err := agentsocket.NewServer(
			slog.Make().Leveled(slog.LevelDebug),
			agentsocket.WithPath(socketPath),
			agentsocket.WithContextManager(cm),
		)
		require.NoError(t, err)
		defer server.Close()

		client := newSocketClient(ctx, t, socketPath)

		snap, err := client.ResyncContext(ctx)
		require.NoError(t, err)
		require.EqualValues(t, 3, snap.Version)
		require.True(t, cm.resynced)
	})

	t.Run("NoManagerErrors", func(t *testing.T) {
		t.Parallel()

		socketPath := testutil.AgentSocketPath(t)
		ctx := testutil.Context(t, testutil.WaitShort)
		// No WithContextManager: the context RPCs must fail cleanly.
		server, err := agentsocket.NewServer(
			slog.Make().Leveled(slog.LevelDebug),
			agentsocket.WithPath(socketPath),
		)
		require.NoError(t, err)
		defer server.Close()

		client := newSocketClient(ctx, t, socketPath)

		// Every context RPC independently guards a nil context manager;
		// exercise all of them so dropping a guard surfaces as a test
		// failure rather than an agent-killing nil dereference in a DRPC
		// handler.
		_, err = client.ContextSources(ctx)
		require.Error(t, err)
		_, err = client.GetContextSource(ctx, "/tmp/x")
		require.Error(t, err)
		_, err = client.AddContextSource(ctx, "/tmp/x")
		require.Error(t, err)
		err = client.RemoveContextSource(ctx, "/tmp/x")
		require.Error(t, err)
		_, err = client.GetContextSnapshot(ctx)
		require.Error(t, err)
		_, err = client.ResyncContext(ctx)
		require.Error(t, err)
	})
}