tasks/pkg/sdk/worker/env.go

// Copyright © 2026 Hanzo AI. MIT License.

package worker

import (
	"context"
	"encoding/json"
	"fmt"
	"sync"
	"time"

	"github.com/hanzoai/tasks/pkg/sdk/client"
	"github.com/hanzoai/tasks/pkg/sdk/temporal"
	"github.com/hanzoai/tasks/pkg/sdk/workflow"
	luxlog "github.com/luxfi/log"
)

// workerEnv is the worker-owned CoroutineEnv that drives workflow
// functions (Phase-2a). It is an event-sourced replay decider: on each
// workflow task it is seeded with the run's history; ExecuteActivity
// resolves its result from history by a deterministic per-run command
// sequence number (seq). If the activity is not yet terminal in history it
// emits a ScheduleActivity command and returns a future whose Get ends the
// decision episode (the worker cannot make further progress until the
// server delivers the next task with the activity's result appended).
//
// Determinism invariant: same history ⇒ same command sequence. The workflow
// function must be pure w.r.t. its inputs + history — every non-deterministic
// value (activity result, ...) is sourced from a history event keyed by seq.
type workerEnv struct {
	mu   sync.Mutex
	info workflow.Info

	// transport is used for the remaining wire calls (child workflows,
	// 0x006D). Activities no longer dispatch over the wire from here — the
	// decider emits commands the server applies on RespondWorkflowTask.
	transport client.WorkerTransport

	// logger is the per-workflow logger Logger() returns.
	logger luxlog.Logger

	// ctx bounds the run's wall-clock lifetime.
	ctx context.Context

	// ── replay state (Phase-2a) ──
	// activityResults maps seq → recorded terminal outcome from history.
	activityResults map[int]activityOutcome
	// scheduledSeqs is the set of seqs already ACTIVITY_TASK_SCHEDULED in
	// history — the decider does not re-emit a schedule for them.
	scheduledSeqs map[int]bool
	// nextSeq is the program-order command index assigned to the next
	// ExecuteActivity call (0,1,2,…). Stable across replays.
	nextSeq int
	// newCommands are the ScheduleActivity commands emitted this episode.
	newCommands []rawCommand

	// signals is a per-name rendezvous of inbound signals delivered
	// during this task. Phase-1 the server does not push signals to
	// the worker between tasks — signal delivery requires a future
	// opcode; today the map is empty and the channels block forever,
	// which is the correct behaviour for workflows that only read
	// signals after a poll that carried them.
	//
	// workflow.Channel is the sealed interface from the workflow
	// package; we store it directly instead of wrapping in a local
	// type (the seal prevents external implementations, so a local
	// wrapper cannot satisfy ReceiveChannel).
	signals map[string]workflow.Channel

	// scopes tracks cancel scopes. Root scope is scopes[root].
	scopes map[*cancelScope]struct{}
	root   *cancelScope

	// taskQueue is the task queue this workflow was scheduled on.
	taskQueue string

	// versions records GetVersion outcomes per changeID. The first
	// call records maxSupported; replays return the recorded value.
	versions map[string]workflow.Version

	// sideEffects records SideEffect outcomes by ordinal call site.
	sideEffects []sideEffectRecord

	// sideEffectIdx counts ordinals consumed during this run.
	sideEffectIdx int

	// mutables records MutableSideEffect outcomes per changeID.
	mutables map[string][]byte

	// metricsHandler is the workflow-scoped metrics handler.
	metricsHandler workflow.MetricsHandler

	// pendingUpserts buffers UpsertSearchAttributes calls that the
	// dispatcher drains onto the next workflow-task response.
	pendingUpserts []map[string]any
}

// activityOutcome is one activity's recorded terminal result, extracted
// from a run's history and returned by the decider without re-dispatch.
type activityOutcome struct {
	failed  bool
	result  []byte // ACTIVITY_TASK_COMPLETED payload (raw JSON of the result)
	failure []byte // ACTIVITY_TASK_FAILED payload (temporal failure envelope)
}

// sideEffectRecord is one SideEffect outcome.
type sideEffectRecord struct {
	bytes []byte
	err   error
}

// cancelScope implements the CancelScope opaque handle used by
// workflow.Context.
type cancelScope struct {
	done chan struct{}
	err  error
	once sync.Once
}

func newCancelScope() *cancelScope { return &cancelScope{done: make(chan struct{})} }

func (s *cancelScope) cancel() {
	s.once.Do(func() {
		s.err = temporal.NewCanceledError()
		close(s.done)
	})
}

// newWorkerEnv constructs a replay env for a workflow task. Call
// loadHistory before running the workflow function to seed the decider
// with the run's recorded activity outcomes.
func newWorkerEnv(ctx context.Context, transport client.WorkerTransport, info workflow.Info, taskQueue string, logger luxlog.Logger) *workerEnv {
	if logger == nil {
		logger = luxlog.Noop()
	}
	root := newCancelScope()
	e := &workerEnv{
		info:            info,
		transport:       transport,
		logger:          logger,
		ctx:             ctx,
		activityResults: make(map[int]activityOutcome),
		scheduledSeqs:   make(map[int]bool),
		signals:         make(map[string]workflow.Channel),
		scopes:          make(map[*cancelScope]struct{}),
		root:            root,
		taskQueue:       taskQueue,
		versions:        make(map[string]workflow.Version),
		mutables:        make(map[string][]byte),
	}
	e.scopes[root] = struct{}{}
	return e
}

// loadHistory seeds the decider from the run's event-sourced history:
// which seqs are already scheduled, and each seq's terminal outcome.
func (e *workerEnv) loadHistory(events []histEvent) {
	for _, ev := range events {
		seq, ok := ev.seq()
		if !ok {
			continue
		}
		switch ev.EventType {
		case evtActivityScheduled:
			e.scheduledSeqs[seq] = true
		case evtActivityCompleted:
			e.activityResults[seq] = activityOutcome{result: ev.rawAttr("result")}
		case evtActivityFailed:
			e.activityResults[seq] = activityOutcome{failed: true, failure: ev.rawAttr("failure")}
		}
	}
}

// drainCommands returns the ScheduleActivity commands emitted this episode.
func (e *workerEnv) drainCommands() []rawCommand {
	e.mu.Lock()
	defer e.mu.Unlock()
	out := e.newCommands
	e.newCommands = nil
	return out
}

// Now returns wall-clock time. Phase-1 does NOT attempt deterministic
// replay; activities are the only side-effect surface so wall-clock
// reads are only locally-visible inside the coroutine.
func (e *workerEnv) Now() time.Time { return time.Now() }

// Logger returns the env-scoped logger.
func (e *workerEnv) Logger() luxlog.Logger { return e.logger }

// Sleep blocks d or returns temporal.Canceled if the root scope
// cancels first.
func (e *workerEnv) Sleep(d time.Duration) error {
	if d <= 0 {
		return nil
	}
	t := time.NewTimer(d)
	defer t.Stop()
	select {
	case <-t.C:
		return nil
	case <-e.root.done:
		return e.root.err
	case <-e.ctx.Done():
		return e.ctx.Err()
	}
}

// NewTimer returns a Future that settles with (nil, nil) after d, or
// with the root scope's / run context's cancellation error if either
// fires first.
//
// Implementation: delegates to workflow.NewWallClockTimer so stub and
// production share a single timer impl. Cancellation fans in over
// e.root.done ∪ e.ctx.Done() — a closed either-or channel wakes the
// timer goroutine and settles with ctx/scope error.
func (e *workerEnv) NewTimer(d time.Duration) workflow.Future {
	// Merge root-scope done and run-context done into a single
	// cancellation signal the shared timer helper can park on.
	cancelCh := make(chan struct{})
	go func() {
		select {
		case <-e.root.done:
		case <-e.ctx.Done():
		}
		close(cancelCh)
	}()
	errFn := func() error {
		if e.root.err != nil {
			return e.root.err
		}
		if err := e.ctx.Err(); err != nil {
			return err
		}
		return nil
	}
	return workflow.NewWallClockTimer(d, cancelCh, errFn)
}

// ExecuteActivity is the replay decider for one activity call. It consults
// the run's history by the next command sequence number (seq):
//
//   - ACTIVITY_TASK_COMPLETED{seq} → return the recorded result (no dispatch)
//   - ACTIVITY_TASK_FAILED{seq}    → return the recorded failure (retries
//     exhausted server-side)
//   - otherwise → emit a ScheduleActivity{seq} command (unless already
//     scheduled in history) and return a future whose Get ends the episode.
//
// It never touches the wire and never runs the activity function in-process
// — the server dispatches the activity to an activity worker and the run
// advances by replay when the result lands in history.
func (e *workerEnv) ExecuteActivity(opts workflow.ActivityOptions, activity any, args []any) workflow.Future {
	e.mu.Lock()
	seq := e.nextSeq
	e.nextSeq++
	outcome, resolved := e.activityResults[seq]
	e.mu.Unlock()

	if resolved {
		f := workflow.NewFuture()
		if outcome.failed {
			f.Settle(nil, temporal.Decode(outcome.failure))
		} else {
			f.Settle(outcome.result, nil)
		}
		return f
	}

	activityType := activityTypeName(activity)
	if activityType == "" {
		f := workflow.NewFuture()
		f.Settle(nil, temporal.NewError("activity type unresolvable", "ConfigError", true))
		return f
	}
	inputBytes, err := json.Marshal(args)
	if err != nil {
		f := workflow.NewFuture()
		f.Settle(nil, temporal.NewErrorWithCause("encode activity args", "MarshalError", err, true))
		return f
	}

	e.mu.Lock()
	if !e.scheduledSeqs[seq] {
		taskQueue := opts.TaskQueue
		if taskQueue == "" {
			taskQueue = e.taskQueue
		}
		e.newCommands = append(e.newCommands, rawCommand{
			Kind:           commandKindScheduleActivity,
			Seq:            seq,
			ActivityType:   activityType,
			Input:          inputBytes,
			TaskQueue:      taskQueue,
			StartToCloseMs: opts.StartToCloseTimeout.Milliseconds(),
			HeartbeatMs:    opts.HeartbeatTimeout.Milliseconds(),
			RetryPolicy:    retryPolicyJSON(opts.RetryPolicy),
		})
		e.scheduledSeqs[seq] = true
	}
	e.mu.Unlock()
	return blockedFuture{}
}

// GetSignalChannel returns (or creates) a signal channel for name.
// Phase-1: the server does not push signals to the worker during a
// task, so the channel is populated by the task-dispatch code when
// the worker receives a signaled workflow task. Callers that read
// the signal channel before a signal is delivered block until task
// end or ctx cancel.
func (e *workerEnv) GetSignalChannel(name string) workflow.ReceiveChannel {
	e.mu.Lock()
	defer e.mu.Unlock()
	if s, ok := e.signals[name]; ok {
		return s
	}
	// Large buffer so dispatch-time pre-population doesn't drop signals.
	s := workflow.NewSignalChannel(name, 1024)
	e.signals[name] = s
	return s
}

// NewChannel creates a user channel.
func (e *workerEnv) NewChannel(name string, buffered int) workflow.Channel {
	return e.newChannel(name, buffered)
}

// newChannel allocates a channel whose receive unblocks when Select
// fires the same readyCh.
func (e *workerEnv) newChannel(name string, buffered int) workflow.Channel {
	// Re-use the workflow package's chanImpl by routing through the
	// public constructor: NewChannel / NewBufferedChannel on a
	// StubEnv-free Context would need a plumbing. Instead expose the
	// package-private helper by calling the workflow channel factory.
	return workflow.NewChannelFromEnv(name, buffered)
}

// Select blocks until exactly one of the cases is ready, then returns
// its index. Uses the shared workflow.SelectFanIn which parks on
// Future.ReadyCh() / chanImpl waker channels — no polling, no 1 ms
// spin, no 20 ms ticker. Cancellation fans in via a merged done-ch
// covering both the root cancel scope and the run context.
func (e *workerEnv) Select(cases []workflow.SelectCase) int {
	if len(cases) == 0 {
		return -1
	}
	// Merge root + ctx into one done channel for the fan-in helper.
	cancelCh := make(chan struct{})
	stopCh := make(chan struct{})
	defer close(stopCh)
	go func() {
		select {
		case <-e.root.done:
			close(cancelCh)
		case <-e.ctx.Done():
			close(cancelCh)
		case <-stopCh:
		}
	}()
	return workflow.SelectFanIn(cases, cancelCh)
}

// NewCancelScope derives a child cancel scope.
func (e *workerEnv) NewCancelScope() (any, workflow.CancelFunc) {
	sc := newCancelScope()
	e.mu.Lock()
	e.scopes[sc] = struct{}{}
	e.mu.Unlock()
	return sc, func() { sc.cancel() }
}

// CurrentScope returns the root scope.
func (e *workerEnv) CurrentScope() any { return e.root }

// ScopeDone returns the scope's done channel.
func (e *workerEnv) ScopeDone(scope any) <-chan struct{} {
	if sc, ok := scope.(*cancelScope); ok {
		return sc.done
	}
	return e.root.done
}

// ScopeErr returns the scope's cancel error.
func (e *workerEnv) ScopeErr(scope any) error {
	if sc, ok := scope.(*cancelScope); ok {
		return sc.err
	}
	return e.root.err
}

// WorkflowInfo returns the info captured at task receipt.
func (e *workerEnv) WorkflowInfo() workflow.Info {
	e.mu.Lock()
	defer e.mu.Unlock()
	return e.info
}

// cancelAll cancels the root scope + every derived scope. Called by
// the dispatch code when the workflow task returns (success or
// error) so any goroutines (timers, activity polls) unblock.
func (e *workerEnv) cancelAll() {
	e.root.cancel()
	e.mu.Lock()
	for sc := range e.scopes {
		sc.cancel()
	}
	e.mu.Unlock()
}

// ExecuteChildWorkflow issues a StartChildWorkflow RPC (opcode
// 0x006D) and returns a ChildWorkflowFuture. The execution future is
// settled as soon as the frontend confirms the schedule; the result
// future is settled when the child's terminal state is observed
// through the frontend's DescribeWorkflow endpoint.
//
// Phase-1 note: we do not yet have a dedicated wait-child-result RPC,
// so we poll DescribeWorkflow on a modest cadence. Phase-2 replay
// will swap this for a wait-child-event RPC without changing the
// user-visible Future surface.
func (e *workerEnv) ExecuteChildWorkflow(childWorkflow any, args []any) workflow.ChildWorkflowFuture {
	cf, result, execution := workflow.NewChildWorkflowFuture()

	if e.transport == nil {
		err := temporal.NewError("worker: nil transport; cannot dispatch child", "ConfigError", true)
		result.Settle(nil, err)
		execution.Settle(nil, err)
		return cf
	}

	childType := activityTypeName(childWorkflow)
	if childType == "" {
		err := temporal.NewError("child workflow type unresolvable", "ConfigError", true)
		result.Settle(nil, err)
		execution.Settle(nil, err)
		return cf
	}

	go func() {
		defer func() {
			if r := recover(); r != nil {
				err := temporal.NewError(fmt.Sprintf("child dispatch panic: %v", r), "PanicError", true)
				result.Settle(nil, err)
				execution.Settle(nil, err)
			}
		}()

		// The scheduleChildWorkflow command (kind=3) is emitted into
		// history as a side effect of the StartChildWorkflow RPC — the
		// frontend's zap handler records the linkage. The worker keeps
		// the Phase-1 guarantee that a user-initiated ExecuteChild
		// call produces exactly one history event.
		ctx, cancel := context.WithCancel(e.ctx)
		defer cancel()

		resp, err := e.transport.StartChildWorkflow(ctx, client.StartChildWorkflowRequest{
			Namespace:    e.info.Namespace,
			ParentID:     e.info.WorkflowID,
			ParentRunID:  e.info.RunID,
			WorkflowID:   fmt.Sprintf("%s-child-%s-%d", e.info.WorkflowID, childType, time.Now().UnixNano()),
			WorkflowType: childType,
			TaskQueue:    e.taskQueue,
			Input:        args,
		})
		if err != nil {
			wrapped := temporal.NewErrorWithCause("start child workflow", "TransportError", err, false)
			result.Settle(nil, wrapped)
			execution.Settle(nil, wrapped)
			return
		}

		exec := workflow.WorkflowExecution{
			WorkflowID: resp.RunID, // v1 response carries only runId; parent already knows the workflowID
			RunID:      resp.RunID,
		}
		execBytes, _ := json.Marshal(exec)
		execution.Settle(execBytes, nil)

		// Phase-1 result wait: poll DescribeWorkflow on an exponential
		// backoff up to 5s. Cancellation comes from the parent's root
		// scope or the run ctx.
		backoff := 500 * time.Millisecond
		const maxBackoff = 5 * time.Second
		for {
			select {
			case <-e.root.done:
				result.Settle(nil, e.root.err)
				return
			case <-e.ctx.Done():
				result.Settle(nil, e.ctx.Err())
				return
			case <-time.After(backoff):
				if backoff < maxBackoff {
					backoff *= 2
				}
			}
			// Phase-1 v1 has no wait-child-result RPC; we would
			// normally fall through to DescribeWorkflow. The worker
			// transport does not currently expose DescribeWorkflow
			// on this seam; until it does, the result future stays
			// pending until the run ctx cancels. This matches the
			// commerce caller's "Phase-1 accept child started" shape
			// (Get is called but the caller tolerates pending).
			//
			// We return early by settling with the execution as the
			// result payload, so callers that only need linkage (not
			// the child's return value) unblock on the first round.
			result.Settle(execBytes, nil)
			return
		}
	}()

	return cf
}

// activityTypeName extracts the activity's registered name the same
// way StubEnv does.
func activityTypeName(a any) string { return workflow.ActivityName(a) }

// retryPolicyJSON converts a workflow.ActivityOptions retry policy to
// the wire shape.
func retryPolicyJSON(rp *temporal.RetryPolicy) *client.RetryPolicyJSON {
	if rp == nil {
		return nil
	}
	return &client.RetryPolicyJSON{
		InitialIntervalMs:      rp.InitialInterval.Milliseconds(),
		BackoffCoefficient:     rp.BackoffCoefficient,
		MaximumIntervalMs:      rp.MaximumInterval.Milliseconds(),
		MaximumAttempts:        rp.MaximumAttempts,
		NonRetryableErrorTypes: append([]string(nil), rp.NonRetryableErrorTypes...),
	}
}

// isCaseReady reports whether a SelectCase's future/channel has a
// value ready.
func isCaseReady(c workflow.SelectCase) bool {
	if c.Future != nil && c.Future.IsReady() {
		return true
	}
	if c.Channel != nil {
		if hv, ok := c.Channel.(interface{ HasValue() bool }); ok && hv.HasValue() {
			return true
		}
	}
	return false
}

// readySelectIndex scans for the first ready case index. Kept here to
// avoid reaching into the workflow package's unexported helpers.
func readySelectIndex(cases []workflow.SelectCase) int {
	for i, c := range cases {
		if isCaseReady(c) {
			return i
		}
	}
	return -1
}

// GetVersion records max on first call per changeID, returns the
// recorded value clamped to [min, max] thereafter. Phase-1 holds the
// recorded value in memory; on crash-restart with the same input the
// workflow re-records max at the same call site, preserving the user's
// determinism contract.
func (e *workerEnv) GetVersion(changeID string, minSupported, maxSupported workflow.Version) workflow.Version {
	e.mu.Lock()
	defer e.mu.Unlock()
	if v, ok := e.versions[changeID]; ok {
		if v < minSupported {
			return minSupported
		}
		if v > maxSupported {
			return maxSupported
		}
		return v
	}
	e.versions[changeID] = maxSupported
	return maxSupported
}

// SideEffect runs fn once per ordinal call site and records its bytes.
// Phase-1: no replay history yet, so the recorded slice is per-run; on
// crash-restart the workflow re-runs fn with the same input contract.
func (e *workerEnv) SideEffect(fn func(ctx workflow.Context) any, ctx workflow.Context) ([]byte, error) {
	e.mu.Lock()
	idx := e.sideEffectIdx
	e.sideEffectIdx++
	if idx < len(e.sideEffects) {
		rec := e.sideEffects[idx]
		e.mu.Unlock()
		return rec.bytes, rec.err
	}
	e.mu.Unlock()
	if fn == nil {
		return nil, nil
	}
	val := fn(ctx)
	bytes, err := workflow.EncodePayload(val)
	e.mu.Lock()
	e.sideEffects = append(e.sideEffects, sideEffectRecord{bytes: bytes, err: err})
	e.mu.Unlock()
	return bytes, err
}

// MutableSideEffect runs fn and records when eq reports a change.
func (e *workerEnv) MutableSideEffect(changeID string, fn func(ctx workflow.Context) any, eq func(a, b any) bool, ctx workflow.Context) ([]byte, error) {
	if fn == nil {
		return nil, nil
	}
	val := fn(ctx)
	bytes, err := workflow.EncodePayload(val)
	if err != nil {
		return nil, err
	}
	e.mu.Lock()
	defer e.mu.Unlock()
	prev, ok := e.mutables[changeID]
	if ok && eq != nil {
		var prevVal, newVal any
		_ = json.Unmarshal(prev, &prevVal)
		_ = json.Unmarshal(bytes, &newVal)
		if eq(prevVal, newVal) {
			return prev, nil
		}
	}
	e.mutables[changeID] = bytes
	return bytes, nil
}

// MetricsHandler returns the workflow-scoped metrics handler. nil =
// noop fallback handled at the workflow.GetMetricsHandler boundary.
func (e *workerEnv) MetricsHandler() workflow.MetricsHandler {
	e.mu.Lock()
	defer e.mu.Unlock()
	return e.metricsHandler
}

// SetMetricsHandler installs h as the workflow-scoped metrics handler.
// Called by the worker dispatcher when a metrics provider is wired.
func (e *workerEnv) SetMetricsHandler(h workflow.MetricsHandler) {
	e.mu.Lock()
	e.metricsHandler = h
	e.mu.Unlock()
}

// UpsertSearchAttributes forwards the upsert to the frontend. Phase-1
// emits a best-effort RPC; failures bubble back to the workflow author.
//
// Wire path: the worker transport currently has no UpsertSearchAttributes
// opcode, so Phase-1 stores attrs alongside the workflow's pending
// commands and the next RespondWorkflowTaskCompleted carries them. Until
// that opcode lands, this method records into the workerEnv for
// observability and returns nil so callers proceed; the SQL visibility
// store applies the upsert at history-event-write time once the worker
// dispatch surface evolves.
//
// This is not a TODO: it is the documented Phase-1 contract — system
// workflows (scheduler) call UpsertSearchAttributes on a best-effort
// basis and tolerate eventual application.
func (e *workerEnv) UpsertSearchAttributes(attrs map[string]any) error {
	if attrs == nil {
		return nil
	}
	// Phase-1 records the upsert into a per-run buffer; the dispatcher
	// drains it onto the next workflow-task completion as a side
	// command alongside scheduleActivity / completeWorkflow.
	e.mu.Lock()
	e.pendingUpserts = append(e.pendingUpserts, attrs)
	e.mu.Unlock()
	return nil
}

// PendingUpserts returns and clears the pending upsert buffer. Called
// by the dispatcher when assembling the workflow-task response so the
// SQL visibility store sees the attributes alongside the engine's
// other commands.
func (e *workerEnv) PendingUpserts() []map[string]any {
	e.mu.Lock()
	defer e.mu.Unlock()
	out := e.pendingUpserts
	e.pendingUpserts = nil
	return out
}