package agentsdk

import (

"context"

"encoding/json"

"errors"

"fmt"

"io"

"net/http"

"net/url"

"sync"

"time"

"github.com/google/uuid"

"github.com/hashicorp/yamux"

"golang.org/x/xerrors"

"storj.io/drpc"

"tailscale.com/tailcfg"

"cdr.dev/slog/v3"

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

"github.com/coder/coder/v2/apiversion"

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

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

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

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

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

"github.com/coder/retry"

"github.com/coder/websocket"

)

// ExternalLogSourceID is the statically-defined ID of a log-source that

// appears as "External" in the dashboard.

//

// This is to support legacy API-consumers that do not create their own

// log-source. This should be removed in the future.

var ExternalLogSourceID = uuid.MustParse("3b579bf4-1ed8-4b99-87a8-e9a1e3410410")

// SessionTokenSetup is a function that creates the token provider while setting up the workspace agent. We do it this

// way because cloud instance identity (AWS, Azure, Google, etc.) requires interacting with coderd to exchange tokens.

// This means that the token providers need a codersdk.Client. However, the SessionTokenProvider is itself used by

// the client to authenticate requests. Thus, the dependency is bidirectional. Functions of this type are used in

// New() to ensure that things are set up correctly so there is only one instance of the codersdk.Client created.

// @typescript-ignore SessionTokenSetup

type SessionTokenSetup func(client *codersdk.Client) RefreshableSessionTokenProvider

// New creates a new *Client which can be used by an agent to connect to Coderd. Use a SessionTokenSetup function

// to define the session token provider for the Client. This overrides the SessionTokenProvider on the underlying

// `*codersdk.Client`, so any `codersdk.ClientOptions` passed as `opts` should not set this property.

func New(serverURL *url.URL, setup SessionTokenSetup, opts ...codersdk.ClientOption) *Client {

var provider RefreshableSessionTokenProvider

opts = append(opts, func(c *codersdk.Client) {

provider = setup(c)

c.SessionTokenProvider = provider

})

c := codersdk.New(serverURL, opts...)

return &Client{

SDK: c,

RefreshableSessionTokenProvider: provider,

}

}

// Client wraps `codersdk.Client` with specific functions

// scoped to a workspace agent.

type Client struct {

RefreshableSessionTokenProvider

SDK *codersdk.Client

}

type GitSSHKey struct {

PublicKey string `json:"public_key"`

PrivateKey string `json:"private_key"`

}

// GitSSHKey will return the user's SSH key pair for the workspace.

func (c *Client) GitSSHKey(ctx context.Context) (GitSSHKey, error) {

res, err := c.SDK.Request(ctx, http.MethodGet, "/api/v2/workspaceagents/me/gitsshkey", nil)

if err != nil {

return GitSSHKey{}, xerrors.Errorf("execute request: %w", err)

}

defer res.Body.Close()

if res.StatusCode != http.StatusOK {

return GitSSHKey{}, codersdk.ReadBodyAsError(res)

}

var gitSSHKey GitSSHKey

return gitSSHKey, json.NewDecoder(res.Body).Decode(&gitSSHKey)

}

type Metadata struct {

Key string `json:"key"`

codersdk.WorkspaceAgentMetadataResult

}

type PostMetadataRequest struct {

Metadata []Metadata `json:"metadata"`

}

// In the future, we may want to support sending back multiple values for

// performance.

type PostMetadataRequestDeprecated = codersdk.WorkspaceAgentMetadataResult

// Manifest is the workspace agent's view of its own configuration.

//

// Secrets are intentionally not a field on this struct. The manifest

// may be serialized (JSON, %+v, logger fields, debug endpoints) in

// many places that do not and should not carry secret values.

// Keeping Secrets off of the struct makes leaking them impossible

// via any code path that only holds a *Manifest. Callers that need

// secrets must load them explicitly via SecretsFromProto on the raw

// proto.

type Manifest struct {

ParentID uuid.UUID `json:"parent_id"`

AgentID uuid.UUID `json:"agent_id"`

AgentName string `json:"agent_name"`

// OwnerUsername and WorkspaceID are used by an open-source user to identify the workspace.

// We do not provide insurance that this will not be removed in the future,

// but if it's easy to persist lets keep it around.

OwnerName string `json:"owner_name"`

WorkspaceID uuid.UUID `json:"workspace_id"`

WorkspaceName string `json:"workspace_name"`

// GitAuthConfigs stores the number of Git configurations

// the Coder deployment has. If this number is >0, we

// set up special configuration in the workspace.

GitAuthConfigs int `json:"git_auth_configs"`

VSCodePortProxyURI string `json:"vscode_port_proxy_uri"`

Apps []codersdk.WorkspaceApp `json:"apps"`

DERPMap *tailcfg.DERPMap `json:"derpmap"`

DERPForceWebSockets bool `json:"derp_force_websockets"`

EnvironmentVariables map[string]string `json:"environment_variables"`

Directory string `json:"directory"`

MOTDFile string `json:"motd_file"`

DisableDirectConnections bool `json:"disable_direct_connections"`

Metadata []codersdk.WorkspaceAgentMetadataDescription `json:"metadata"`

Scripts []codersdk.WorkspaceAgentScript `json:"scripts"`

Devcontainers []codersdk.WorkspaceAgentDevcontainer `json:"devcontainers"`

}

// WorkspaceSecret is a user secret for injection into a workspace.

//

// Value carries decrypted secret material and is omitted from JSON

// serialization to protect against future leaking of the secret.

type WorkspaceSecret struct {

EnvName string

FilePath string

Value []byte `json:"-"`

}

type LogSource struct {

ID uuid.UUID `json:"id"`

DisplayName string `json:"display_name"`

Icon string `json:"icon"`

}

type Script struct {

Script string `json:"script"`

}

// RewriteDERPMap rewrites the DERP map to use the configured access URL of the

// agent as the "embedded relay" access URL.

//

// See tailnet.RewriteDERPMapDefaultRelay for more details on why this is

// necessary.

func (c *Client) RewriteDERPMap(derpMap *tailcfg.DERPMap) {

tailnet.RewriteDERPMapDefaultRelay(context.Background(), c.SDK.Logger(), derpMap, c.SDK.URL)

}

// ConnectRPC20 returns a dRPC client to the Agent API v2.0. Notably, it is missing

// GetAnnouncementBanners, but is useful when you want to be maximally compatible with Coderd

// Release Versions from 2.9+

// Deprecated: use ConnectRPC20WithTailnet

func (c *Client) ConnectRPC20(ctx context.Context) (proto.DRPCAgentClient20, error) {

conn, err := c.connectRPCVersion(ctx, apiversion.New(2, 0), "")

if err != nil {

return nil, err

}

return proto.NewDRPCAgentClient(conn), nil

}

// ConnectRPC20WithTailnet returns a dRPC client to the Agent API v2.0. Notably, it is missing

// GetAnnouncementBanners, but is useful when you want to be maximally compatible with Coderd

// Release Versions from 2.9+

func (c *Client) ConnectRPC20WithTailnet(ctx context.Context) (

proto.DRPCAgentClient20, tailnetproto.DRPCTailnetClient20, error,

) {

conn, err := c.connectRPCVersion(ctx, apiversion.New(2, 0), "")

if err != nil {

return nil, nil, err

}

return proto.NewDRPCAgentClient(conn), tailnetproto.NewDRPCTailnetClient(conn), nil

}

// ConnectRPC21 returns a dRPC client to the Agent API v2.1. It is useful when you want to be

// maximally compatible with Coderd Release Versions from 2.12+

// Deprecated: use ConnectRPC21WithTailnet

func (c *Client) ConnectRPC21(ctx context.Context) (proto.DRPCAgentClient21, error) {

conn, err := c.connectRPCVersion(ctx, apiversion.New(2, 1), "")

if err != nil {

return nil, err

}

return proto.NewDRPCAgentClient(conn), nil

}

// ConnectRPC21WithTailnet returns a dRPC client to the Agent API v2.1. It is useful when you want to be

// maximally compatible with Coderd Release Versions from 2.12+

func (c *Client) ConnectRPC21WithTailnet(ctx context.Context) (

proto.DRPCAgentClient21, tailnetproto.DRPCTailnetClient21, error,

) {

conn, err := c.connectRPCVersion(ctx, apiversion.New(2, 1), "")

if err != nil {

return nil, nil, err

}

return proto.NewDRPCAgentClient(conn), tailnetproto.NewDRPCTailnetClient(conn), nil

}

// ConnectRPC22 returns a dRPC client to the Agent API v2.2. It is useful when you want to be

// maximally compatible with Coderd Release Versions from 2.13+

func (c *Client) ConnectRPC22(ctx context.Context) (

proto.DRPCAgentClient22, tailnetproto.DRPCTailnetClient22, error,

) {

conn, err := c.connectRPCVersion(ctx, apiversion.New(2, 2), "")

if err != nil {

return nil, nil, err

}

return proto.NewDRPCAgentClient(conn), tailnetproto.NewDRPCTailnetClient(conn), nil

}

// ConnectRPC23 returns a dRPC client to the Agent API v2.3. It is useful when you want to be

// maximally compatible with Coderd Release Versions from 2.18+

func (c *Client) ConnectRPC23(ctx context.Context) (

proto.DRPCAgentClient23, tailnetproto.DRPCTailnetClient23, error,

) {

conn, err := c.connectRPCVersion(ctx, apiversion.New(2, 3), "")

if err != nil {

return nil, nil, err

}

return proto.NewDRPCAgentClient(conn), tailnetproto.NewDRPCTailnetClient(conn), nil

}

// ConnectRPC24 returns a dRPC client to the Agent API v2.4. It is useful when you want to be

// maximally compatible with Coderd Release Versions from 2.20+

func (c *Client) ConnectRPC24(ctx context.Context) (

proto.DRPCAgentClient24, tailnetproto.DRPCTailnetClient24, error,

) {

conn, err := c.connectRPCVersion(ctx, apiversion.New(2, 4), "")

if err != nil {

return nil, nil, err

}

return proto.NewDRPCAgentClient(conn), tailnetproto.NewDRPCTailnetClient(conn), nil

}

// ConnectRPC25 returns a dRPC client to the Agent API v2.5. It is useful when you want to be

// maximally compatible with Coderd Release Versions from 2.23+

func (c *Client) ConnectRPC25(ctx context.Context) (

proto.DRPCAgentClient25, tailnetproto.DRPCTailnetClient25, error,

) {

conn, err := c.connectRPCVersion(ctx, apiversion.New(2, 5), "")

if err != nil {

return nil, nil, err

}

return proto.NewDRPCAgentClient(conn), tailnetproto.NewDRPCTailnetClient(conn), nil

}

// ConnectRPC26 returns a dRPC client to the Agent API v2.6. It is useful when you want to be

// maximally compatible with Coderd Release Versions from 2.24+

func (c *Client) ConnectRPC26(ctx context.Context) (

proto.DRPCAgentClient26, tailnetproto.DRPCTailnetClient26, error,

) {

conn, err := c.connectRPCVersion(ctx, apiversion.New(2, 6), "")

if err != nil {

return nil, nil, err

}

return proto.NewDRPCAgentClient(conn), tailnetproto.NewDRPCTailnetClient(conn), nil

}

// ConnectRPC27 returns a dRPC client to the Agent API v2.7. It is useful when you want to be

// maximally compatible with Coderd Release Versions from 2.30+

func (c *Client) ConnectRPC27(ctx context.Context) (

proto.DRPCAgentClient27, tailnetproto.DRPCTailnetClient27, error,

) {

conn, err := c.connectRPCVersion(ctx, apiversion.New(2, 7), "")

if err != nil {

return nil, nil, err

}

return proto.NewDRPCAgentClient(conn), tailnetproto.NewDRPCTailnetClient(conn), nil

}

// ConnectRPC28 returns a dRPC client to the Agent API v2.8. It is useful when you want to be

// maximally compatible with Coderd Release Versions from 2.31+

func (c *Client) ConnectRPC28(ctx context.Context) (

proto.DRPCAgentClient28, tailnetproto.DRPCTailnetClient28, error,

) {

conn, err := c.connectRPCVersion(ctx, apiversion.New(2, 8), "")

if err != nil {

return nil, nil, err

}

return proto.NewDRPCAgentClient(conn), tailnetproto.NewDRPCTailnetClient(conn), nil

}

// ConnectRPC28WithRole is like ConnectRPC28 but sends an explicit role

// query parameter to the server. Use "agent" for workspace agents to

// enable connection monitoring.

func (c *Client) ConnectRPC28WithRole(ctx context.Context, role string) (

proto.DRPCAgentClient28, tailnetproto.DRPCTailnetClient28, error,

) {

conn, err := c.connectRPCVersion(ctx, apiversion.New(2, 8), role)

if err != nil {

return nil, nil, err

}

return proto.NewDRPCAgentClient(conn), tailnetproto.NewDRPCTailnetClient(conn), nil

}

// ConnectRPC connects to the workspace agent API and tailnet API.

// It does not send a role query parameter, so the server will apply

// its default behavior (currently: enable connection monitoring for

// backward compatibility). Use ConnectRPCWithRole to explicitly

// identify the caller's role.

func (c *Client) ConnectRPC(ctx context.Context) (drpc.Conn, error) {

return c.connectRPCVersion(ctx, proto.CurrentVersion, "")

}

// ConnectRPCWithRole connects to the workspace agent RPC API with an

// explicit role. The role parameter is sent to the server to identify

// the type of client. Use "agent" for workspace agents to enable

// connection monitoring.

func (c *Client) ConnectRPCWithRole(ctx context.Context, role string) (drpc.Conn, error) {

return c.connectRPCVersion(ctx, proto.CurrentVersion, role)

}

func (c *Client) connectRPCVersion(ctx context.Context, version *apiversion.APIVersion, role string) (drpc.Conn, error) {

rpcURL, err := c.SDK.URL.Parse("/api/v2/workspaceagents/me/rpc")

if err != nil {

return nil, xerrors.Errorf("parse url: %w", err)

}

q := rpcURL.Query()

q.Add("version", version.String())

if role != "" {

q.Add("role", role)

}

rpcURL.RawQuery = q.Encode()

httpClient := &http.Client{

Transport: c.SDK.HTTPClient.Transport,

}

// nolint:bodyclose

conn, res, err := websocket.Dial(ctx, rpcURL.String(), &websocket.DialOptions{

HTTPClient: httpClient,

HTTPHeader: http.Header{

codersdk.SessionTokenHeader: []string{c.SDK.SessionToken()},

},

})

if err != nil {

if res == nil {

return nil, err

}

return nil, codersdk.ReadBodyAsError(res)

}

// Set the read limit to 4 MiB -- about the limit for protobufs. This needs to be larger than

// the default because some of our protocols can include large messages like startup scripts.

conn.SetReadLimit(1 << 22)

netConn := websocket.NetConn(ctx, conn, websocket.MessageBinary)

config := yamux.DefaultConfig()

config.LogOutput = nil

config.Logger = slog.Stdlib(ctx, c.SDK.Logger(), slog.LevelInfo)

session, err := yamux.Client(netConn, config)

if err != nil {

return nil, xerrors.Errorf("multiplex client: %w", err)

}

return drpcsdk.MultiplexedConn(session), nil

}

type PostAppHealthsRequest struct {

// Healths is a map of the workspace app name and the health of the app.

Healths map[uuid.UUID]codersdk.WorkspaceAppHealth

}

// BatchUpdateAppHealthsClient is a partial interface of proto.DRPCAgentClient.

type BatchUpdateAppHealthsClient interface {

BatchUpdateAppHealths(ctx context.Context, req *proto.BatchUpdateAppHealthRequest) (*proto.BatchUpdateAppHealthResponse, error)

}

func AppHealthPoster(aAPI BatchUpdateAppHealthsClient) func(ctx context.Context, req PostAppHealthsRequest) error {

return func(ctx context.Context, req PostAppHealthsRequest) error {

pReq, err := ProtoFromAppHealthsRequest(req)

if err != nil {

return xerrors.Errorf("convert AppHealthsRequest: %w", err)

}

_, err = aAPI.BatchUpdateAppHealths(ctx, pReq)

if err != nil {

return xerrors.Errorf("batch update app healths: %w", err)

}

return nil

}

}

// AuthenticateResponse is returned when an instance ID

// has been exchanged for a session token.

// @typescript-ignore AuthenticateResponse

type AuthenticateResponse struct {

SessionToken string `json:"session_token"`

}

// RefreshableSessionTokenProvider is a SessionTokenProvider that can be refreshed, for example, via token exchange.

// @typescript-ignore RefreshableSessionTokenProvider

type RefreshableSessionTokenProvider interface {

codersdk.SessionTokenProvider

RefreshToken(ctx context.Context) error

}

// InstanceIdentitySessionTokenProvider implements RefreshableSessionTokenProvider via token exchange for a cloud

// compute instance identity.

// @typescript-ignore InstanceIdentitySessionTokenProvider

type InstanceIdentitySessionTokenProvider struct {

TokenExchanger TokenExchanger

logger slog.Logger

// cache so we don't request each time

mu sync.Mutex

sessionToken string

}

// TokenExchanger obtains a session token by exchanging a cloud instance identity credential for a Coder session token.

// @typescript-ignore TokenExchanger

type TokenExchanger interface {

exchange(ctx context.Context) (AuthenticateResponse, error)

}

func (i *InstanceIdentitySessionTokenProvider) AsRequestOption() codersdk.RequestOption {

t := i.GetSessionToken()

return func(req *http.Request) {

req.Header.Set(codersdk.SessionTokenHeader, t)

}

}

func (i *InstanceIdentitySessionTokenProvider) SetDialOption(opts *websocket.DialOptions) {

t := i.GetSessionToken()

if opts.HTTPHeader == nil {

opts.HTTPHeader = http.Header{}

}

if opts.HTTPHeader.Get(codersdk.SessionTokenHeader) == "" {

opts.HTTPHeader.Set(codersdk.SessionTokenHeader, t)

}

}

func (i *InstanceIdentitySessionTokenProvider) GetSessionToken() string {

i.mu.Lock()

defer i.mu.Unlock()

if i.sessionToken != "" {

return i.sessionToken

}

ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)

defer cancel()

resp, err := i.TokenExchanger.exchange(ctx)

if err != nil {

i.logger.Error(ctx, "failed to exchange session token", slog.Error(err))

return ""

}

i.sessionToken = resp.SessionToken

return i.sessionToken

}

func (i *InstanceIdentitySessionTokenProvider) RefreshToken(ctx context.Context) error {

i.mu.Lock()

defer i.mu.Unlock()

resp, err := i.TokenExchanger.exchange(ctx)

if err != nil {

return err

}

i.sessionToken = resp.SessionToken

return nil

}

// FixedSessionTokenProvider wraps the codersdk variant to add a no-op RefreshToken method to satisfy the

// RefreshableSessionTokenProvider interface.

// @typescript-ignore FixedSessionTokenProvider

type FixedSessionTokenProvider struct {

codersdk.FixedSessionTokenProvider

}

func (FixedSessionTokenProvider) RefreshToken(_ context.Context) error {

return nil

}

// InstanceIdentityConfig holds optional configuration for cloud

// instance-identity authentication.

type InstanceIdentityConfig struct {

AgentName string

}

// InstanceIdentityOption configures instance-identity authentication.

type InstanceIdentityOption func(*InstanceIdentityConfig)

// WithInstanceIdentityAgentName sets the agent name selector sent with

// the instance-identity authentication request.

func WithInstanceIdentityAgentName(name string) InstanceIdentityOption {

return func(c *InstanceIdentityConfig) {

c.AgentName = name

}

}

// applyInstanceIdentityOptions applies the given options and returns

// the resulting configuration.

func applyInstanceIdentityOptions(opts []InstanceIdentityOption) InstanceIdentityConfig {

var cfg InstanceIdentityConfig

for _, o := range opts {

o(&cfg)

}

return cfg

}

func WithFixedToken(token string) SessionTokenSetup {

return func(_ *codersdk.Client) RefreshableSessionTokenProvider {

return FixedSessionTokenProvider{FixedSessionTokenProvider: codersdk.FixedSessionTokenProvider{SessionToken: token}}

}

}

// Stats records the Agent's network connection statistics for use in

// user-facing metrics and debugging.

type Stats struct {

// ConnectionsByProto is a count of connections by protocol.

ConnectionsByProto map[string]int64 `json:"connections_by_proto"`

// ConnectionCount is the number of connections received by an agent.

ConnectionCount int64 `json:"connection_count"`

// ConnectionMedianLatencyMS is the median latency of all connections in milliseconds.

ConnectionMedianLatencyMS float64 `json:"connection_median_latency_ms"`

// RxPackets is the number of received packets.

RxPackets int64 `json:"rx_packets"`

// RxBytes is the number of received bytes.

RxBytes int64 `json:"rx_bytes"`

// TxPackets is the number of transmitted bytes.

TxPackets int64 `json:"tx_packets"`

// TxBytes is the number of transmitted bytes.

TxBytes int64 `json:"tx_bytes"`

// SessionCountVSCode is the number of connections received by an agent

// that are from our VS Code extension.

SessionCountVSCode int64 `json:"session_count_vscode"`

// SessionCountJetBrains is the number of connections received by an agent

// that are from our JetBrains extension.

SessionCountJetBrains int64 `json:"session_count_jetbrains"`

// SessionCountReconnectingPTY is the number of connections received by an agent

// that are from the reconnecting web terminal.

SessionCountReconnectingPTY int64 `json:"session_count_reconnecting_pty"`

// SessionCountSSH is the number of connections received by an agent

// that are normal, non-tagged SSH sessions.

SessionCountSSH int64 `json:"session_count_ssh"`

// Metrics collected by the agent

Metrics []AgentMetric `json:"metrics"`

}

func (s Stats) SessionCount() int64 {

return s.SessionCountVSCode + s.SessionCountJetBrains + s.SessionCountReconnectingPTY + s.SessionCountSSH

}

type AgentMetricType string

const (

AgentMetricTypeCounter AgentMetricType = "counter"

AgentMetricTypeGauge AgentMetricType = "gauge"

)

type AgentMetric struct {

Name string `json:"name" validate:"required"`

Type AgentMetricType `json:"type" validate:"required" enums:"counter,gauge"`

Value float64 `json:"value" validate:"required"`

Labels []AgentMetricLabel `json:"labels,omitempty"`

}

type AgentMetricLabel struct {

Name string `json:"name" validate:"required"`

Value string `json:"value" validate:"required"`

}

type StatsResponse struct {

// ReportInterval is the duration after which the agent should send stats

// again.

ReportInterval time.Duration `json:"report_interval"`

}

type PostLifecycleRequest struct {

State codersdk.WorkspaceAgentLifecycle `json:"state"`

ChangedAt time.Time `json:"changed_at"`

}

type PostStartupRequest struct {

Version string `json:"version"`

ExpandedDirectory string `json:"expanded_directory"`

Subsystems []codersdk.AgentSubsystem `json:"subsystems"`

}

type Log struct {

CreatedAt time.Time `json:"created_at"`

Output string `json:"output"`

Level codersdk.LogLevel `json:"level"`

}

type PatchLogs struct {

LogSourceID uuid.UUID `json:"log_source_id"`

Logs []Log `json:"logs"`

}

// PatchLogs writes log messages to the agent startup script.

// Log messages are limited to 1MB in total.

//

// Deprecated: use the DRPCAgentClient.BatchCreateLogs instead

func (c *Client) PatchLogs(ctx context.Context, req PatchLogs) error {

res, err := c.SDK.Request(ctx, http.MethodPatch, "/api/v2/workspaceagents/me/logs", req)

if err != nil {

return err

}

defer res.Body.Close()

if res.StatusCode != http.StatusOK {

return codersdk.ReadBodyAsError(res)

}

return nil

}

// PatchAppStatus updates the status of a workspace app.

type PatchAppStatus struct {

AppSlug string `json:"app_slug"`

State codersdk.WorkspaceAppStatusState `json:"state"`

Message string `json:"message"`

URI string `json:"uri"`

// Deprecated: this field is unused and will be removed in a future version.

Icon string `json:"icon"`

// Deprecated: this field is unused and will be removed in a future version.

NeedsUserAttention bool `json:"needs_user_attention"`

}

// PatchAppStatus updates the status of a workspace app.

// Deprecated: use the DRPCAgentClient.UpdateAppStatus instead

func (c *Client) PatchAppStatus(ctx context.Context, req PatchAppStatus) error {

res, err := c.SDK.Request(ctx, http.MethodPatch, "/api/v2/workspaceagents/me/app-status", req)

if err != nil {

return err

}

defer res.Body.Close()

if res.StatusCode != http.StatusOK {

return codersdk.ReadBodyAsError(res)

}

return nil

}

type PostLogSourceRequest struct {

// ID is a unique identifier for the log source.

// It is scoped to a workspace agent, and can be statically

// defined inside code to prevent duplicate sources from being

// created for the same agent.

ID uuid.UUID `json:"id"`

DisplayName string `json:"display_name"`

Icon string `json:"icon"`

}

func (c *Client) PostLogSource(ctx context.Context, req PostLogSourceRequest) (codersdk.WorkspaceAgentLogSource, error) {

res, err := c.SDK.Request(ctx, http.MethodPost, "/api/v2/workspaceagents/me/log-source", req)

if err != nil {

return codersdk.WorkspaceAgentLogSource{}, err

}

defer res.Body.Close()

if res.StatusCode != http.StatusCreated {

return codersdk.WorkspaceAgentLogSource{}, codersdk.ReadBodyAsError(res)

}

var logSource codersdk.WorkspaceAgentLogSource

return logSource, json.NewDecoder(res.Body).Decode(&logSource)

}

type ExternalAuthResponse struct {

AccessToken string `json:"access_token"`

TokenExtra map[string]interface{} `json:"token_extra"`

URL string `json:"url"`

Type string `json:"type"`

// Deprecated: Only supported on `/workspaceagents/me/gitauth`

// for backwards compatibility.

Username string `json:"username"`

Password string `json:"password"`

}

// ExternalAuthRequest is used to request an access token for a provider.

// Either ID or Match must be specified, but not both.

type ExternalAuthRequest struct {

// ID is the ID of a provider to request authentication for.

ID string

// Match is an arbitrary string matched against the regex of the provider.

Match string

// GitBranch is the current git branch in the working directory.

// Sent by the agent so the control plane can resolve diffs

// without SSHing into the workspace.

GitBranch string

// GitRemoteOrigin is the remote origin URL of the git repository.

// Sent by the agent so the control plane can resolve diffs

// without SSHing into the workspace.

GitRemoteOrigin string

// ChatID identifies which chat initiated the git operation.

ChatID string

// Listen indicates that the request should be long-lived and listen for

// a new token to be requested.

Listen bool

}

// ExternalAuth submits a URL or provider ID to fetch an access token for.

// nolint:revive

func (c *Client) ExternalAuth(ctx context.Context, req ExternalAuthRequest) (ExternalAuthResponse, error) {

q := url.Values{

"id": []string{req.ID},

"match": []string{req.Match},

}

if req.Listen {

q.Set("listen", "true")

}

if req.GitBranch != "" {

q.Set("git_branch", req.GitBranch)

}

if req.GitRemoteOrigin != "" {

q.Set("git_remote_origin", req.GitRemoteOrigin)

}

if req.ChatID != "" {

q.Set("chat_id", req.ChatID)

}

reqURL := "/api/v2/workspaceagents/me/external-auth?" + q.Encode()

res, err := c.SDK.Request(ctx, http.MethodGet, reqURL, nil)

if err != nil {

return ExternalAuthResponse{}, xerrors.Errorf("execute request: %w", err)

}

defer res.Body.Close()

if res.StatusCode != http.StatusOK {

return ExternalAuthResponse{}, codersdk.ReadBodyAsError(res)

}

var authResp ExternalAuthResponse

return authResp, json.NewDecoder(res.Body).Decode(&authResp)

}

// LogsNotifyChannel returns the channel name responsible for notifying

// of new logs.

func LogsNotifyChannel(agentID uuid.UUID) string {

return fmt.Sprintf("agent-logs:%s", agentID)

}

type LogsNotifyMessage struct {

CreatedAfter int64 `json:"created_after"`

}

type ReinitializationReason string

const (

ReinitializeReasonPrebuildClaimed ReinitializationReason = "prebuild_claimed"

)

type ReinitializationEvent struct {

WorkspaceID uuid.UUID `json:"workspace_id" format:"uuid"`

Reason ReinitializationReason `json:"reason"`

OwnerID uuid.UUID `json:"owner_id,omitzero" format:"uuid"`

}

func PrebuildClaimedChannel(id uuid.UUID) string {

return fmt.Sprintf("prebuild_claimed_%s", id)

}

// WaitForReinit polls a SSE endpoint, and receives an event back under the following conditions:

// - ping: ignored, keepalive

// - prebuild claimed: a prebuilt workspace is claimed, so the agent must reinitialize.

func (c *Client) WaitForReinit(ctx context.Context) (*ReinitializationEvent, error) {

rpcURL, err := c.SDK.URL.Parse("/api/v2/workspaceagents/me/reinit")

if err != nil {

return nil, xerrors.Errorf("parse url: %w", err)

}

q := rpcURL.Query()

q.Set("wait", "true")

rpcURL.RawQuery = q.Encode()

httpClient := &http.Client{

Transport: c.SDK.HTTPClient.Transport,

}

req, err := http.NewRequestWithContext(ctx, http.MethodGet, rpcURL.String(), nil)

if err != nil {

return nil, xerrors.Errorf("build request: %w", err)

}

req.Header[codersdk.SessionTokenHeader] = []string{c.SDK.SessionToken()}

res, err := httpClient.Do(req)

if err != nil {

return nil, xerrors.Errorf("execute request: %w", err)

}

defer res.Body.Close()

if res.StatusCode != http.StatusOK {

return nil, codersdk.ReadBodyAsError(res)

}

reinitEvent, err := NewSSEAgentReinitReceiver(res.Body).Receive(ctx)

if err != nil {

return nil, xerrors.Errorf("listening for reinitialization events: %w", err)

}

return reinitEvent, nil

}

// WaitForReinitLoop polls the /reinit SSE endpoint in a retry loop and

// forwards received reinitialization events to the returned channel. The

// channel is closed when ctx is canceled or the server returns 409

// Conflict (indicating the workspace is not a prebuilt workspace or the

// claim build failed permanently). The caller should select on both the

// channel and ctx.Done().

func WaitForReinitLoop(ctx context.Context, logger slog.Logger, client *Client) <-chan ReinitializationEvent {

reinitEvents := make(chan ReinitializationEvent)

go func() {

defer close(reinitEvents)

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

logger.Debug(ctx, "waiting for agent reinitialization instructions")

reinitEvent, err := client.WaitForReinit(ctx)

if err != nil {

var sdkErr *codersdk.Error

if errors.As(err, &sdkErr) && sdkErr.StatusCode() == http.StatusConflict {

logger.Info(ctx, "received terminal 409, stopping reinit polling",

slog.Error(sdkErr))

return

}

logger.Error(ctx, "failed to wait for agent reinitialization instructions", slog.Error(err))

continue

}

retrier.Reset()

select {

case <-ctx.Done():

return

case reinitEvents <- *reinitEvent:

}

}

}()

return reinitEvents

}

func NewSSEAgentReinitTransmitter(logger slog.Logger, rw http.ResponseWriter, r *http.Request) *SSEAgentReinitTransmitter {

return &SSEAgentReinitTransmitter{logger: logger, rw: rw, r: r}

}

type SSEAgentReinitTransmitter struct {

rw http.ResponseWriter

r *http.Request

logger slog.Logger

}

var (

ErrTransmissionSourceClosed = xerrors.New("transmission source closed")

ErrTransmissionTargetClosed = xerrors.New("transmission target closed")

)

// Transmit will read from the given chan and send events for as long as:

// * the chan remains open

// * the context has not been canceled

// * not timed out

// * the connection to the receiver remains open

func (s *SSEAgentReinitTransmitter) Transmit(ctx context.Context, reinitEvents <-chan ReinitializationEvent) error {

select {

case <-ctx.Done():

return ctx.Err()

default:

}

sseSendEvent, sseSenderClosed, err := httpapi.ServerSentEventSender(s.rw, s.r)

if err != nil {

return xerrors.Errorf("failed to create sse transmitter: %w", err)

}

defer func() {

// Block returning until the ServerSentEventSender is closed

// to avoid a race condition where we might write or flush to rw after the handler returns.

<-sseSenderClosed

}()

for {

select {

case <-ctx.Done():

return ctx.Err()

case <-sseSenderClosed:

return ErrTransmissionTargetClosed

case reinitEvent, ok := <-reinitEvents:

if !ok {

return ErrTransmissionSourceClosed

}

err := sseSendEvent(codersdk.ServerSentEvent{

Type: codersdk.ServerSentEventTypeData,

Data: reinitEvent,

})

if err != nil {

return err

}

}

}

}

func NewSSEAgentReinitReceiver(r io.ReadCloser) *SSEAgentReinitReceiver {

return &SSEAgentReinitReceiver{r: r}

}

type SSEAgentReinitReceiver struct {

r io.ReadCloser

}

func (s *SSEAgentReinitReceiver) Receive(ctx context.Context) (*ReinitializationEvent, error) {

nextEvent := codersdk.ServerSentEventReader(ctx, s.r)

for {

select {

case <-ctx.Done():

return nil, ctx.Err()

default:

}

sse, err := nextEvent()

switch {

case err != nil:

return nil, xerrors.Errorf("failed to read server-sent event: %w", err)

case sse.Type == codersdk.ServerSentEventTypeError:

return nil, xerrors.Errorf("unexpected server sent event type error")

case sse.Type == codersdk.ServerSentEventTypePing:

continue

case sse.Type != codersdk.ServerSentEventTypeData:

return nil, xerrors.Errorf("unexpected server sent event type: %s", sse.Type)

}

// At this point we know that the sent event is of type codersdk.ServerSentEventTypeData

var reinitEvent ReinitializationEvent

b, ok := sse.Data.([]byte)

if !ok {

return nil, xerrors.Errorf("expected data as []byte, got %T", sse.Data)

}

err = json.Unmarshal(b, &reinitEvent)

if err != nil {

return nil, xerrors.Errorf("unmarshal reinit response: %w", err)

}

return &reinitEvent, nil

}

}

// AddChatContextRequest is the request body for adding chat context.

type AddChatContextRequest struct {

// ChatID optionally identifies the chat to add context to.

// If empty, auto-detection is used (CODER_CHAT_ID env, the

// only active chat, or the only top-level active chat for this

// agent).

ChatID uuid.UUID `json:"chat_id,omitempty"`

// Parts are the context-file and skill parts to add.

Parts []codersdk.ChatMessagePart `json:"parts"`

}

// AddChatContextResponse is the response for adding chat context.

type AddChatContextResponse struct {

ChatID uuid.UUID `json:"chat_id"`

Count int `json:"count"`

}

// ClearChatContextRequest is the request body for clearing chat context.

type ClearChatContextRequest struct {

// ChatID optionally identifies the chat to clear context from.

// If empty, auto-detection is used (CODER_CHAT_ID env, the

// only active chat, or the only top-level active chat for this

// agent).

ChatID uuid.UUID `json:"chat_id,omitempty"`

}

// ClearChatContextResponse is the response for clearing chat context.

type ClearChatContextResponse struct {

ChatID uuid.UUID `json:"chat_id"`

}

// AddChatContext adds context-file and skill parts to an active chat.

func (c *Client) AddChatContext(ctx context.Context, req AddChatContextRequest) (AddChatContextResponse, error) {

res, err := c.SDK.Request(ctx, http.MethodPost, "/api/v2/workspaceagents/me/experimental/chat-context", req)

if err != nil {

return AddChatContextResponse{}, xerrors.Errorf("execute request: %w", err)

}

defer res.Body.Close()

if res.StatusCode != http.StatusOK {

return AddChatContextResponse{}, codersdk.ReadBodyAsError(res)

}

var resp AddChatContextResponse

return resp, json.NewDecoder(res.Body).Decode(&resp)

}

// ClearChatContext soft-deletes context-file and skill messages from an active chat.

func (c *Client) ClearChatContext(ctx context.Context, req ClearChatContextRequest) (ClearChatContextResponse, error) {

res, err := c.SDK.Request(ctx, http.MethodDelete, "/api/v2/workspaceagents/me/experimental/chat-context", req)

if err != nil {

return ClearChatContextResponse{}, xerrors.Errorf("execute request: %w", err)

}

defer res.Body.Close()

if res.StatusCode != http.StatusOK {

return ClearChatContextResponse{}, codersdk.ReadBodyAsError(res)

}