diff --git a/pkg/coredata/agent_run.go b/pkg/coredata/agent_run.go new file mode 100644 index 000000000..d7ea4df4d --- /dev/null +++ b/pkg/coredata/agent_run.go @@ -0,0 +1,789 @@ +// Copyright (c) 2026 Probo Inc . +// +// Permission to use, copy, modify, and/or distribute this software for any +// purpose with or without fee is hereby granted, provided that the above +// copyright notice and this permission notice appear in all copies. +// +// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH +// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY +// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, +// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM +// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR +// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR +// PERFORMANCE OF THIS SOFTWARE. + +package coredata + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "maps" + "time" + + "github.com/jackc/pgx/v5" + "go.gearno.de/kit/pg" + "go.probo.inc/probo/pkg/agent" + "go.probo.inc/probo/pkg/gid" + "go.probo.inc/probo/pkg/page" +) + +type ( + AgentRunStatus string + + AgentRun struct { + ID gid.GID `db:"id"` + OrganizationID gid.GID `db:"organization_id"` + StartAgentName string `db:"start_agent_name"` + Status AgentRunStatus `db:"status"` + Checkpoint json.RawMessage `db:"checkpoint"` + InputMessages json.RawMessage `db:"input_messages"` + Result json.RawMessage `db:"result"` + ErrorMessage *string `db:"error_message"` + StopRequested bool `db:"stop_requested"` + StartedAt *time.Time `db:"started_at"` + LeaseOwner *string `db:"lease_owner"` + LeaseExpiresAt *time.Time `db:"lease_expires_at"` + CreatedAt time.Time `db:"created_at"` + UpdatedAt time.Time `db:"updated_at"` + } + + AgentRuns []*AgentRun +) + +const ( + AgentRunStatusPending AgentRunStatus = "PENDING" + AgentRunStatusRunning AgentRunStatus = "RUNNING" + AgentRunStatusSuspended AgentRunStatus = "SUSPENDED" + AgentRunStatusAwaitingApproval AgentRunStatus = "AWAITING_APPROVAL" + AgentRunStatusCompleted AgentRunStatus = "COMPLETED" + AgentRunStatusFailed AgentRunStatus = "FAILED" +) + +func (e AgentRun) CursorKey(orderBy AgentRunOrderField) page.CursorKey { + switch orderBy { + case AgentRunOrderFieldCreatedAt: + return page.NewCursorKey(e.ID, e.CreatedAt) + } + + panic(fmt.Sprintf("unsupported order by: %s", orderBy)) +} + +func (e *AgentRun) AuthorizationAttributes(ctx context.Context, conn pg.Querier) (map[string]string, error) { + q := `SELECT organization_id FROM agent_runs WHERE id = @id LIMIT 1;` + + args := pgx.StrictNamedArgs{"id": e.ID.String()} + + rows, err := conn.Query(ctx, q, args) + if err != nil { + return nil, fmt.Errorf("cannot query agent run authorization attributes: %w", err) + } + + type row struct { + OrganizationID gid.GID `db:"organization_id"` + } + + r, err := pgx.CollectExactlyOneRow(rows, pgx.RowToStructByName[row]) + if err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return nil, ErrResourceNotFound + } + return nil, fmt.Errorf("cannot load agent run authorization attributes: %w", err) + } + + return map[string]string{"organization_id": r.OrganizationID.String()}, nil +} + +func (e *AgentRun) LoadByID( + ctx context.Context, + conn pg.Querier, + scope Scoper, + id gid.GID, +) error { + q := ` +SELECT + id, + organization_id, + start_agent_name, + status, + checkpoint, + input_messages, + result, + error_message, + stop_requested, + started_at, + lease_owner, + lease_expires_at, + created_at, + updated_at +FROM + agent_runs +WHERE + %s + AND id = @id +LIMIT 1; +` + + q = fmt.Sprintf(q, scope.SQLFragment()) + + args := pgx.StrictNamedArgs{"id": id.String()} + maps.Copy(args, scope.SQLArguments()) + + rows, err := conn.Query(ctx, q, args) + if err != nil { + return fmt.Errorf("cannot query agent run: %w", err) + } + + entity, err := pgx.CollectExactlyOneRow(rows, pgx.RowToStructByName[AgentRun]) + if err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return ErrResourceNotFound + } + return fmt.Errorf("cannot load agent run: %w", err) + } + + *e = entity + + return nil +} + +func (e *AgentRun) LoadByIDForUpdate( + ctx context.Context, + tx pg.Tx, + scope Scoper, + id gid.GID, +) error { + q := ` +SELECT + id, + organization_id, + start_agent_name, + status, + checkpoint, + input_messages, + result, + error_message, + stop_requested, + started_at, + lease_owner, + lease_expires_at, + created_at, + updated_at +FROM + agent_runs +WHERE + %s + AND id = @id +LIMIT 1 +FOR UPDATE; +` + + q = fmt.Sprintf(q, scope.SQLFragment()) + + args := pgx.StrictNamedArgs{"id": id.String()} + maps.Copy(args, scope.SQLArguments()) + + rows, err := tx.Query(ctx, q, args) + if err != nil { + return fmt.Errorf("cannot query agent run: %w", err) + } + + entity, err := pgx.CollectExactlyOneRow(rows, pgx.RowToStructByName[AgentRun]) + if err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return ErrResourceNotFound + } + return fmt.Errorf("cannot load agent run: %w", err) + } + + *e = entity + + return nil +} + +func (rs *AgentRuns) LoadByOrganizationID( + ctx context.Context, + conn pg.Querier, + scope Scoper, + organizationID gid.GID, + cursor *page.Cursor[AgentRunOrderField], +) error { + q := ` +SELECT + id, + organization_id, + start_agent_name, + status, + checkpoint, + input_messages, + result, + error_message, + stop_requested, + started_at, + lease_owner, + lease_expires_at, + created_at, + updated_at +FROM + agent_runs +WHERE + %s + AND organization_id = @organization_id + AND %s +` + + q = fmt.Sprintf(q, scope.SQLFragment(), cursor.SQLFragment()) + + args := pgx.StrictNamedArgs{"organization_id": organizationID.String()} + maps.Copy(args, scope.SQLArguments()) + maps.Copy(args, cursor.SQLArguments()) + + rows, err := conn.Query(ctx, q, args) + if err != nil { + return fmt.Errorf("cannot query agent runs: %w", err) + } + + entities, err := pgx.CollectRows(rows, pgx.RowToAddrOfStructByName[AgentRun]) + if err != nil { + return fmt.Errorf("cannot collect agent runs: %w", err) + } + + *rs = entities + + return nil +} + +func (rs *AgentRuns) CountByOrganizationID( + ctx context.Context, + conn pg.Querier, + scope Scoper, + organizationID gid.GID, +) (int, error) { + q := ` +SELECT + COUNT(id) +FROM + agent_runs +WHERE + %s + AND organization_id = @organization_id; +` + + q = fmt.Sprintf(q, scope.SQLFragment()) + + args := pgx.StrictNamedArgs{"organization_id": organizationID.String()} + maps.Copy(args, scope.SQLArguments()) + + var count int + if err := conn.QueryRow(ctx, q, args).Scan(&count); err != nil { + return 0, fmt.Errorf("cannot count agent runs: %w", err) + } + + return count, nil +} + +func (e *AgentRun) Insert( + ctx context.Context, + tx pg.Tx, + scope Scoper, +) error { + q := ` +INSERT INTO agent_runs ( + id, + tenant_id, + organization_id, + start_agent_name, + status, + input_messages, + created_at, + updated_at +) VALUES ( + @id, + @tenant_id, + @organization_id, + @start_agent_name, + @status, + @input_messages, + @created_at, + @updated_at +) +RETURNING + id, + organization_id, + start_agent_name, + status, + checkpoint, + input_messages, + result, + error_message, + stop_requested, + started_at, + lease_owner, + lease_expires_at, + created_at, + updated_at; +` + + args := pgx.StrictNamedArgs{ + "id": e.ID.String(), + "tenant_id": scope.GetTenantID(), + "organization_id": e.OrganizationID.String(), + "start_agent_name": e.StartAgentName, + "status": e.Status, + "input_messages": e.InputMessages, + "created_at": e.CreatedAt, + "updated_at": e.UpdatedAt, + } + + rows, err := tx.Query(ctx, q, args) + if err != nil { + return fmt.Errorf("cannot insert agent run: %w", err) + } + + entity, err := pgx.CollectExactlyOneRow(rows, pgx.RowToStructByName[AgentRun]) + if err != nil { + return fmt.Errorf("cannot insert agent run: %w", err) + } + + *e = entity + + return nil +} + +func (e *AgentRun) Update( + ctx context.Context, + tx pg.Tx, + scope Scoper, +) error { + q := ` +UPDATE agent_runs +SET + status = @status, + result = @result, + error_message = @error_message, + stop_requested = @stop_requested, + started_at = @started_at, + lease_owner = @lease_owner, + lease_expires_at = @lease_expires_at, + updated_at = @updated_at +WHERE + %s + AND id = @id +RETURNING + id, + organization_id, + start_agent_name, + status, + checkpoint, + input_messages, + result, + error_message, + stop_requested, + started_at, + lease_owner, + lease_expires_at, + created_at, + updated_at; +` + + q = fmt.Sprintf(q, scope.SQLFragment()) + + args := pgx.StrictNamedArgs{ + "id": e.ID.String(), + "status": e.Status, + "result": e.Result, + "error_message": e.ErrorMessage, + "stop_requested": e.StopRequested, + "started_at": e.StartedAt, + "lease_owner": e.LeaseOwner, + "lease_expires_at": e.LeaseExpiresAt, + "updated_at": e.UpdatedAt, + } + maps.Copy(args, scope.SQLArguments()) + + rows, err := tx.Query(ctx, q, args) + if err != nil { + return fmt.Errorf("cannot update agent run: %w", err) + } + + entity, err := pgx.CollectExactlyOneRow(rows, pgx.RowToStructByName[AgentRun]) + if err != nil { + return fmt.Errorf("cannot update agent run: %w", err) + } + + *e = entity + + return nil +} + +// ClearCheckpoint is the explicit path for removing persisted checkpoint +// data. AgentRun.Update intentionally does not write checkpoint so status +// commits cannot erase a checkpoint saved by PGCheckpointStore.Save. +func (e *AgentRun) ClearCheckpoint( + ctx context.Context, + tx pg.Tx, + scope Scoper, +) error { + q := ` +UPDATE agent_runs +SET + checkpoint = NULL, + updated_at = now() +WHERE + %s + AND id = @id; +` + + q = fmt.Sprintf(q, scope.SQLFragment()) + + args := pgx.StrictNamedArgs{"id": e.ID.String()} + maps.Copy(args, scope.SQLArguments()) + + _, err := tx.Exec(ctx, q, args) + if err != nil { + return fmt.Errorf("cannot clear agent run checkpoint: %w", err) + } + + e.Checkpoint = nil + + return nil +} + +// SaveCheckpoint writes checkpoint inside the caller's transaction. Use +// this from service methods that must update checkpoint and status +// atomically, such as Approve. +func (e *AgentRun) SaveCheckpoint( + ctx context.Context, + tx pg.Tx, + scope Scoper, + cp *agent.Checkpoint, +) error { + data, err := marshalAgentCheckpoint(cp) + if err != nil { + return err + } + + q := ` +UPDATE agent_runs +SET + checkpoint = @checkpoint, + updated_at = now() +WHERE + %s + AND id = @id +RETURNING + id, + organization_id, + start_agent_name, + status, + checkpoint, + input_messages, + result, + error_message, + stop_requested, + started_at, + lease_owner, + lease_expires_at, + created_at, + updated_at; +` + + q = fmt.Sprintf(q, scope.SQLFragment()) + + args := pgx.StrictNamedArgs{ + "id": e.ID.String(), + "checkpoint": json.RawMessage(data), + } + maps.Copy(args, scope.SQLArguments()) + + rows, err := tx.Query(ctx, q, args) + if err != nil { + return fmt.Errorf("cannot save agent run checkpoint: %w", err) + } + + entity, err := pgx.CollectExactlyOneRow(rows, pgx.RowToStructByName[AgentRun]) + if err != nil { + return fmt.Errorf("cannot save agent run checkpoint: %w", err) + } + + *e = entity + + return nil +} + +func (e *AgentRun) LoadNextPendingForUpdateSkipLocked( + ctx context.Context, + tx pg.Tx, +) error { + q := ` +SELECT + id, + organization_id, + start_agent_name, + status, + checkpoint, + input_messages, + result, + error_message, + stop_requested, + started_at, + lease_owner, + lease_expires_at, + created_at, + updated_at +FROM + agent_runs +WHERE + status = 'PENDING' + AND stop_requested = false +ORDER BY created_at ASC +LIMIT 1 +FOR UPDATE SKIP LOCKED; +` + + rows, err := tx.Query(ctx, q) + if err != nil { + return fmt.Errorf("cannot query pending agent run: %w", err) + } + + entity, err := pgx.CollectExactlyOneRow(rows, pgx.RowToStructByName[AgentRun]) + if err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return ErrResourceNotFound + } + return fmt.Errorf("cannot load pending agent run: %w", err) + } + + *e = entity + + return nil +} + +// ResetStaleAgentRuns resets agent runs whose worker lease has expired. +// The worker refreshes lease_expires_at from a separate heartbeat goroutine, +// so a long LLM or tool call is not considered stale while the process is alive. +// Stale recovery returns rows to PENDING so the supervisor auto-resumes +// from checkpoint when one exists. User-requested suspension remains +// SUSPENDED and is resumed only by the Resume service method. +func ResetStaleAgentRuns(ctx context.Context, conn pg.Querier) error { + q := ` +UPDATE agent_runs +SET + status = 'PENDING', + started_at = NULL, + lease_owner = NULL, + lease_expires_at = NULL, + stop_requested = false, + updated_at = now() +WHERE + status = 'RUNNING' + AND lease_expires_at IS NOT NULL + AND lease_expires_at < now(); +` + + _, err := conn.Exec(ctx, q) + if err != nil { + return fmt.Errorf("cannot reset stale agent runs: %w", err) + } + + return nil +} + +// HeartbeatAgentRunLease refreshes the lease for a running agent run. +// Returns the number of rows affected (0 if lease was lost). +func HeartbeatAgentRunLease( + ctx context.Context, + conn pg.Querier, + runID string, + leaseOwner string, + expiresAt time.Time, +) (int64, error) { + q := ` +UPDATE agent_runs +SET + lease_expires_at = @lease_expires_at, + updated_at = now() +WHERE + id = @id + AND status = 'RUNNING' + AND lease_owner = @lease_owner; +` + + args := pgx.StrictNamedArgs{ + "id": runID, + "lease_owner": leaseOwner, + "lease_expires_at": expiresAt, + } + + tag, err := conn.Exec(ctx, q, args) + if err != nil { + return 0, fmt.Errorf("cannot heartbeat agent run lease: %w", err) + } + + return tag.RowsAffected(), nil +} + +// LoadRunningStopRequestedIDs returns IDs of running agent runs with +// stop_requested = true. +func LoadRunningStopRequestedIDs(ctx context.Context, conn pg.Querier) ([]string, error) { + q := `SELECT id FROM agent_runs WHERE status = 'RUNNING' AND stop_requested = true;` + + rows, err := conn.Query(ctx, q) + if err != nil { + return nil, fmt.Errorf("cannot query stop-requested agent runs: %w", err) + } + defer rows.Close() + + var ids []string + for rows.Next() { + var id string + if err := rows.Scan(&id); err != nil { + return nil, fmt.Errorf("cannot scan stop-requested agent run ID: %w", err) + } + ids = append(ids, id) + } + + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("cannot iterate stop-requested agent runs: %w", err) + } + + return ids, nil +} + +// PGCheckpointStore implements agent.CheckpointStore backed by the +// agent_runs table checkpoint column. It is supervisor-internal and +// intentionally uses raw run IDs with no tenant scope; public service/API +// methods must load AgentRun through scoped coredata methods before invoking +// lifecycle transitions. +type PGCheckpointStore struct { + pg *pg.Client +} + +func NewPGCheckpointStore(pgClient *pg.Client) *PGCheckpointStore { + return &PGCheckpointStore{pg: pgClient} +} + +func (s *PGCheckpointStore) Save(ctx context.Context, runID string, cp *agent.Checkpoint) error { + data, err := marshalAgentCheckpoint(cp) + if err != nil { + return err + } + + return s.pg.WithConn( + ctx, + func(ctx context.Context, conn pg.Querier) error { + q := `UPDATE agent_runs SET checkpoint = @checkpoint, updated_at = now() WHERE id = @id;` + + args := pgx.StrictNamedArgs{ + "id": runID, + "checkpoint": json.RawMessage(data), + } + + tag, err := conn.Exec(ctx, q, args) + if err != nil { + return fmt.Errorf("cannot save checkpoint: %w", err) + } + + if tag.RowsAffected() == 0 { + return fmt.Errorf("cannot save checkpoint: agent run %s not found", runID) + } + + return nil + }, + ) +} + +func (s *PGCheckpointStore) Load(ctx context.Context, runID string) (*agent.Checkpoint, error) { + var cp *agent.Checkpoint + + err := s.pg.WithConn( + ctx, + func(ctx context.Context, conn pg.Querier) error { + q := `SELECT checkpoint FROM agent_runs WHERE id = @id;` + + args := pgx.StrictNamedArgs{"id": runID} + + rows, err := conn.Query(ctx, q, args) + if err != nil { + return fmt.Errorf("cannot query checkpoint: %w", err) + } + + type row struct { + Checkpoint json.RawMessage `db:"checkpoint"` + } + + r, err := pgx.CollectExactlyOneRow(rows, pgx.RowToStructByName[row]) + if err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return ErrResourceNotFound + } + return fmt.Errorf("cannot load checkpoint: %w", err) + } + + if r.Checkpoint == nil { + return nil + } + + cp = new(agent.Checkpoint) + if err := json.Unmarshal(r.Checkpoint, cp); err != nil { + return fmt.Errorf("cannot unmarshal checkpoint: %w", err) + } + + if cp.Version != agent.CheckpointVersion { + return fmt.Errorf("cannot load checkpoint: unsupported checkpoint version %d", cp.Version) + } + + return nil + }, + ) + + return cp, err +} + +func (s *PGCheckpointStore) Delete(ctx context.Context, runID string) error { + return s.pg.WithConn( + ctx, + func(ctx context.Context, conn pg.Querier) error { + q := `UPDATE agent_runs SET checkpoint = NULL, updated_at = now() WHERE id = @id;` + + args := pgx.StrictNamedArgs{"id": runID} + + _, err := conn.Exec(ctx, q, args) + if err != nil { + return fmt.Errorf("cannot delete checkpoint: %w", err) + } + + // Delete is intentionally idempotent: callers use it for cleanup after + // completion, and a concurrently-cleared checkpoint is already the + // desired state. + return nil + }, + ) +} + +func marshalAgentCheckpoint(cp *agent.Checkpoint) ([]byte, error) { + if cp == nil { + return nil, fmt.Errorf("cannot marshal checkpoint: checkpoint is required") + } + + next := *cp + if next.Version == 0 { + next.Version = agent.CheckpointVersion + } + + if next.Version != agent.CheckpointVersion { + return nil, fmt.Errorf("cannot marshal checkpoint: unsupported checkpoint version %d", next.Version) + } + + data, err := json.Marshal(&next) + if err != nil { + return nil, fmt.Errorf("cannot marshal checkpoint: %w", err) + } + + if len(data) > agent.MaxCheckpointBytes { + return nil, fmt.Errorf("cannot marshal checkpoint: size %d exceeds limit %d", len(data), agent.MaxCheckpointBytes) + } + + return data, nil +} diff --git a/pkg/coredata/agent_run_order_field.go b/pkg/coredata/agent_run_order_field.go new file mode 100644 index 000000000..96eaa6d74 --- /dev/null +++ b/pkg/coredata/agent_run_order_field.go @@ -0,0 +1,58 @@ +// Copyright (c) 2026 Probo Inc . +// +// Permission to use, copy, modify, and/or distribute this software for any +// purpose with or without fee is hereby granted, provided that the above +// copyright notice and this permission notice appear in all copies. +// +// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH +// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY +// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, +// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM +// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR +// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR +// PERFORMANCE OF THIS SOFTWARE. + +package coredata + +import "fmt" + +type ( + AgentRunOrderField string +) + +const ( + AgentRunOrderFieldCreatedAt AgentRunOrderField = "CREATED_AT" +) + +func (p AgentRunOrderField) Column() string { + switch p { + case AgentRunOrderFieldCreatedAt: + return "created_at" + } + + panic(fmt.Sprintf("unsupported order by: %s", p)) +} + +func (p AgentRunOrderField) IsValid() bool { + switch p { + case AgentRunOrderFieldCreatedAt: + return true + } + return false +} + +func (p AgentRunOrderField) String() string { + return string(p) +} + +func (p *AgentRunOrderField) UnmarshalText(text []byte) error { + *p = AgentRunOrderField(text) + if !p.IsValid() { + return fmt.Errorf("%s is not a valid AgentRunOrderField", string(text)) + } + return nil +} + +func (p AgentRunOrderField) MarshalText() ([]byte, error) { + return []byte(p.String()), nil +} diff --git a/pkg/coredata/entity_type_reg.go b/pkg/coredata/entity_type_reg.go index 076ea6e59..c2492478b 100644 --- a/pkg/coredata/entity_type_reg.go +++ b/pkg/coredata/entity_type_reg.go @@ -110,6 +110,7 @@ const ( OAuth2DeviceCodeEntityType uint16 = 84 CookieEntityType uint16 = 85 CookieBannerTranslationEntityType uint16 = 86 + AgentRunEntityType uint16 = 87 ) func NewEntityFromID(id gid.GID) (any, bool) { @@ -278,6 +279,8 @@ func NewEntityFromID(id gid.GID) (any, bool) { return &Cookie{ID: id}, true case CookieBannerTranslationEntityType: return &CookieBannerTranslation{ID: id}, true + case AgentRunEntityType: + return &AgentRun{ID: id}, true default: return nil, false } diff --git a/pkg/coredata/migrations/20260424T120000Z.sql b/pkg/coredata/migrations/20260424T120000Z.sql new file mode 100644 index 000000000..f3e6a1c6b --- /dev/null +++ b/pkg/coredata/migrations/20260424T120000Z.sql @@ -0,0 +1,37 @@ +-- Copyright (c) 2026 Probo Inc . +-- +-- Permission to use, copy, modify, and/or distribute this software for any +-- purpose with or without fee is hereby granted, provided that the above +-- copyright notice and this permission notice appear in all copies. +-- +-- THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH +-- REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY +-- AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, +-- INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM +-- LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR +-- OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR +-- PERFORMANCE OF THIS SOFTWARE. + +CREATE TABLE agent_runs ( + id TEXT NOT NULL PRIMARY KEY, + tenant_id TEXT NOT NULL, + organization_id TEXT NOT NULL, + start_agent_name TEXT NOT NULL, + status TEXT NOT NULL DEFAULT 'PENDING', + checkpoint JSONB, + input_messages JSONB NOT NULL, + result JSONB, + error_message TEXT, + stop_requested BOOLEAN NOT NULL DEFAULT FALSE, + started_at TIMESTAMPTZ, + lease_owner TEXT, + lease_expires_at TIMESTAMPTZ, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now() +); + +CREATE INDEX idx_agent_runs_status ON agent_runs (status) + WHERE status IN ('PENDING', 'RUNNING', 'SUSPENDED'); +CREATE INDEX idx_agent_runs_organization_status ON agent_runs (organization_id, status, created_at); +CREATE INDEX idx_agent_runs_running_lease ON agent_runs (lease_expires_at) + WHERE status = 'RUNNING';