package agent

import (

"bytes"

"context"

"crypto/tls"

"encoding/json"

"errors"

"fmt"

"hash/fnv"

"io"

"maps"

"net"

"net/http"

"net/netip"

"os"

"path/filepath"

"slices"

"strconv"

"strings"

"sync"

"time"

"github.com/go-chi/chi/v5"

"github.com/google/uuid"

"github.com/prometheus/client_golang/prometheus"

"github.com/prometheus/common/expfmt"

"github.com/spf13/afero"

"go.uber.org/atomic"

"golang.org/x/sync/errgroup"

"golang.org/x/xerrors"

googleproto "google.golang.org/protobuf/proto"

"google.golang.org/protobuf/types/known/timestamppb"

"tailscale.com/net/speedtest"

"tailscale.com/tailcfg"

"tailscale.com/types/netlogtype"

"tailscale.com/util/clientmetric"

"cdr.dev/slog/v3"

"github.com/coder/clistat"

"github.com/coder/coder/v2/agent/agentcontainers"

"github.com/coder/coder/v2/agent/agentcontext"

"github.com/coder/coder/v2/agent/agentcontextconfig"

"github.com/coder/coder/v2/agent/agentexec"

"github.com/coder/coder/v2/agent/agentfiles"

"github.com/coder/coder/v2/agent/agentgit"

"github.com/coder/coder/v2/agent/agentproc"

"github.com/coder/coder/v2/agent/agentscripts"

"github.com/coder/coder/v2/agent/agentsocket"

"github.com/coder/coder/v2/agent/agentssh"

"github.com/coder/coder/v2/agent/agenttoolcall"

"github.com/coder/coder/v2/agent/boundarylogproxy"

"github.com/coder/coder/v2/agent/proto"

"github.com/coder/coder/v2/agent/proto/resourcesmonitor"

"github.com/coder/coder/v2/agent/reconnectingpty"

"github.com/coder/coder/v2/agent/usershell"

"github.com/coder/coder/v2/agent/x/agentacp"

"github.com/coder/coder/v2/agent/x/agentdesktop"

"github.com/coder/coder/v2/agent/x/agentmcp"

"github.com/coder/coder/v2/buildinfo"

"github.com/coder/coder/v2/cli/gitauth"

"github.com/coder/coder/v2/coderd/database/dbtime"

"github.com/coder/coder/v2/codersdk"

"github.com/coder/coder/v2/codersdk/agentsdk"

"github.com/coder/coder/v2/codersdk/workspacesdk"

"github.com/coder/coder/v2/tailnet"

tailnetproto "github.com/coder/coder/v2/tailnet/proto"

"github.com/coder/quartz"

"github.com/coder/retry"

)

const (

ProtocolReconnectingPTY = "reconnecting-pty"

ProtocolSSH = "ssh"

ProtocolDial = "dial"

)

// EnvProcPrioMgmt determines whether we attempt to manage

// process CPU and OOM Killer priority.

const (

EnvProcPrioMgmt = "CODER_PROC_PRIO_MGMT"

EnvProcOOMScore = "CODER_PROC_OOM_SCORE"

)

var ErrAgentClosing = xerrors.New("agent is closing")

type Options struct {

Filesystem afero.Fs

LogDir string

TempDir string

ScriptDataDir string

Client Client

ReconnectingPTYTimeout time.Duration

EnvironmentVariables map[string]string

// EnvInfo overrides the session command environment source. Only

// tests set this. Nil defaults to usershell.SystemEnvInfo.

EnvInfo usershell.EnvInfoer

Logger slog.Logger

// IgnorePorts tells the api handler which ports to ignore when

// listing all listening ports. This is helpful to hide ports that

// are used by the agent, that the user does not care about.

IgnorePorts map[int]string

// ListeningPortsGetter is used to get the list of listening ports. Only

// tests should set this. If unset, a default that queries the OS will be used.

ListeningPortsGetter ListeningPortsGetter

SSHMaxTimeout time.Duration

TailnetListenPort uint16

Subsystems []codersdk.AgentSubsystem

PrometheusRegistry *prometheus.Registry

ReportMetadataInterval time.Duration

ServiceBannerRefreshInterval time.Duration

BlockFileTransfer bool

BlockReversePortForwarding bool

BlockLocalPortForwarding bool

Execer agentexec.Execer

Devcontainers bool

DevcontainerAPIOptions []agentcontainers.Option // Enable Devcontainers for these to be effective.

GitAPIOptions []agentgit.Option

Clock quartz.Clock

SocketServerEnabled bool

SocketPath string // Path for the agent socket server socket

AgentFirewallLogProxySocketPath string

ContextConfig agentcontextconfig.Config

// DERPTLSConfig is an optional TLS config for DERP connections.

DERPTLSConfig *tls.Config

// StatsReportInterval is the interval for the connstats callback

// installed at statsReporter creation.

StatsReportInterval time.Duration

}

type Client interface {

// ConnectRPC217WithRole connects to the Agent API v2.17. The workspace

// agent should use role "agent" to enable connection monitoring.

ConnectRPC217WithRole(ctx context.Context, role string) (

proto.DRPCAgentClient217, tailnetproto.DRPCTailnetClient28, error,

)

tailnet.DERPMapRewriter

agentsdk.RefreshableSessionTokenProvider

}

type Agent interface {

HTTPDebug() http.Handler

// TailnetConn may be nil.

TailnetConn() *tailnet.Conn

io.Closer

}

func New(options Options) Agent {

if options.Filesystem == nil {

options.Filesystem = afero.NewOsFs()

}

if options.EnvInfo == nil {

options.EnvInfo = &usershell.SystemEnvInfo{}

}

if options.TempDir == "" {

options.TempDir = os.TempDir()

}

if options.LogDir == "" {

if options.TempDir != os.TempDir() {

options.Logger.Debug(context.Background(), "log dir not set, using temp dir", slog.F("temp_dir", options.TempDir))

} else {

options.Logger.Debug(context.Background(), "using log dir", slog.F("log_dir", options.LogDir))

}

options.LogDir = options.TempDir

}

if options.ScriptDataDir == "" {

if options.TempDir != os.TempDir() {

options.Logger.Debug(context.Background(), "script data dir not set, using temp dir", slog.F("temp_dir", options.TempDir))

} else {

options.Logger.Debug(context.Background(), "using script data dir", slog.F("script_data_dir", options.ScriptDataDir))

}

options.ScriptDataDir = options.TempDir

}

if options.ReportMetadataInterval == 0 {

options.ReportMetadataInterval = time.Second

}

if options.ServiceBannerRefreshInterval == 0 {

options.ServiceBannerRefreshInterval = 2 * time.Minute

}

if options.Clock == nil {

options.Clock = quartz.NewReal()

}

prometheusRegistry := options.PrometheusRegistry

if prometheusRegistry == nil {

prometheusRegistry = prometheus.NewRegistry()

}

if options.Execer == nil {

options.Execer = agentexec.DefaultExecer

}

if options.StatsReportInterval == 0 {

options.StatsReportInterval = DefaultStatsReportInterval

}

if options.ListeningPortsGetter == nil {

options.ListeningPortsGetter = &osListeningPortsGetter{

cacheDuration: 1 * time.Second,

}

}

hardCtx, hardCancel := context.WithCancel(context.Background())

gracefulCtx, gracefulCancel := context.WithCancel(hardCtx)

a := &agent{

clock: options.Clock,

tailnetListenPort: options.TailnetListenPort,

reconnectingPTYTimeout: options.ReconnectingPTYTimeout,

logger: options.Logger,

gracefulCtx: gracefulCtx,

gracefulCancel: gracefulCancel,

hardCtx: hardCtx,

hardCancel: hardCancel,

coordDisconnected: make(chan struct{}),

environmentVariables: options.EnvironmentVariables,

client: options.Client,

filesystem: options.Filesystem,

logDir: options.LogDir,

tempDir: options.TempDir,

scriptDataDir: options.ScriptDataDir,

lifecycleUpdate: make(chan struct{}, 1),

lifecycleReported: make(chan codersdk.WorkspaceAgentLifecycle, 1),

lifecycleStates: []agentsdk.PostLifecycleRequest{{State: codersdk.WorkspaceAgentLifecycleCreated}},

connectionReporter: newConnectionReporter(hardCtx, options.Logger),

listeningPortsHandler: listeningPortsHandler{

getter: options.ListeningPortsGetter,

ignorePorts: maps.Clone(options.IgnorePorts),

},

reportMetadataInterval: options.ReportMetadataInterval,

statsReportInterval: options.StatsReportInterval,

announcementBannersRefreshInterval: options.ServiceBannerRefreshInterval,

sshMaxTimeout: options.SSHMaxTimeout,

envInfo: options.EnvInfo,

subsystems: options.Subsystems,

logSender: agentsdk.NewLogSender(options.Logger),

blockFileTransfer: options.BlockFileTransfer,

blockReversePortForwarding: options.BlockReversePortForwarding,

blockLocalPortForwarding: options.BlockLocalPortForwarding,

prometheusRegistry: prometheusRegistry,

metrics: newAgentMetrics(prometheusRegistry),

execer: options.Execer,

devcontainers: options.Devcontainers,

containerAPIOptions: options.DevcontainerAPIOptions,

gitAPIOptions: options.GitAPIOptions,

socketPath: options.SocketPath,

socketServerEnabled: options.SocketServerEnabled,

agentFirewallLogProxySocketPath: options.AgentFirewallLogProxySocketPath,

contextConfig: options.ContextConfig,

derpTLSConfig: options.DERPTLSConfig,

}

// Initially, we have a closed channel, reflecting the fact that we are not initially connected.

// Each time we connect we replace the channel (while holding the closeMutex) with a new one

// that gets closed on disconnection. This is used to wait for graceful disconnection from the

// coordinator during shut down.

close(a.coordDisconnected)

a.announcementBanners.Store(new([]codersdk.BannerConfig))

a.init()

return a

}

type agent struct {

clock quartz.Clock

logger slog.Logger

client Client

tailnetListenPort uint16

filesystem afero.Fs

logDir string

tempDir string

scriptDataDir string

listeningPortsHandler listeningPortsHandler

subsystems []codersdk.AgentSubsystem

reconnectingPTYTimeout time.Duration

reconnectingPTYServer *reconnectingpty.Server

// we track 2 contexts and associated cancel functions: "graceful" which is Done when it is time

// to start gracefully shutting down and "hard" which is Done when it is time to close

// everything down (regardless of whether graceful shutdown completed).

gracefulCtx context.Context

gracefulCancel context.CancelFunc

hardCtx context.Context

hardCancel context.CancelFunc

// closeMutex protects the following:

closeMutex sync.Mutex

closeWaitGroup sync.WaitGroup

coordDisconnected chan struct{}

closing bool

// note that once the network is set to non-nil, it is never modified, as with the statsReporter. So, routines

// that run after createOrUpdateNetwork and check the networkOK checkpoint do not need to hold the lock to use them.

network *tailnet.Conn

statsReporter *statsReporter

// end fields protected by closeMutex

environmentVariables map[string]string

manifest atomic.Pointer[agentsdk.Manifest] // manifest is atomic because values can change after reconnection.

// secrets are held separately from the manifest so that code paths that

// only need manifest data cannot accidentally access or leak secret

// values. Callers that need secrets must explicitly load this.

secrets atomic.Pointer[[]agentsdk.WorkspaceSecret]

reportMetadataInterval time.Duration

statsReportInterval time.Duration

scriptRunner *agentscripts.Runner

announcementBanners atomic.Pointer[[]codersdk.BannerConfig] // announcementBanners is atomic because it is periodically updated.

announcementBannersRefreshInterval time.Duration

sshServer *agentssh.Server

sshMaxTimeout time.Duration

envInfo usershell.EnvInfoer

blockFileTransfer bool

blockReversePortForwarding bool

blockLocalPortForwarding bool

lifecycleUpdate chan struct{}

lifecycleReported chan codersdk.WorkspaceAgentLifecycle

lifecycleMu sync.RWMutex // Protects following.

lifecycleStates []agentsdk.PostLifecycleRequest

lifecycleLastReportedIndex int // Keeps track of the last lifecycle state we successfully reported.

connectionReporter *connectionReporter

logSender *agentsdk.LogSender

// agentFirewallLogProxy is a socket server that forwards Agent Firewall audit logs to coderd.

// It may be nil if there is a problem starting the server.

agentFirewallLogProxy *boundarylogproxy.Server

agentFirewallLogProxySocketPath string

contextConfig agentcontextconfig.Config

prometheusRegistry *prometheus.Registry

// metrics are prometheus registered metrics that will be collected and

// labeled in Coder with the agent + workspace.

metrics *agentMetrics

execer agentexec.Execer

devcontainers bool

containerAPIOptions []agentcontainers.Option

containerAPI *agentcontainers.API

gitAPIOptions []agentgit.Option

filesAPI *agentfiles.API

gitAPI *agentgit.API

processAPI *agentproc.API

toolCalls *agenttoolcall.Table

desktopAPI *agentdesktop.API

mcpManager *agentmcp.Manager

acpManager *agentacp.Manager

acpAPI *agentacp.API

mcpAPI *agentmcp.API

contextConfigAPI *agentcontextconfig.API

contextManager *agentcontext.Manager

contextAPI *agentcontext.API

socketServerEnabled bool

socketPath string

socketServer *agentsocket.Server

derpTLSConfig *tls.Config

}

func (a *agent) TailnetConn() *tailnet.Conn {

a.closeMutex.Lock()

defer a.closeMutex.Unlock()

return a.network

}

// initialContextSources translates the boot-time

// CODER_AGENT_EXP_*_DIRS env vars into agentcontext.Source

// entries. This preserves the "set it on the template" workflow

// while the user-facing CLI for source CRUD ships in a

// follow-up.

func initialContextSources(cfg agentcontextconfig.Config, workingDir func() string) []agentcontext.Source {

base := ""

if workingDir != nil {

base = workingDir()

}

seen := make(map[string]struct{})

var sources []agentcontext.Source

add := func(path string) {

if path == "" {

return

}

if _, ok := seen[path]; ok {

return

}

seen[path] = struct{}{}

sources = append(sources, agentcontext.Source{Path: path})

}

for _, p := range agentcontextconfig.ResolvePaths(cfg.InstructionsDirs, base) {

add(p)

}

for _, p := range agentcontextconfig.ResolvePaths(cfg.SkillsDirs, base) {

add(p)

}

for _, p := range agentcontextconfig.ResolvePaths(cfg.MCPConfigFiles, base) {

add(p)

}

return sources

}

func (a *agent) init() {

// pass the "hard" context because we explicitly close the SSH server as part of graceful shutdown.

sshSrv, err := agentssh.NewServer(a.hardCtx, a.logger.Named("ssh-server"), a.prometheusRegistry, a.filesystem, a.execer, &agentssh.Config{

MaxTimeout: a.sshMaxTimeout,

MOTDFile: func() string { return a.manifest.Load().MOTDFile },

AnnouncementBanners: func() *[]codersdk.BannerConfig { return a.announcementBanners.Load() },

UpdateEnv: a.updateCommandEnv,

WorkingDirectory: func() string { return a.manifest.Load().Directory },

EnvInfo: a.envInfo,

BlockFileTransfer: a.blockFileTransfer,

BlockReversePortForwarding: a.blockReversePortForwarding,

BlockLocalPortForwarding: a.blockLocalPortForwarding,

ConnectionReporter: a.connectionReporter,

ExperimentalContainers: a.devcontainers,

})

if err != nil {

panic(err)

}

a.sshServer = sshSrv

a.scriptRunner = agentscripts.New(agentscripts.Options{

LogDir: a.logDir,

DataDirBase: a.scriptDataDir,

Logger: a.logger,

Clock: a.clock,

SSHServer: sshSrv,

Filesystem: a.filesystem,

GetScriptLogger: func(logSourceID uuid.UUID) agentscripts.ScriptLogger {

return a.logSender.GetScriptLogger(logSourceID)

},

})

// Register runner metrics. If the prom registry is nil, the metrics

// will not report anywhere.

a.scriptRunner.RegisterMetrics(a.prometheusRegistry)

containerAPIOpts := []agentcontainers.Option{

agentcontainers.WithExecer(a.execer),

agentcontainers.WithCommandEnv(a.sshServer.CommandEnv),

agentcontainers.WithScriptLogger(func(logSourceID uuid.UUID) agentcontainers.ScriptLogger {

return a.logSender.GetScriptLogger(logSourceID)

}),

}

containerAPIOpts = append(containerAPIOpts, a.containerAPIOptions...)

a.containerAPI = agentcontainers.NewAPI(a.logger.Named("containers"), containerAPIOpts...)

pathStore := agentgit.NewPathStore()

a.filesAPI = agentfiles.NewAPI(a.logger.Named("files"), a.filesystem, pathStore, agentfiles.WithEnvInfo(a.envInfo))

// workingDirFn reports the workspace directory ("" before the first manifest).

workingDirFn := func() string {

if m := a.manifest.Load(); m != nil {

return m.Directory

}

return ""

}

a.processAPI = agentproc.NewAPI(a.logger.Named("processes"), a.execer, a.filesystem, pathStore, a.envInfo, a.updateCommandEnv, workingDirFn)

a.toolCalls = agenttoolcall.New(a.clock, a.processAPI.CancelToolCall)

gitOpts := append([]agentgit.Option{agentgit.WithClock(a.clock), agentgit.WithExecer(a.execer)}, a.gitAPIOptions...)

a.gitAPI = agentgit.NewAPI(a.logger.Named("git"), pathStore, gitOpts...)

desktop := agentdesktop.NewPortableDesktop(

a.logger.Named("desktop"), a.execer, a.scriptRunner.ScriptBinDir(), nil,

)

a.desktopAPI = agentdesktop.NewAPI(a.logger.Named("desktop"), desktop, a.clock)

a.mcpManager = agentmcp.NewManager(a.gracefulCtx, a.logger.Named("mcp"), a.execer, a.filesystem, a.envInfo, a.updateCommandEnv, workingDirFn)

a.acpManager = agentacp.NewManager(a.gracefulCtx, agentacp.Options{Logger: a.logger.Named("acp"), Clock: a.clock, Execer: a.execer, Filesystem: a.filesystem, EnvInfo: a.envInfo, UpdateEnv: a.updateCommandEnv, WorkingDir: workingDirFn})

a.acpAPI = agentacp.NewAPI(a.acpManager)

a.contextConfigAPI = agentcontextconfig.NewAPI(workingDirFn, a.contextConfig)

a.mcpAPI = agentmcp.NewAPI(a.mcpManager)

// agentcontext.Manager is the new consolidated resolver,

// watcher, and pusher. It coexists with contextConfigAPI

// and the MCP manager during rollout. Initial sources are

// seeded from the existing CODER_AGENT_EXP_* env vars and

// from the agent's working directory at scan time.

a.contextManager = agentcontext.NewManager(agentcontext.ManagerOptions{

Logger: a.logger.Named("agentcontext"),

Clock: a.clock,

WorkingDir: workingDirFn,

InitialSources: initialContextSources(a.contextConfig, workingDirFn),

// The manager surfaces MCP servers and their tools as

// KindMCPServer resources by reading the shared MCP engine's

// catalog (a.mcpManager). That engine owns the single set of

// MCP server connections used for both discovery and tool-call

// execution, so each declared server is launched once.

MCPCatalog: func() []agentcontext.MCPServerStatus {

return mcpCatalogToContext(a.mcpManager.Catalog())

},

MCPReconnect: a.mcpManager.Reconnect,

ACPResources: a.acpManager.ContextResources,

})

a.contextAPI = agentcontext.NewAPI(a.contextManager)

// Re-resolve and re-push KindMCPServer resources whenever the MCP

// engine's catalog changes (startup connect, .mcp.json edits).

a.mcpManager.SetOnReload(a.contextManager.Trigger)

a.acpManager.SetOnChange(a.contextManager.Trigger)

a.reconnectingPTYServer = reconnectingpty.NewServer(

a.logger.Named("reconnecting-pty"),

a.sshServer,

a.connectionReporter,

a.metrics.connectionsTotal, a.metrics.reconnectingPTYErrors,

a.reconnectingPTYTimeout,

func(s *reconnectingpty.Server) {

s.ExperimentalContainers = a.devcontainers

},

)

a.initSocketServer()

a.startAgentFirewallLogProxyServer()

// Start the agentcontext manager's resolver/watcher loop.

// It runs for the lifetime of the agent and is closed in

// agent.Close. The push goroutine is started per-connection

// inside run() so it picks up the right drpc client.

go func() {

if err := a.contextManager.Run(a.gracefulCtx); err != nil && !errors.Is(err, context.Canceled) {

a.logger.Warn(a.gracefulCtx, "agentcontext manager run exited", slog.Error(err))

}

}()

go a.runLoop()

}

// initSocketServer initializes server that allows direct communication with a workspace agent using IPC.

func (a *agent) initSocketServer() {

if !a.socketServerEnabled {

a.logger.Info(a.hardCtx, "socket server is disabled")

return

}

server, err := agentsocket.NewServer(

a.logger.Named("socket"),

agentsocket.WithPath(a.socketPath),

agentsocket.WithContextManager(a.contextManager),

)

if err != nil {

a.logger.Error(a.hardCtx, "failed to create socket server", slog.Error(err), slog.F("path", a.socketPath))

return

}

a.socketServer = server

a.logger.Debug(a.hardCtx, "socket server started", slog.F("path", a.socketPath))

}

func (a *agent) startAgentFirewallLogProxyServer() {

if a.agentFirewallLogProxySocketPath == "" {

a.logger.Warn(a.hardCtx, "agent firewall log proxy socket path not defined; not starting proxy")

return

}

proxy := boundarylogproxy.NewServer(a.logger, a.agentFirewallLogProxySocketPath, a.prometheusRegistry)

if err := proxy.Start(); err != nil {

a.logger.Warn(a.hardCtx, "failed to start agent firewall log proxy", slog.Error(err))

return

}

a.agentFirewallLogProxy = proxy

a.logger.Info(a.hardCtx, "agent firewall log proxy server started",

slog.F("socket_path", a.agentFirewallLogProxySocketPath))

}

// runLoop attempts to start the agent in a retry loop.

// Coder may be offline temporarily, a connection issue

// may be happening, but regardless after the intermittent

// failure, you'll want the agent to reconnect.

func (a *agent) runLoop() {

// need to keep retrying up to the hardCtx so that we can send graceful shutdown-related

// messages.

ctx := a.hardCtx

defer a.logger.Info(ctx, "agent main loop exited")

for retrier := retry.New(100*time.Millisecond, 10*time.Second); retrier.Wait(ctx); {

a.logger.Info(ctx, "connecting to coderd")

err := a.run()

if err == nil {

continue

}

if ctx.Err() != nil {

// Context canceled errors may come from websocket pings, so we

// don't want to use `errors.Is(err, context.Canceled)` here.

a.logger.Warn(ctx, "runLoop exited with error", slog.Error(ctx.Err()))

return

}

if a.isClosed() {

a.logger.Warn(ctx, "runLoop exited because agent is closed")

return

}

if errors.Is(err, io.EOF) {

// Lost the connection to coderd. Flush the recorded debug history so the

// detail leading up to the disconnect is available in the logs.

a.logger.Flush(ctx)

a.logger.Warn(ctx, "disconnected from coderd",

codersdk.ConnectionDirectionServerToAgent.SlogField(),

codersdk.DisconnectReasonNetworkError.SlogField(),

codersdk.DisconnectReasonNetworkError.SlogExpectedField(),

codersdk.DisconnectInitiatorNetwork.SlogField(),

)

continue

}

// Flush the recorded debug history so the detail leading up to the failure

// is available in the logs.

a.logger.Flush(ctx)

a.logger.Warn(ctx, "run exited with error", slog.Error(err))

}

}

func (a *agent) collectMetadata(ctx context.Context, md codersdk.WorkspaceAgentMetadataDescription, now time.Time) *codersdk.WorkspaceAgentMetadataResult {

var out bytes.Buffer

result := &codersdk.WorkspaceAgentMetadataResult{

// CollectedAt is set here for testing purposes and overrode by

// coderd to the time of server receipt to solve clock skew.

//

// In the future, the server may accept the timestamp from the agent

// if it can guarantee the clocks are synchronized.

CollectedAt: now,

}

cmdPty, err := a.sshServer.CreateCommand(ctx, md.Script, nil, nil)

if err != nil {

result.Error = fmt.Sprintf("create cmd: %+v", err)

return result

}

cmd := cmdPty.AsExec()

cmd.Stdout = &out

cmd.Stderr = &out

cmd.Stdin = io.LimitReader(nil, 0)

// We split up Start and Wait instead of calling Run so that we can return a more precise error.

err = cmd.Start()

if err != nil {

result.Error = fmt.Sprintf("start cmd: %+v", err)

return result

}

// This error isn't mutually exclusive with useful output.

err = cmd.Wait()

const bufLimit = 10 << 10

if out.Len() > bufLimit {

err = errors.Join(

err,

xerrors.Errorf("output truncated from %v to %v bytes", out.Len(), bufLimit),

)

out.Truncate(bufLimit)

}

// Important: if the command times out, we may see a misleading error like

// "exit status 1", so it's important to include the context error.

err = errors.Join(err, ctx.Err())

if err != nil {

result.Error = fmt.Sprintf("run cmd: %+v", err)

}

result.Value = out.String()

return result

}

type metadataResultAndKey struct {

result *codersdk.WorkspaceAgentMetadataResult

key string

}

type trySingleflight struct {

mu sync.Mutex

m map[string]struct{}

}

func (t *trySingleflight) Do(key string, fn func()) {

t.mu.Lock()

_, ok := t.m[key]

if ok {

t.mu.Unlock()

return

}

t.m[key] = struct{}{}

t.mu.Unlock()

defer func() {

t.mu.Lock()

delete(t.m, key)

t.mu.Unlock()

}()

fn()

}

func (a *agent) reportMetadata(ctx context.Context, aAPI proto.DRPCAgentClient28) error {

tickerDone := make(chan struct{})

collectDone := make(chan struct{})

ctx, cancel := context.WithCancel(ctx)

defer func() {

cancel()

<-collectDone

<-tickerDone

}()

var (

logger = a.logger.Named("metadata")

report = make(chan struct{}, 1)

collect = make(chan struct{}, 1)

metadataResults = make(chan metadataResultAndKey, 1)

)

// Set up collect and report as a single ticker with two channels,

// this is to allow collection and reporting to be triggered

// independently of each other.

go func() {

t := time.NewTicker(a.reportMetadataInterval)

defer func() {

t.Stop()

close(report)

close(collect)

close(tickerDone)

}()

wake := func(c chan<- struct{}) {

select {

case c <- struct{}{}:

default:

}

}

wake(collect) // Start immediately.

for {

select {

case <-ctx.Done():

return

case <-t.C:

wake(report)

wake(collect)

}

}

}()

go func() {

defer close(collectDone)

var (

// We use a custom singleflight that immediately returns if there is already

// a goroutine running for a given key. This is to prevent a build-up of

// goroutines waiting on Do when the script takes many multiples of

// baseInterval to run.

flight = trySingleflight{m: map[string]struct{}{}}

lastCollectedAtMu sync.RWMutex

lastCollectedAts = make(map[string]time.Time)

)

for {

select {

case <-ctx.Done():

return

case <-collect:

}

manifest := a.manifest.Load()

if manifest == nil {

continue

}

// If the manifest changes (e.g. on agent reconnect) we need to

// purge old cache values to prevent lastCollectedAt from growing

// boundlessly.

lastCollectedAtMu.Lock()

for key := range lastCollectedAts {

if slices.IndexFunc(manifest.Metadata, func(md codersdk.WorkspaceAgentMetadataDescription) bool {

return md.Key == key

}) < 0 {

logger.Debug(ctx, "deleting lastCollected key, missing from manifest",

slog.F("key", key),

)

delete(lastCollectedAts, key)

}

}

lastCollectedAtMu.Unlock()

// Spawn a goroutine for each metadata collection, and use a

// channel to synchronize the results and avoid both messy

// mutex logic and overloading the API.

for _, md := range manifest.Metadata {

// We send the result to the channel in the goroutine to avoid

// sending the same result multiple times. So, we don't care about

// the return values.

go flight.Do(md.Key, func() {

ctx := slog.With(ctx, slog.F("key", md.Key))

lastCollectedAtMu.RLock()

collectedAt, ok := lastCollectedAts[md.Key]

lastCollectedAtMu.RUnlock()

if ok {

// If the interval is zero, we assume the user just wants

// a single collection at startup, not a spinning loop.

if md.Interval == 0 {

return

}

intervalUnit := time.Second

// reportMetadataInterval is only less than a second in tests,

// so adjust the interval unit for them.

if a.reportMetadataInterval < time.Second {

intervalUnit = 100 * time.Millisecond

}

// The last collected value isn't quite stale yet, so we skip it.

if collectedAt.Add(time.Duration(md.Interval) * intervalUnit).After(time.Now()) {

return

}

}

timeout := md.Timeout

if timeout == 0 {

if md.Interval != 0 {

timeout = md.Interval

} else if interval := int64(a.reportMetadataInterval.Seconds()); interval != 0 {

// Fallback to the report interval

timeout = interval * 3

} else {

// If the interval is still 0 (possible if the interval

// is less than a second), default to 5. This was

// randomly picked.

timeout = 5

}

}

ctxTimeout := time.Duration(timeout) * time.Second

ctx, cancel := context.WithTimeout(ctx, ctxTimeout)

defer cancel()

now := time.Now()

select {

case <-ctx.Done():

logger.Warn(ctx, "metadata collection timed out", slog.F("timeout", ctxTimeout))

case metadataResults <- metadataResultAndKey{

key: md.Key,

result: a.collectMetadata(ctx, md, now),

}:

lastCollectedAtMu.Lock()

lastCollectedAts[md.Key] = now

lastCollectedAtMu.Unlock()

}

})

}

}

}()

// Gather metadata updates and report them once every interval. If a

// previous report is in flight, wait for it to complete before

// sending a new one. If the network conditions are bad, we won't

// benefit from canceling the previous send and starting a new one.

var (

updatedMetadata = make(map[string]*codersdk.WorkspaceAgentMetadataResult)

reportTimeout = 30 * time.Second

reportError = make(chan error, 1)

reportInFlight = false

)

for {

select {

case <-ctx.Done():

return ctx.Err()

case mr := <-metadataResults:

// This can overwrite unsent values, but that's fine because

// we're only interested about up-to-date values.

updatedMetadata[mr.key] = mr.result

continue

case err := <-reportError:

logMsg := "batch update metadata complete"

if err != nil {

a.logger.Debug(ctx, logMsg, slog.Error(err))

return xerrors.Errorf("failed to report metadata: %w", err)

}

a.logger.Debug(ctx, logMsg)

reportInFlight = false

case <-report:

if len(updatedMetadata) == 0 {

continue

}

if reportInFlight {

// If there's already a report in flight, don't send

// another one, wait for next tick instead.

a.logger.Debug(ctx, "skipped metadata report tick because report is in flight")

continue

}

metadata := make([]*proto.Metadata, 0, len(updatedMetadata))

for key, result := range updatedMetadata {

pr := agentsdk.ProtoFromMetadataResult(*result)

metadata = append(metadata, &proto.Metadata{

Key: key,

Result: pr,

})

delete(updatedMetadata, key)

}

reportInFlight = true

go func() {

a.logger.Debug(ctx, "batch updating metadata")

ctx, cancel := context.WithTimeout(ctx, reportTimeout)

defer cancel()

_, err := aAPI.BatchUpdateMetadata(ctx, &proto.BatchUpdateMetadataRequest{Metadata: metadata})

reportError <- err

}()

}

}

}

// reportLifecycle reports the current lifecycle state once. All state

// changes are reported in order.

func (a *agent) reportLifecycle(ctx context.Context, aAPI proto.DRPCAgentClient28) error {

for {

select {

case <-a.lifecycleUpdate:

case <-ctx.Done():

return ctx.Err()

}

for {

a.lifecycleMu.RLock()

lastIndex := len(a.lifecycleStates) - 1

report := a.lifecycleStates[a.lifecycleLastReportedIndex]

if len(a.lifecycleStates) > a.lifecycleLastReportedIndex+1 {

report = a.lifecycleStates[a.lifecycleLastReportedIndex+1]

}

a.lifecycleMu.RUnlock()

if lastIndex == a.lifecycleLastReportedIndex {

break

}

l, err := agentsdk.ProtoFromLifecycle(report)

if err != nil {

a.logger.Critical(ctx, "failed to convert lifecycle state", slog.F("report", report))

// Skip this report; there is no point retrying. Maybe we can successfully convert the next one?

a.lifecycleLastReportedIndex++

continue

}

payload := &proto.UpdateLifecycleRequest{Lifecycle: l}

logger := a.logger.With(slog.F("payload", payload))

logger.Debug(ctx, "reporting lifecycle state")

_, err = aAPI.UpdateLifecycle(ctx, payload)

if err != nil {

return xerrors.Errorf("failed to update lifecycle: %w", err)

}

logger.Debug(ctx, "successfully reported lifecycle state")

a.lifecycleLastReportedIndex++

select {

case a.lifecycleReported <- report.State:

case <-a.lifecycleReported:

a.lifecycleReported <- report.State

}

if a.lifecycleLastReportedIndex < lastIndex {

// Keep reporting until we've sent all messages, we can't

// rely on the channel triggering us before the backlog is

// consumed.

continue

}

break

}

}

}

// setLifecycle sets the lifecycle state and notifies the lifecycle loop.

// The state is only updated if it's a valid state transition.

func (a *agent) setLifecycle(state codersdk.WorkspaceAgentLifecycle) {

report := agentsdk.PostLifecycleRequest{

State: state,

ChangedAt: dbtime.Now(),

}

a.lifecycleMu.Lock()

lastReport := a.lifecycleStates[len(a.lifecycleStates)-1]

if slices.Index(codersdk.WorkspaceAgentLifecycleOrder, lastReport.State) >= slices.Index(codersdk.WorkspaceAgentLifecycleOrder, report.State) {

a.logger.Warn(context.Background(), "attempted to set lifecycle state to a previous state", slog.F("last", lastReport), slog.F("current", report))

a.lifecycleMu.Unlock()

return

}

a.lifecycleStates = append(a.lifecycleStates, report)

a.logger.Debug(context.Background(), "set lifecycle state", slog.F("current", report), slog.F("last", lastReport))

a.lifecycleMu.Unlock()

select {

case a.lifecycleUpdate <- struct{}{}:

default:

}

}

// fetchServiceBannerLoop fetches the service banner on an interval. It will

// not be fetched immediately; the expectation is that it is primed elsewhere

// (and must be done before the session actually starts).

func (a *agent) fetchServiceBannerLoop(ctx context.Context, aAPI proto.DRPCAgentClient28) error {

ticker := time.NewTicker(a.announcementBannersRefreshInterval)

defer ticker.Stop()

for {

select {

case <-ctx.Done():

return ctx.Err()

case <-ticker.C:

bannersProto, err := aAPI.GetAnnouncementBanners(ctx, &proto.GetAnnouncementBannersRequest{})

if err != nil {

if ctx.Err() != nil {

return ctx.Err()

}

a.logger.Error(ctx, "failed to update notification banners", slog.Error(err))

return err

}

banners := make([]codersdk.BannerConfig, 0, len(bannersProto.AnnouncementBanners))

for _, bannerProto := range bannersProto.AnnouncementBanners {