poimen-workflows/internal/memory/worker_setup.go

package memory

import (
	"context"
	"fmt"
	"time"

	"go.temporal.io/sdk/activity"
	"go.temporal.io/sdk/temporal"
	"go.temporal.io/sdk/worker"
	"go.temporal.io/sdk/workflow"
)

// RegisterMemoryActivities registers all memory service activities with worker
func RegisterMemoryActivities(w worker.Worker, service *Service) {
	activities := NewActivities(service)

	// Register activities (activity name = "ActivityName" → "activityName")
	w.RegisterActivity(activities.CreateKnowledgeActivity)
	w.RegisterActivity(activities.UpdateKnowledgeActivity)
	w.RegisterActivity(activities.SearchKnowledgeActivity)
	w.RegisterActivity(activities.GetContextActivity)
	w.RegisterActivity(activities.GetVaultActivity)
	w.RegisterActivity(activities.HealthCheckActivity)
	w.RegisterActivity(activities.LearnFromExecutionActivity)
	w.RegisterActivity(activities.DiagnoseIssueActivity)
	w.RegisterActivity(activities.AnalyzeErrorActivity)
	w.RegisterActivity(activities.DocumentDecisionActivity)
	w.RegisterActivity(activities.SearchAndApplyActivity)
	w.RegisterActivity(activities.RefreshMemoryActivity)
}

// ActivityOptions memory service activity options
type ActivityOptions struct {
	RetryAttempts int
	RetryBackoff  time.Duration
	StartTimeout  time.Duration
	HeartbeatRate time.Duration
}

// DefaultActivityOptions returns sensible defaults
func DefaultActivityOptions() *ActivityOptions {
	return &ActivityOptions{
		RetryAttempts: 3,
		RetryBackoff:  time.Second,
		StartTimeout:  30 * time.Second,
		HeartbeatRate: 10 * time.Second,
	}
}

// ExecuteCreateKnowledge wrapper for CreateKnowledgeActivity
func ExecuteCreateKnowledge(
	ctx workflow.Context,
	record *KnowledgeRecord,
	opts *ActivityOptions,
) (string, error) {
	if opts == nil {
		opts = DefaultActivityOptions()
	}

	activityCtx := workflow.WithActivityOptions(ctx, workflow.ActivityOptions{
		ScheduleToCloseTimeout: 2 * time.Minute,
		StartToCloseTimeout:    time.Minute,
		RetryPolicy: &temporal.RetryPolicy{
			InitialInterval:        opts.RetryBackoff,
			BackoffCoefficient:     2.0,
			MaximumInterval:        30 * time.Second,
			MaximumAttempts:        int32(opts.RetryAttempts),
			NonRetryableErrorTypes: []string{},
		},
	})

	var result string
	err := workflow.ExecuteActivity(activityCtx, "CreateKnowledgeActivity", record).Get(activityCtx, &result)
	return result, err
}

// ExecuteSearchKnowledge wrapper for SearchKnowledgeActivity
func ExecuteSearchKnowledge(
	ctx workflow.Context,
	query string,
	opts *RetrievalOptions,
	activityOpts *ActivityOptions,
) ([]KnowledgeRecord, error) {
	if activityOpts == nil {
		activityOpts = DefaultActivityOptions()
	}

	activityCtx := workflow.WithActivityOptions(ctx, workflow.ActivityOptions{
		ScheduleToCloseTimeout: 3 * time.Minute,
		StartToCloseTimeout:    2 * time.Minute,
		RetryPolicy: &temporal.RetryPolicy{
			InitialInterval:    activityOpts.RetryBackoff,
			BackoffCoefficient: 2.0,
			MaximumInterval:    30 * time.Second,
			MaximumAttempts:    int32(activityOpts.RetryAttempts),
		},
	})

	var result []KnowledgeRecord
	err := workflow.ExecuteActivity(activityCtx, "SearchKnowledgeActivity", query, opts).Get(activityCtx, &result)
	return result, err
}

// ExecuteGetContext wrapper for GetContextActivity
func ExecuteGetContext(
	ctx workflow.Context,
	tool, task string,
	budget int,
	opts *ActivityOptions,
) (*ServiceContext, error) {
	if opts == nil {
		opts = DefaultActivityOptions()
	}

	activityCtx := workflow.WithActivityOptions(ctx, workflow.ActivityOptions{
		ScheduleToCloseTimeout: 3 * time.Minute,
		StartToCloseTimeout:    2 * time.Minute,
		RetryPolicy: &temporal.RetryPolicy{
			InitialInterval:    opts.RetryBackoff,
			BackoffCoefficient: 2.0,
			MaximumInterval:    30 * time.Second,
			MaximumAttempts:    int32(opts.RetryAttempts),
		},
	})

	var result *ServiceContext
	err := workflow.ExecuteActivity(activityCtx, "GetContextActivity", tool, task, budget).Get(activityCtx, &result)
	return result, err
}

// ExecuteDiagnoseIssue wrapper for DiagnoseIssueActivity
func ExecuteDiagnoseIssue(
	ctx workflow.Context,
	tool, issue string,
	opts *ActivityOptions,
) ([]string, error) {
	if opts == nil {
		opts = DefaultActivityOptions()
	}

	activityCtx := workflow.WithActivityOptions(ctx, workflow.ActivityOptions{
		ScheduleToCloseTimeout: 2 * time.Minute,
		StartToCloseTimeout:    time.Minute,
		RetryPolicy: &temporal.RetryPolicy{
			InitialInterval:    opts.RetryBackoff,
			BackoffCoefficient: 2.0,
			MaximumInterval:    30 * time.Second,
			MaximumAttempts:    int32(opts.RetryAttempts),
		},
	})

	var result []string
	err := workflow.ExecuteActivity(activityCtx, "DiagnoseIssueActivity", tool, issue).Get(activityCtx, &result)
	return result, err
}

// ExecuteAnalyzeError wrapper for AnalyzeErrorActivity
func ExecuteAnalyzeError(
	ctx workflow.Context,
	errorMsg string,
	opts *ActivityOptions,
) ([]KnowledgeRecord, error) {
	if opts == nil {
		opts = DefaultActivityOptions()
	}

	activityCtx := workflow.WithActivityOptions(ctx, workflow.ActivityOptions{
		ScheduleToCloseTimeout: 2 * time.Minute,
		StartToCloseTimeout:    time.Minute,
		RetryPolicy: &temporal.RetryPolicy{
			InitialInterval:    opts.RetryBackoff,
			BackoffCoefficient: 2.0,
			MaximumInterval:    30 * time.Second,
			MaximumAttempts:    int32(opts.RetryAttempts),
		},
	})

	var result []KnowledgeRecord
	err := workflow.ExecuteActivity(activityCtx, "AnalyzeErrorActivity", errorMsg).Get(activityCtx, &result)
	return result, err
}

// ExecuteHealthCheck wrapper for HealthCheckActivity
func ExecuteHealthCheck(
	ctx workflow.Context,
	opts *ActivityOptions,
) (bool, error) {
	if opts == nil {
		opts = DefaultActivityOptions()
	}

	activityCtx := workflow.WithActivityOptions(ctx, workflow.ActivityOptions{
		ScheduleToCloseTimeout: 1 * time.Minute,
		StartToCloseTimeout:    30 * time.Second,
		RetryPolicy: &temporal.RetryPolicy{
			InitialInterval:    opts.RetryBackoff,
			BackoffCoefficient: 2.0,
			MaximumInterval:    15 * time.Second,
			MaximumAttempts:    int32(opts.RetryAttempts),
		},
	})

	var result bool
	err := workflow.ExecuteActivity(activityCtx, "HealthCheckActivity").Get(activityCtx, &result)
	return result, err
}

// ExecuteLearnFromExecution wrapper for LearnFromExecutionActivity
func ExecuteLearnFromExecution(
	ctx workflow.Context,
	taskID, result string,
	tags []string,
	opts *ActivityOptions,
) (string, error) {
	if opts == nil {
		opts = DefaultActivityOptions()
	}

	activityCtx := workflow.WithActivityOptions(ctx, workflow.ActivityOptions{
		ScheduleToCloseTimeout: 2 * time.Minute,
		StartToCloseTimeout:    time.Minute,
		RetryPolicy: &temporal.RetryPolicy{
			InitialInterval:    opts.RetryBackoff,
			BackoffCoefficient: 2.0,
			MaximumInterval:    30 * time.Second,
			MaximumAttempts:    int32(opts.RetryAttempts),
		},
	})

	var recordID string
	err := workflow.ExecuteActivity(activityCtx, "LearnFromExecutionActivity", taskID, result, tags).Get(activityCtx, &recordID)
	return recordID, err
}

// ExecuteDocumentDecision wrapper for DocumentDecisionActivity
func ExecuteDocumentDecision(
	ctx workflow.Context,
	decisionType, decision, reasoning string,
	opts *ActivityOptions,
) (string, error) {
	if opts == nil {
		opts = DefaultActivityOptions()
	}

	activityCtx := workflow.WithActivityOptions(ctx, workflow.ActivityOptions{
		ScheduleToCloseTimeout: 2 * time.Minute,
		StartToCloseTimeout:    time.Minute,
		RetryPolicy: &temporal.RetryPolicy{
			InitialInterval:    opts.RetryBackoff,
			BackoffCoefficient: 2.0,
			MaximumInterval:    30 * time.Second,
			MaximumAttempts:    int32(opts.RetryAttempts),
		},
	})

	var recordID string
	err := workflow.ExecuteActivity(activityCtx, "DocumentDecisionActivity", decisionType, decision, reasoning).Get(activityCtx, &recordID)
	return recordID, err
}

// ExecuteRefreshMemory wrapper for RefreshMemoryActivity
func ExecuteRefreshMemory(
	ctx workflow.Context,
	opts *ActivityOptions,
) (map[string]interface{}, error) {
	if opts == nil {
		opts = DefaultActivityOptions()
	}

	activityCtx := workflow.WithActivityOptions(ctx, workflow.ActivityOptions{
		ScheduleToCloseTimeout: 2 * time.Minute,
		StartToCloseTimeout:    time.Minute,
		RetryPolicy: &temporal.RetryPolicy{
			InitialInterval:    opts.RetryBackoff,
			BackoffCoefficient: 2.0,
			MaximumInterval:    30 * time.Second,
			MaximumAttempts:    int32(opts.RetryAttempts),
		},
	})

	var result map[string]interface{}
	err := workflow.ExecuteActivity(activityCtx, "RefreshMemoryActivity").Get(activityCtx, &result)
	return result, err
}

// HeartbeatMemoryActivity sends heartbeat every N seconds
// Usage: Long-running memory operations
func HeartbeatMemoryActivity(ctx context.Context, maxDuration time.Duration) error {
	ticker := time.NewTicker(10 * time.Second)
	defer ticker.Stop()

	deadline := time.Now().Add(maxDuration)

	for {
		select {
		case <-ticker.C:
			activity.RecordHeartbeat(ctx, time.Now())
		case <-ctx.Done():
			return ctx.Err()
		}

		if time.Now().After(deadline) {
			break
		}
	}

	return nil
}

// MemoryActivityError wraps errors with activity context
type MemoryActivityError struct {
	ActivityName string
	Attempt      int
	Err          error
}

func (e *MemoryActivityError) Error() string {
	return fmt.Sprintf("memory activity %s (attempt %d): %v", e.ActivityName, e.Attempt, e.Err)
}

// CaptureActivityError captures activity execution errors
func CaptureActivityError(activityName string, err error) error {
	if err == nil {
		return nil
	}

	return &MemoryActivityError{
		ActivityName: activityName,
		Attempt:      1,
		Err:          err,
	}
}