coder/codersdk/aibridge.go

package codersdk

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"time"

	"github.com/google/uuid"
	"golang.org/x/xerrors"

	"github.com/coder/coder/v2/coderd/util/slice"
)

// AIBudgetLimitSource identifies which tier produced the user's
// effective budget limit.
type AIBudgetLimitSource string

const (
	// AIBudgetLimitSourceUserOverride indicates the limit came from a
	// per-user override.
	AIBudgetLimitSourceUserOverride AIBudgetLimitSource = "user_override"
	// AIBudgetLimitSourceGroup indicates the limit came from a group
	// budget selected by the deployment budget policy.
	AIBudgetLimitSourceGroup AIBudgetLimitSource = "group"
)

// AIGroupBudget is an AI spend limit and the tier that produced it. Both
// fields are always populated together.
type AIGroupBudget struct {
	SpendLimitMicros int64               `json:"spend_limit_micros"`
	LimitSource      AIBudgetLimitSource `json:"limit_source"`
}

// UserAIBudgetSummary is the effective AI budget for a user. When no budget
// applies, the effective group falls back to the Everyone group with a null
// limit and source.
type UserAIBudgetSummary struct {
	UserID uuid.UUID `json:"user_id" format:"uuid"`
	// EffectiveGroupID is the group the spend is attributed to, falling back to
	// the Everyone group when no budget applies. Null only when the user has no
	// organization membership.
	EffectiveGroupID *uuid.UUID `json:"effective_group_id" format:"uuid"`
	// SpendLimitMicros is the effective spend limit in micro-units.
	// Null when no budget applies to the user (unlimited).
	SpendLimitMicros *int64 `json:"spend_limit_micros"`
	// LimitSource identifies which tier produced the limit. Null when no
	// budget applies.
	LimitSource *AIBudgetLimitSource `json:"limit_source"`
}

// AISpendPeriodWindow is the [Start, End) window over which AI spend is
// aggregated.
type AISpendPeriodWindow struct {
	// PeriodStart is the inclusive lower bound of the current budget
	// period.
	PeriodStart time.Time `json:"period_start" format:"date-time"`
	// PeriodEnd is the exclusive upper bound of the current budget
	// period.
	PeriodEnd time.Time `json:"period_end" format:"date-time"`
}

// UserAISpendStatus is the current AI spend snapshot for a user within
// the active budget period.
type UserAISpendStatus struct {
	UserAIBudgetSummary
	AISpendPeriodWindow
	// CurrentSpendMicros is the user's spend on their effective group over
	// the current budget period.
	CurrentSpendMicros int64 `json:"current_spend_micros"`
}

// OrganizationGroupsAISpend reports AI spend for a set of groups in the
// active budget period.
type OrganizationGroupsAISpend struct {
	AISpendPeriodWindow
	Groups []OrganizationGroupAISpend `json:"groups"`
}

// OrganizationGroupAISpend is the current AI spend snapshot for a group
// within the active budget period.
type OrganizationGroupAISpend struct {
	GroupID uuid.UUID `json:"group_id" format:"uuid"`
	// SpendLimitMicros is the group's configured AI spend limit. Null when
	// the group has no configured budget.
	SpendLimitMicros *int64 `json:"spend_limit_micros"`
	// CurrentSpendMicros is the group's spend over the current budget
	// period.
	CurrentSpendMicros int64 `json:"current_spend_micros"`
}

// GroupMembersAISpend reports per-member AI spend attributed to a specific
// group in the active budget period.
type GroupMembersAISpend struct {
	AISpendPeriodWindow
	Members []GroupMemberAISpend `json:"members"`
}

// GroupMemberAISpend is a single member's AI spend attributed to the queried
// group in the current budget period.
type GroupMemberAISpend struct {
	UserID uuid.UUID `json:"user_id" format:"uuid"`
	// EffectiveGroupID is the user's effective budget group within the queried
	// group's organization, falling back to the Everyone group when no budget
	// applies. Null when the effective group belongs to a different organization
	// than the queried group.
	EffectiveGroupID *uuid.UUID `json:"effective_group_id" format:"uuid"`
	// GroupBudget is the budget when the queried group is this user's
	// effective budget source. Null when the user's budget resolves to another
	// group or no budget applies to the user.
	GroupBudget *AIGroupBudget `json:"group_budget"`
	// GroupSpendMicros is the user's spend attributed to the queried group
	// over the current budget period.
	GroupSpendMicros int64 `json:"group_spend_micros"`
}

type AIBridgeSession struct {
	ID                string                           `json:"id"`
	Initiator         MinimalUser                      `json:"initiator"`
	Providers         []string                         `json:"providers"`
	Models            []string                         `json:"models"`
	Client            *string                          `json:"client"`
	Metadata          map[string]any                   `json:"metadata"`
	StartedAt         time.Time                        `json:"started_at" format:"date-time"`
	EndedAt           *time.Time                       `json:"ended_at,omitempty" format:"date-time"`
	Threads           int64                            `json:"threads"`
	TokenUsageSummary AIBridgeSessionTokenUsageSummary `json:"token_usage_summary"`
	// NetworkCalls summarizes the Agent Firewall network calls made during the
	// session. A nil value means the session did not pass through Agent
	// Firewall, so network call monitoring was not active, which the UI
	// surfaces as "Disabled".
	NetworkCalls *AIBridgeSessionNetworkCallSummary `json:"network_calls,omitempty"`
	LastPrompt   *string                            `json:"last_prompt,omitempty"`
	LastActiveAt time.Time                          `json:"last_active_at" format:"date-time"`
}

type AIBridgeSessionTokenUsageSummary struct {
	InputTokens           int64 `json:"input_tokens"`
	OutputTokens          int64 `json:"output_tokens"`
	CacheReadInputTokens  int64 `json:"cache_read_input_tokens"`
	CacheWriteInputTokens int64 `json:"cache_write_input_tokens"`
}

// AIBridgeSessionNetworkCallSummary aggregates the Agent Firewall network
// calls made during a session. Blocked counts calls denied by the firewall
// allow-list.
type AIBridgeSessionNetworkCallSummary struct {
	Total   int64 `json:"total"`
	Blocked int64 `json:"blocked"`
}

type AIBridgeListSessionsResponse struct {
	Count    int64             `json:"count"`
	Sessions []AIBridgeSession `json:"sessions"`
}

// AIBridgeSessionThreadsResponse is the response for GET
// /api/v2/ai-gateway/sessions/{session_id} which returns a single
// session with fully expanded threads.
type AIBridgeSessionThreadsResponse struct {
	ID                string                           `json:"id"`
	Initiator         MinimalUser                      `json:"initiator"`
	Providers         []string                         `json:"providers"`
	Models            []string                         `json:"models"`
	Client            *string                          `json:"client,omitempty"`
	Metadata          map[string]any                   `json:"metadata"`
	PageStartedAt     *time.Time                       `json:"page_started_at,omitempty" format:"date-time"`
	PageEndedAt       *time.Time                       `json:"page_ended_at,omitempty" format:"date-time"`
	StartedAt         time.Time                        `json:"started_at" format:"date-time"`
	EndedAt           *time.Time                       `json:"ended_at,omitempty" format:"date-time"`
	TokenUsageSummary AIBridgeSessionThreadsTokenUsage `json:"token_usage_summary"`
	Threads           []AIBridgeThread                 `json:"threads"`
}

// AIBridgeSessionThreadsTokenUsage represents aggregated token usage
// with metadata containing provider-specific fields.
type AIBridgeSessionThreadsTokenUsage struct {
	InputTokens           int64          `json:"input_tokens"`
	OutputTokens          int64          `json:"output_tokens"`
	CacheReadInputTokens  int64          `json:"cache_read_input_tokens"`
	CacheWriteInputTokens int64          `json:"cache_write_input_tokens"`
	Metadata              map[string]any `json:"metadata"`
}

// AIBridgeThread represents a single thread within a session.
// A thread groups interceptions by their thread_root_id.
type AIBridgeThread struct {
	ID             uuid.UUID                        `json:"id" format:"uuid"`
	Prompt         *string                          `json:"prompt,omitempty"`
	Model          string                           `json:"model"`
	Provider       string                           `json:"provider"`
	CredentialKind string                           `json:"credential_kind"`
	CredentialHint string                           `json:"credential_hint"`
	StartedAt      time.Time                        `json:"started_at" format:"date-time"`
	EndedAt        *time.Time                       `json:"ended_at,omitempty" format:"date-time"`
	TokenUsage     AIBridgeSessionThreadsTokenUsage `json:"token_usage"`
	AgenticActions []AIBridgeAgenticAction          `json:"agentic_actions"`
	// ErrorType is the categorized terminal upstream error from the root
	// interception, or nil when the interception succeeded. See the
	// aibridge_interception_error_type enum for possible values.
	ErrorType *string `json:"error_type,omitempty"`
	// ErrorMessage is the raw terminal upstream error message from the root
	// interception. Nil when the interception succeeded.
	ErrorMessage *string `json:"error_message,omitempty"`
	// AgentFirewallSessionID links this thread to an agent firewall
	// confinement session. Nil when the request did not pass through
	// the agent firewall.
	AgentFirewallSessionID *uuid.UUID `json:"agent_firewall_session_id,omitempty" format:"uuid"`
	// AgentFirewallSequenceNumber is the firewall sequence number from
	// the root interception. Used to determine the position of this
	// LLM request in the firewall event stream. Nil when the request
	// did not pass through the agent firewall.
	AgentFirewallSequenceNumber *int32 `json:"agent_firewall_sequence_number,omitempty"`
}

// AIBridgeAgenticAction represents a tool call with associated
// thinking blocks and token usage from one or more interceptions.
type AIBridgeAgenticAction struct {
	Model      string                           `json:"model"`
	TokenUsage AIBridgeSessionThreadsTokenUsage `json:"token_usage"`
	Thinking   []AIBridgeModelThought           `json:"thinking"`
	ToolCalls  []AIBridgeToolCall               `json:"tool_calls"`
}

// AIBridgeModelThought represents a single thinking block from
// the model.
type AIBridgeModelThought struct {
	Text string `json:"text"`
}

// AIBridgeToolCall represents a tool call recorded during an
// interception.
type AIBridgeToolCall struct {
	ID                 uuid.UUID      `json:"id" format:"uuid"`
	InterceptionID     uuid.UUID      `json:"interception_id" format:"uuid"`
	ProviderResponseID string         `json:"provider_response_id"`
	ServerURL          string         `json:"server_url"`
	Tool               string         `json:"tool"`
	Injected           bool           `json:"injected"`
	Input              string         `json:"input"`
	Metadata           map[string]any `json:"metadata"`
	CreatedAt          time.Time      `json:"created_at" format:"date-time"`
}

// @typescript-ignore AIBridgeListSessionsFilter
type AIBridgeListSessionsFilter struct {
	// Limit defaults to 100, max is 1000.
	Pagination Pagination `json:"pagination,omitempty"`

	// Initiator is a user ID, username, or "me".
	Initiator     string    `json:"initiator,omitempty"`
	StartedBefore time.Time `json:"started_before,omitempty" format:"date-time"`
	StartedAfter  time.Time `json:"started_after,omitempty" format:"date-time"`
	// Provider matches the runtime provider type column (openai,
	// anthropic, copilot). The runtime type collapses the configured
	// ai_provider_type: azure, google, openai-compat, openrouter, and
	// vercel route through openai; bedrock routes through anthropic.
	// Retained for backward compatibility; new clients should prefer
	// ProviderName, which scopes to a specific configured row.
	Provider     string `json:"provider,omitempty"`
	ProviderName string `json:"provider_name,omitempty"`
	Model        string `json:"model,omitempty"`
	Client       string `json:"client,omitempty"`
	SessionID    string `json:"session_id,omitempty"`

	// AfterSessionID is a cursor for pagination. It is the session ID of the
	// last session in the previous page.
	AfterSessionID string `json:"after_session_id,omitempty"`

	FilterQuery string `json:"q,omitempty"`
}

// asRequestOption returns a function that can be used in (*Client).Request.
func (f AIBridgeListSessionsFilter) asRequestOption() RequestOption {
	return func(r *http.Request) {
		var params []string
		if f.Initiator != "" {
			params = append(params, fmt.Sprintf("initiator:%q", f.Initiator))
		}
		if !f.StartedBefore.IsZero() {
			params = append(params, fmt.Sprintf("started_before:%q", f.StartedBefore.Format(time.RFC3339Nano)))
		}
		if !f.StartedAfter.IsZero() {
			params = append(params, fmt.Sprintf("started_after:%q", f.StartedAfter.Format(time.RFC3339Nano)))
		}
		if f.Provider != "" {
			params = append(params, fmt.Sprintf("provider:%q", f.Provider))
		}
		if f.ProviderName != "" {
			params = append(params, fmt.Sprintf("provider_name:%q", f.ProviderName))
		}
		if f.Model != "" {
			params = append(params, fmt.Sprintf("model:%q", f.Model))
		}
		if f.Client != "" {
			params = append(params, fmt.Sprintf("client:%q", f.Client))
		}
		if f.SessionID != "" {
			params = append(params, fmt.Sprintf("session_id:%q", f.SessionID))
		}
		if f.FilterQuery != "" {
			params = append(params, f.FilterQuery)
		}

		q := r.URL.Query()
		q.Set("q", strings.Join(params, " "))
		if f.AfterSessionID != "" {
			q.Set("after_session_id", f.AfterSessionID)
		}
		r.URL.RawQuery = q.Encode()
	}
}

// AIBridgeListSessions returns AI Bridge sessions with the given filter.
func (c *Client) AIBridgeListSessions(ctx context.Context, filter AIBridgeListSessionsFilter) (AIBridgeListSessionsResponse, error) {
	res, err := c.Request(ctx, http.MethodGet, "/api/v2/ai-gateway/sessions", nil, filter.asRequestOption(), filter.Pagination.asRequestOption())
	if err != nil {
		return AIBridgeListSessionsResponse{}, err
	}
	defer res.Body.Close()
	if res.StatusCode != http.StatusOK {
		return AIBridgeListSessionsResponse{}, ReadBodyAsError(res)
	}
	var resp AIBridgeListSessionsResponse
	return resp, json.NewDecoder(res.Body).Decode(&resp)
}

// AIBridgeGetSessionThreads returns a single session with expanded
// thread details including agentic actions and thinking blocks.
func (c *Client) AIBridgeGetSessionThreads(ctx context.Context, sessionID string, afterID, beforeID uuid.UUID, limit int32) (AIBridgeSessionThreadsResponse, error) {
	res, err := c.Request(ctx, http.MethodGet, fmt.Sprintf("/api/v2/ai-gateway/sessions/%s", sessionID), nil, func(r *http.Request) {
		q := r.URL.Query()
		if afterID != uuid.Nil {
			q.Set("after_id", afterID.String())
		}
		if beforeID != uuid.Nil {
			q.Set("before_id", beforeID.String())
		}
		if limit > 0 {
			q.Set("limit", fmt.Sprintf("%d", limit))
		}
		r.URL.RawQuery = q.Encode()
	})
	if err != nil {
		return AIBridgeSessionThreadsResponse{}, err
	}
	defer res.Body.Close()
	if res.StatusCode != http.StatusOK {
		return AIBridgeSessionThreadsResponse{}, ReadBodyAsError(res)
	}
	var resp AIBridgeSessionThreadsResponse
	return resp, json.NewDecoder(res.Body).Decode(&resp)
}

// AIBridgeListClients returns the distinct AI clients visible to the caller.
func (c *Client) AIBridgeListClients(ctx context.Context) ([]string, error) {
	res, err := c.Request(ctx, http.MethodGet, "/api/v2/ai-gateway/clients", nil)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()
	if res.StatusCode != http.StatusOK {
		return nil, ReadBodyAsError(res)
	}
	var clients []string
	return clients, json.NewDecoder(res.Body).Decode(&clients)
}

type GroupAIBudget struct {
	GroupID          uuid.UUID `json:"group_id" format:"uuid"`
	SpendLimitMicros int64     `json:"spend_limit_micros"`
	CreatedAt        time.Time `json:"created_at" format:"date-time"`
	UpdatedAt        time.Time `json:"updated_at" format:"date-time"`
}

type UpsertGroupAIBudgetRequest struct {
	SpendLimitMicros int64 `json:"spend_limit_micros" validate:"gte=0"`
}

// GroupAIBudget returns the AI spend budget configured for the given group.
func (c *Client) GroupAIBudget(ctx context.Context, group uuid.UUID) (GroupAIBudget, error) {
	res, err := c.Request(ctx, http.MethodGet,
		fmt.Sprintf("/api/v2/groups/%s/ai/budget", group.String()),
		nil,
	)
	if err != nil {
		return GroupAIBudget{}, xerrors.Errorf("make request: %w", err)
	}
	defer res.Body.Close()

	if res.StatusCode != http.StatusOK {
		return GroupAIBudget{}, ReadBodyAsError(res)
	}
	var resp GroupAIBudget
	return resp, json.NewDecoder(res.Body).Decode(&resp)
}

// UpsertGroupAIBudget creates or updates the AI spend budget for the given group.
func (c *Client) UpsertGroupAIBudget(ctx context.Context, group uuid.UUID, req UpsertGroupAIBudgetRequest) (GroupAIBudget, error) {
	res, err := c.Request(ctx, http.MethodPut,
		fmt.Sprintf("/api/v2/groups/%s/ai/budget", group.String()),
		req,
	)
	if err != nil {
		return GroupAIBudget{}, xerrors.Errorf("make request: %w", err)
	}
	defer res.Body.Close()

	if res.StatusCode != http.StatusOK {
		return GroupAIBudget{}, ReadBodyAsError(res)
	}
	var resp GroupAIBudget
	return resp, json.NewDecoder(res.Body).Decode(&resp)
}

// DeleteGroupAIBudget removes the AI spend budget for the given group.
func (c *Client) DeleteGroupAIBudget(ctx context.Context, group uuid.UUID) error {
	res, err := c.Request(ctx, http.MethodDelete,
		fmt.Sprintf("/api/v2/groups/%s/ai/budget", group.String()),
		nil,
	)
	if err != nil {
		return xerrors.Errorf("make request: %w", err)
	}
	defer res.Body.Close()

	if res.StatusCode != http.StatusNoContent {
		return ReadBodyAsError(res)
	}
	return nil
}

type UserAIBudgetOverride struct {
	UserID           uuid.UUID `json:"user_id" format:"uuid"`
	GroupID          uuid.UUID `json:"group_id" format:"uuid"`
	SpendLimitMicros int64     `json:"spend_limit_micros"`
	CreatedAt        time.Time `json:"created_at" format:"date-time"`
	UpdatedAt        time.Time `json:"updated_at" format:"date-time"`
}

type UpsertUserAIBudgetOverrideRequest struct {
	// GroupID is the group the user's spend is attributed to. The user must
	// be a member of this group.
	GroupID          uuid.UUID `json:"group_id" format:"uuid" validate:"required"`
	SpendLimitMicros int64     `json:"spend_limit_micros" validate:"gte=0"`
}

// UserAIBudgetOverride returns the AI spend budget override configured for the given user.
func (c *Client) UserAIBudgetOverride(ctx context.Context, user uuid.UUID) (UserAIBudgetOverride, error) {
	res, err := c.Request(ctx, http.MethodGet,
		fmt.Sprintf("/api/v2/users/%s/ai/budget", user.String()),
		nil,
	)
	if err != nil {
		return UserAIBudgetOverride{}, xerrors.Errorf("make request: %w", err)
	}
	defer res.Body.Close()

	if res.StatusCode != http.StatusOK {
		return UserAIBudgetOverride{}, ReadBodyAsError(res)
	}
	var resp UserAIBudgetOverride
	return resp, json.NewDecoder(res.Body).Decode(&resp)
}

// UpsertUserAIBudgetOverride creates or updates the AI spend budget override for the given user.
func (c *Client) UpsertUserAIBudgetOverride(ctx context.Context, user uuid.UUID, req UpsertUserAIBudgetOverrideRequest) (UserAIBudgetOverride, error) {
	res, err := c.Request(ctx, http.MethodPut,
		fmt.Sprintf("/api/v2/users/%s/ai/budget", user.String()),
		req,
	)
	if err != nil {
		return UserAIBudgetOverride{}, xerrors.Errorf("make request: %w", err)
	}
	defer res.Body.Close()

	if res.StatusCode != http.StatusOK {
		return UserAIBudgetOverride{}, ReadBodyAsError(res)
	}
	var resp UserAIBudgetOverride
	return resp, json.NewDecoder(res.Body).Decode(&resp)
}

// DeleteUserAIBudgetOverride removes the AI spend budget override for the given user.
func (c *Client) DeleteUserAIBudgetOverride(ctx context.Context, user uuid.UUID) error {
	res, err := c.Request(ctx, http.MethodDelete,
		fmt.Sprintf("/api/v2/users/%s/ai/budget", user.String()),
		nil,
	)
	if err != nil {
		return xerrors.Errorf("make request: %w", err)
	}
	defer res.Body.Close()

	if res.StatusCode != http.StatusNoContent {
		return ReadBodyAsError(res)
	}
	return nil
}

// UserAISpendStatus returns the current AI spend snapshot for the given user
// within the active budget period.
func (c *Client) UserAISpendStatus(ctx context.Context, user uuid.UUID) (UserAISpendStatus, error) {
	res, err := c.Request(ctx, http.MethodGet,
		fmt.Sprintf("/api/v2/users/%s/ai/spend", user.String()),
		nil,
	)
	if err != nil {
		return UserAISpendStatus{}, xerrors.Errorf("make request: %w", err)
	}
	defer res.Body.Close()

	if res.StatusCode != http.StatusOK {
		return UserAISpendStatus{}, ReadBodyAsError(res)
	}
	var resp UserAISpendStatus
	return resp, json.NewDecoder(res.Body).Decode(&resp)
}

// OrganizationGroupsAISpend returns AI spend for the given groups within the
// organization for the active budget period. At most 100 group IDs may be
// requested per call, and callers with more groups are expected to batch
// across multiple requests.
func (c *Client) OrganizationGroupsAISpend(ctx context.Context, organization uuid.UUID, groupIDs []uuid.UUID) (OrganizationGroupsAISpend, error) {
	ids := slice.List(groupIDs, func(id uuid.UUID) string { return id.String() })
	res, err := c.Request(ctx, http.MethodGet,
		fmt.Sprintf("/api/v2/organizations/%s/groups/ai/spend", organization.String()),
		nil,
		func(r *http.Request) {
			q := r.URL.Query()
			q.Set("group_ids", strings.Join(ids, ","))
			r.URL.RawQuery = q.Encode()
		},
	)
	if err != nil {
		return OrganizationGroupsAISpend{}, xerrors.Errorf("make request: %w", err)
	}
	defer res.Body.Close()

	if res.StatusCode != http.StatusOK {
		return OrganizationGroupsAISpend{}, ReadBodyAsError(res)
	}
	var resp OrganizationGroupsAISpend
	return resp, json.NewDecoder(res.Body).Decode(&resp)
}

// GroupMembersAISpend returns AI spend attributed to the given group for the
// specified users within the active budget period. At most 100 user IDs may be
// requested per call, and callers with more members are expected to batch
// across multiple requests.
func (c *Client) GroupMembersAISpend(ctx context.Context, group uuid.UUID, userIDs []uuid.UUID) (GroupMembersAISpend, error) {
	ids := slice.List(userIDs, func(id uuid.UUID) string { return id.String() })
	res, err := c.Request(ctx, http.MethodGet,
		fmt.Sprintf("/api/v2/groups/%s/members/ai/spend", group.String()),
		nil,
		func(r *http.Request) {
			q := r.URL.Query()
			q.Set("user_ids", strings.Join(ids, ","))
			r.URL.RawQuery = q.Encode()
		},
	)
	if err != nil {
		return GroupMembersAISpend{}, xerrors.Errorf("make request: %w", err)
	}
	defer res.Body.Close()

	if res.StatusCode != http.StatusOK {
		return GroupMembersAISpend{}, ReadBodyAsError(res)
	}
	var resp GroupMembersAISpend
	return resp, json.NewDecoder(res.Body).Decode(&resp)
}