package agentapi

import (

"context"

"database/sql"

"slices"

"sync"

"time"

"github.com/google/uuid"

"golang.org/x/mod/semver"

"golang.org/x/xerrors"

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

"cdr.dev/slog/v3"

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

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

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

"github.com/coder/coder/v2/coderd/wspubsub"

)

type contextKeyAPIVersion struct{}

func WithAPIVersion(ctx context.Context, version string) context.Context {

return context.WithValue(ctx, contextKeyAPIVersion{}, version)

}

type LifecycleAPI struct {

AgentFn func(context.Context) (database.WorkspaceAgent, error)

WorkspaceID uuid.UUID

Database database.Store

Log slog.Logger

PublishWorkspaceUpdateFn func(context.Context, uuid.UUID, wspubsub.WorkspaceEventKind) error

TimeNowFn func() time.Time // defaults to dbtime.Now()

Metrics *LifecycleMetrics

emitMetricsOnce sync.Once

}

func (a *LifecycleAPI) now() time.Time {

if a.TimeNowFn != nil {

return a.TimeNowFn()

}

return dbtime.Now()

}

func (a *LifecycleAPI) UpdateLifecycle(ctx context.Context, req *agentproto.UpdateLifecycleRequest) (*agentproto.Lifecycle, error) {

workspaceAgent, err := a.AgentFn(ctx)

if err != nil {

return nil, err

}

logger := a.Log.With(

slog.F("workspace_id", a.WorkspaceID),

slog.F("payload", req),

)

logger.Debug(ctx, "workspace agent state report")

var lifecycleState database.WorkspaceAgentLifecycleState

switch req.Lifecycle.State {

case agentproto.Lifecycle_CREATED:

lifecycleState = database.WorkspaceAgentLifecycleStateCreated

case agentproto.Lifecycle_STARTING:

lifecycleState = database.WorkspaceAgentLifecycleStateStarting

case agentproto.Lifecycle_START_TIMEOUT:

lifecycleState = database.WorkspaceAgentLifecycleStateStartTimeout

case agentproto.Lifecycle_START_ERROR:

lifecycleState = database.WorkspaceAgentLifecycleStateStartError

case agentproto.Lifecycle_READY:

lifecycleState = database.WorkspaceAgentLifecycleStateReady

case agentproto.Lifecycle_SHUTTING_DOWN:

lifecycleState = database.WorkspaceAgentLifecycleStateShuttingDown

case agentproto.Lifecycle_SHUTDOWN_TIMEOUT:

lifecycleState = database.WorkspaceAgentLifecycleStateShutdownTimeout

case agentproto.Lifecycle_SHUTDOWN_ERROR:

lifecycleState = database.WorkspaceAgentLifecycleStateShutdownError

case agentproto.Lifecycle_OFF:

lifecycleState = database.WorkspaceAgentLifecycleStateOff

default:

return nil, xerrors.Errorf("unknown lifecycle state %q", req.Lifecycle.State)

}

if !lifecycleState.Valid() {

return nil, xerrors.Errorf("unknown lifecycle state %q", req.Lifecycle.State)

}

changedAt := req.Lifecycle.ChangedAt.AsTime()

if changedAt.IsZero() {

changedAt = a.now()

req.Lifecycle.ChangedAt = timestamppb.New(changedAt)

}

dbChangedAt := sql.NullTime{Time: changedAt, Valid: true}

startedAt := workspaceAgent.StartedAt

readyAt := workspaceAgent.ReadyAt

switch lifecycleState {

case database.WorkspaceAgentLifecycleStateStarting:

startedAt = dbChangedAt

// This agent is (re)starting, so it's not ready yet.

readyAt.Time = time.Time{}

readyAt.Valid = false

case database.WorkspaceAgentLifecycleStateReady,

database.WorkspaceAgentLifecycleStateStartTimeout,

database.WorkspaceAgentLifecycleStateStartError:

if !startedAt.Valid {

startedAt = dbChangedAt

}

readyAt = dbChangedAt

}

err = a.Database.UpdateWorkspaceAgentLifecycleStateByID(ctx, database.UpdateWorkspaceAgentLifecycleStateByIDParams{

ID: workspaceAgent.ID,

LifecycleState: lifecycleState,

StartedAt: startedAt,

ReadyAt: readyAt,

})

if err != nil {

if !database.IsQueryCanceledError(err) {

// not an error if we are canceled

logger.Error(ctx, "failed to update lifecycle state", slog.Error(err))

}

return nil, xerrors.Errorf("update workspace agent lifecycle state: %w", err)

}

if a.PublishWorkspaceUpdateFn != nil {

err = a.PublishWorkspaceUpdateFn(ctx, workspaceAgent.ID, wspubsub.WorkspaceEventKindAgentLifecycleUpdate)

if err != nil {

return nil, xerrors.Errorf("publish workspace update: %w", err)

}

}

// Emit build duration metric when agent transitions to a terminal startup state.

// We only emit once per agent connection to avoid duplicate metrics.

switch lifecycleState {

case database.WorkspaceAgentLifecycleStateReady,

database.WorkspaceAgentLifecycleStateStartTimeout,

database.WorkspaceAgentLifecycleStateStartError:

// Only emit metrics for the parent agent, this metric is not intended to measure devcontainer durations.

if !workspaceAgent.ParentID.Valid {

a.emitMetricsOnce.Do(func() {

a.emitBuildDurationMetric(ctx, workspaceAgent.ResourceID)

})

}

}

return req.Lifecycle, nil

}

func (a *LifecycleAPI) UpdateStartup(ctx context.Context, req *agentproto.UpdateStartupRequest) (*agentproto.Startup, error) {

apiVersion, ok := ctx.Value(contextKeyAPIVersion{}).(string)

if !ok {

return nil, xerrors.Errorf("internal error; api version unspecified")

}

workspaceAgent, err := a.AgentFn(ctx)

if err != nil {

return nil, err

}

a.Log.Debug(

ctx,

"post workspace agent version",

slog.F("workspace_id", a.WorkspaceID),

slog.F("agent_version", req.Startup.Version),

)

if !semver.IsValid(req.Startup.Version) {

return nil, xerrors.Errorf("invalid agent semver version %q", req.Startup.Version)

}

// Validate subsystems.

dbSubsystems := make([]database.WorkspaceAgentSubsystem, 0, len(req.Startup.Subsystems))

seenSubsystems := make(map[database.WorkspaceAgentSubsystem]struct{}, len(req.Startup.Subsystems))

for _, s := range req.Startup.Subsystems {

var dbSubsystem database.WorkspaceAgentSubsystem

switch s {

case agentproto.Startup_ENVBOX:

dbSubsystem = database.WorkspaceAgentSubsystemEnvbox

case agentproto.Startup_ENVBUILDER:

dbSubsystem = database.WorkspaceAgentSubsystemEnvbuilder

case agentproto.Startup_EXECTRACE:

dbSubsystem = database.WorkspaceAgentSubsystemExectrace

default:

return nil, xerrors.Errorf("invalid agent subsystem %q", s)

}

if _, ok := seenSubsystems[dbSubsystem]; !ok {

seenSubsystems[dbSubsystem] = struct{}{}

dbSubsystems = append(dbSubsystems, dbSubsystem)

}

}

slices.Sort(dbSubsystems)

err = a.Database.UpdateWorkspaceAgentStartupByID(ctx, database.UpdateWorkspaceAgentStartupByIDParams{

ID: workspaceAgent.ID,

Version: req.Startup.Version,

ExpandedDirectory: req.Startup.ExpandedDirectory,

Subsystems: dbSubsystems,

APIVersion: apiVersion,

})

if err != nil {

return nil, xerrors.Errorf("update workspace agent startup in database: %w", err)

}

return req.Startup, nil

}