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) }) }