230 lines
10 KiB
Go
230 lines
10 KiB
Go
package orchestration
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"time"
|
|
|
|
"data-terminal-cloud/internal/conversation"
|
|
"data-terminal-cloud/internal/deviceauth"
|
|
"data-terminal-cloud/internal/rtc"
|
|
|
|
"github.com/jackc/pgx/v5"
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
)
|
|
|
|
type HeartbeatResult struct {
|
|
Response HeartbeatResponse
|
|
InternalID string
|
|
DevicePublicID string
|
|
Status conversation.Status
|
|
SequenceNo int64
|
|
Updated bool
|
|
Credential *CredentialRenewal
|
|
}
|
|
|
|
type CredentialRenewal struct {
|
|
RoomID, DeviceUserID, AgentUserID string
|
|
IssuedAt, ExpiresAt time.Time
|
|
}
|
|
|
|
type HeartbeatServiceOptions struct {
|
|
RTCAppID string
|
|
LeaseTTL time.Duration
|
|
TokenTTL time.Duration
|
|
CredentialRenewalWindow time.Duration
|
|
}
|
|
|
|
type HeartbeatService struct {
|
|
pool *pgxpool.Pool
|
|
repository *Repository
|
|
leases LeaseCache
|
|
tokens rtc.TokenIssuer
|
|
options HeartbeatServiceOptions
|
|
now func() time.Time
|
|
}
|
|
|
|
func NewHeartbeatService(pool *pgxpool.Pool, repository *Repository, leases LeaseCache, tokens rtc.TokenIssuer, options HeartbeatServiceOptions) (*HeartbeatService, error) {
|
|
if pool == nil || repository == nil || leases == nil || tokens == nil || options.RTCAppID == "" || options.LeaseTTL <= 0 || options.TokenTTL <= 0 || options.CredentialRenewalWindow <= 0 || options.CredentialRenewalWindow >= options.TokenTTL {
|
|
return nil, errors.New("heartbeat database, repository, lease cache, token issuer and valid TTL options are required")
|
|
}
|
|
return &HeartbeatService{pool: pool, repository: repository, leases: leases, tokens: tokens, options: options, now: time.Now}, nil
|
|
}
|
|
|
|
func (s *HeartbeatService) Heartbeat(ctx context.Context, identity deviceauth.Identity, conversationID string, input HeartbeatInput) (HeartbeatResponse, error) {
|
|
tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{})
|
|
if err != nil {
|
|
return HeartbeatResponse{}, fmt.Errorf("begin heartbeat transaction: %w", err)
|
|
}
|
|
defer tx.Rollback(context.Background())
|
|
now := s.now().UTC()
|
|
result, err := s.repository.Heartbeat(ctx, tx, identity, conversationID, input, now, s.options.LeaseTTL, s.options.TokenTTL, s.options.CredentialRenewalWindow)
|
|
if err != nil {
|
|
return HeartbeatResponse{}, err
|
|
}
|
|
if result.Credential != nil {
|
|
token, issueErr := s.tokens.Issue(result.Credential.RoomID, result.Credential.DeviceUserID, result.Credential.ExpiresAt)
|
|
if issueErr != nil {
|
|
return HeartbeatResponse{}, fmt.Errorf("renew RTC credential: %w", issueErr)
|
|
}
|
|
result.Response.Actions = append(result.Response.Actions, HeartbeatAction{Type: "renew_rtc_credential", RTC: &RTCGrant{
|
|
AppID: s.options.RTCAppID, RoomID: result.Credential.RoomID, UserID: result.Credential.DeviceUserID,
|
|
AgentUserID: result.Credential.AgentUserID, Token: token, ExpiresAt: result.Credential.ExpiresAt,
|
|
}})
|
|
}
|
|
if err := tx.Commit(ctx); err != nil {
|
|
return HeartbeatResponse{}, fmt.Errorf("commit heartbeat: %w", err)
|
|
}
|
|
if result.Updated && result.Response.LeaseExpiresAt != nil {
|
|
// PostgreSQL is authoritative. A missing Redis lease is recoverable and must not
|
|
// turn an already committed, sequence-protected heartbeat into an ambiguous retry.
|
|
_ = s.leases.Refresh(ctx, LeaseRecord{
|
|
ConversationID: conversationID, DeviceID: result.DevicePublicID,
|
|
SequenceNo: result.SequenceNo, Status: string(result.Status),
|
|
ExpiresAt: *result.Response.LeaseExpiresAt, ObservedAt: now,
|
|
})
|
|
}
|
|
return result.Response, nil
|
|
}
|
|
|
|
func (r *Repository) Heartbeat(
|
|
ctx context.Context,
|
|
tx pgx.Tx,
|
|
identity deviceauth.Identity,
|
|
conversationPublicID string,
|
|
input HeartbeatInput,
|
|
now time.Time,
|
|
leaseTTL, tokenTTL, credentialRenewalWindow time.Duration,
|
|
) (HeartbeatResult, error) {
|
|
var internalID string
|
|
var status conversation.Status
|
|
var lastSequence, desiredRevision int64
|
|
var currentLease, startedAt *time.Time
|
|
var credentialIssuedAt, credentialExpiresAt, conversationCreatedAt time.Time
|
|
var roomID, deviceRTCUserID, agentRTCUserID string
|
|
var credentialRenewalSupported bool
|
|
var maxDurationSec *int
|
|
var inProgressAction *string
|
|
err := tx.QueryRow(ctx, `
|
|
SELECT c.id::text, c.status, c.last_heartbeat_seq, c.lease_expires_at,
|
|
dt.desired_revision,c.started_at,cub.max_call_duration_sec,cub.in_progress_action,
|
|
r.credential_issued_at,r.credential_expires_at,c.created_at,r.room_id,r.device_rtc_user_id,r.agent_rtc_user_id,
|
|
c.rtc_credential_renewal_supported
|
|
FROM conversation c
|
|
JOIN device_twin dt ON dt.device_id=c.device_id AND dt.tenant_id=c.tenant_id
|
|
JOIN rtc_task r ON r.conversation_id=c.id
|
|
LEFT JOIN conversation_usage_budget cub ON cub.conversation_id=c.id AND cub.tenant_id=c.tenant_id
|
|
WHERE c.public_id=$1::uuid AND c.device_id=$2::uuid AND c.tenant_id=$3::uuid
|
|
FOR UPDATE OF c, dt, r`, conversationPublicID, identity.ID, identity.TenantID,
|
|
).Scan(&internalID, &status, &lastSequence, ¤tLease, &desiredRevision, &startedAt, &maxDurationSec, &inProgressAction,
|
|
&credentialIssuedAt, &credentialExpiresAt, &conversationCreatedAt, &roomID, &deviceRTCUserID, &agentRTCUserID,
|
|
&credentialRenewalSupported)
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return HeartbeatResult{}, ErrConversationNotFound
|
|
}
|
|
if err != nil {
|
|
return HeartbeatResult{}, fmt.Errorf("read heartbeat state: %w", err)
|
|
}
|
|
result := HeartbeatResult{
|
|
InternalID: internalID, DevicePublicID: identity.PublicID, Status: status,
|
|
SequenceNo: input.SequenceNo,
|
|
Response: HeartbeatResponse{
|
|
LeaseExpiresAt: currentLease, Actions: heartbeatActions(status),
|
|
DesiredConfigRevision: desiredRevision, Commands: []json.RawMessage{},
|
|
},
|
|
}
|
|
if !input.RTC.Joined || conversation.IsTerminal(input.ConversationStatus) || input.ConversationStatus == conversation.StatusStopping {
|
|
result.Response.Actions = []HeartbeatAction{{Type: "stop"}}
|
|
}
|
|
if startedAt != nil && maxDurationSec != nil && !now.Before(startedAt.Add(time.Duration(*maxDurationSec)*time.Second)) {
|
|
action := "stop_after_turn"
|
|
if inProgressAction != nil && *inProgressAction == "STOP_IMMEDIATELY" {
|
|
action = "stop"
|
|
}
|
|
result.Response.Actions = []HeartbeatAction{{Type: action}}
|
|
}
|
|
if input.ReportedConfigRevision > desiredRevision {
|
|
return HeartbeatResult{}, ErrConfigVersionConflict
|
|
}
|
|
canRenewCredential := credentialRenewalSupported && len(result.Response.Actions) == 1 && result.Response.Actions[0].Type == "continue"
|
|
storedCredentialIsNewer := canRenewCredential && credentialIssuedAt.After(conversationCreatedAt) &&
|
|
(input.RTCCredentialExpiresAt == nil || input.RTCCredentialExpiresAt.Before(credentialExpiresAt))
|
|
if storedCredentialIsNewer {
|
|
result.Credential = &CredentialRenewal{RoomID: roomID, DeviceUserID: deviceRTCUserID, AgentUserID: agentRTCUserID, IssuedAt: credentialIssuedAt, ExpiresAt: credentialExpiresAt}
|
|
}
|
|
if input.SequenceNo <= lastSequence || conversation.IsTerminal(status) {
|
|
return result, nil
|
|
}
|
|
if status != conversation.StatusStartingAgent && status != conversation.StatusActive && status != conversation.StatusStopping {
|
|
return HeartbeatResult{}, ErrConversationStateConflict
|
|
}
|
|
if canRenewCredential && result.Credential == nil && !now.Add(credentialRenewalWindow).Before(credentialExpiresAt) {
|
|
credentialIssuedAt = now
|
|
credentialExpiresAt = now.Add(tokenTTL)
|
|
if _, err := tx.Exec(ctx, `UPDATE rtc_task SET credential_issued_at=$2,credential_expires_at=$3,updated_at=$2 WHERE conversation_id=$1::uuid`, internalID, credentialIssuedAt, credentialExpiresAt); err != nil {
|
|
return HeartbeatResult{}, fmt.Errorf("rotate RTC credential metadata: %w", err)
|
|
}
|
|
result.Credential = &CredentialRenewal{RoomID: roomID, DeviceUserID: deviceRTCUserID, AgentUserID: agentRTCUserID, IssuedAt: credentialIssuedAt, ExpiresAt: credentialExpiresAt}
|
|
}
|
|
|
|
reportedHeartbeat, err := json.Marshal(map[string]any{
|
|
"sequence_no": input.SequenceNo, "device_time": input.DeviceTime,
|
|
"monotonic_ms": input.MonotonicMS, "conversation_status": input.ConversationStatus,
|
|
"rtc": input.RTC, "network": input.Network, "recording_queue": input.RecordingQueue,
|
|
"last_error": input.LastError,
|
|
})
|
|
if err != nil {
|
|
return HeartbeatResult{}, fmt.Errorf("encode reported heartbeat: %w", err)
|
|
}
|
|
if status == conversation.StatusStartingAgent || status == conversation.StatusActive {
|
|
leaseExpiresAt := now.Add(leaseTTL)
|
|
result.Response.LeaseExpiresAt = &leaseExpiresAt
|
|
_, err = tx.Exec(ctx, `
|
|
UPDATE conversation SET last_heartbeat_seq=$2, lease_expires_at=$3,
|
|
version=version+1, updated_at=$4 WHERE id=$1::uuid`,
|
|
internalID, input.SequenceNo, leaseExpiresAt, now)
|
|
if err == nil {
|
|
_, err = tx.Exec(ctx, `
|
|
UPDATE rtc_task SET lease_expires_at=$2, last_checked_at=$3, updated_at=$3
|
|
WHERE conversation_id=$1::uuid`, internalID, leaseExpiresAt, now)
|
|
}
|
|
} else {
|
|
_, err = tx.Exec(ctx, `
|
|
UPDATE conversation SET last_heartbeat_seq=$2, version=version+1, updated_at=$3
|
|
WHERE id=$1::uuid`, internalID, input.SequenceNo, now)
|
|
}
|
|
if err != nil {
|
|
return HeartbeatResult{}, fmt.Errorf("refresh conversation heartbeat: %w", err)
|
|
}
|
|
_, err = tx.Exec(ctx, `
|
|
UPDATE device_twin SET
|
|
reported_revision=CASE
|
|
WHEN reported_revision IS NULL OR reported_revision <= $2 THEN $2
|
|
ELSE reported_revision END,
|
|
reported_config=CASE
|
|
WHEN reported_revision IS NULL OR reported_revision <= $2
|
|
THEN reported_config || jsonb_build_object('heartbeat', $3::jsonb)
|
|
ELSE reported_config END,
|
|
firmware_version=CASE WHEN NULLIF($4,'') IS NULL THEN firmware_version ELSE $4 END,
|
|
last_reported_at=$5, updated_at=$5
|
|
WHERE device_id=$1::uuid AND tenant_id=$6::uuid`,
|
|
identity.ID, input.ReportedConfigRevision, reportedHeartbeat, input.FirmwareVersion, now, identity.TenantID)
|
|
if err != nil {
|
|
return HeartbeatResult{}, fmt.Errorf("update device heartbeat report: %w", err)
|
|
}
|
|
if _, err = tx.Exec(ctx, `SELECT record_device_presence($1::uuid,$2::uuid,$3)`, identity.TenantID, identity.ID, now); err != nil {
|
|
return HeartbeatResult{}, fmt.Errorf("record device presence: %w", err)
|
|
}
|
|
result.Updated = true
|
|
return result, nil
|
|
}
|
|
|
|
func heartbeatActions(status conversation.Status) []HeartbeatAction {
|
|
if status == conversation.StatusStartingAgent || status == conversation.StatusActive {
|
|
return []HeartbeatAction{{Type: "continue"}}
|
|
}
|
|
return []HeartbeatAction{{Type: "stop"}}
|
|
}
|