WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

internal/application/service/tenant_skill_service.go

@@ -0,0 +1,227 @@
+package service
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"sync"
+	"time"
+
+	"github.com/redis/go-redis/v9"
+	"github.com/robfig/cron/v3"
+	"golang.org/x/sync/singleflight"
+
+	"github.com/Tencent/WeKnora/internal/application/repository"
+	"github.com/Tencent/WeKnora/internal/common/redislock"
+	"github.com/Tencent/WeKnora/internal/sandbox"
+	"github.com/Tencent/WeKnora/internal/types/interfaces"
+)
+
+// skillImageLockLease bounds how long one install/remove may hold the config
+// lock without renewing. Installs run for minutes, so the lease is renewed by
+// redislock rather than being set long.
+const (
+	skillImageLockLease = 30 * time.Second
+	skillImageLockRenew = 10 * time.Second
+
+	// skillInstallStuckTTL is how long a run may go without a heartbeat
+	// before the reaper treats it as abandoned. It is a silence budget, not
+	// a duration budget: a legitimate install that spends two hours in the
+	// agent keeps beating and is left alone.
+	skillInstallStuckTTL = 60 * time.Minute
+
+	// skillInstallHeartbeatInterval is how often a running install stamps
+	// InstallingSince to say its process is still alive. Everything that has
+	// to tell "still working" from "died" reads that timestamp.
+	skillInstallHeartbeatInterval = 30 * time.Second
+
+	// skillInstallInFlightSkip is how much heartbeat silence makes a second
+	// upload of the same archive stop deferring to the run that owns the row.
+	// It is a multiple of the heartbeat so a slow install is never mistaken
+	// for a dead one, and short enough that a re-upload recovers a dead
+	// process in minutes instead of waiting for skillInstallStuckTTL.
+	skillInstallInFlightSkip = 3 * time.Minute
+
+	// skillSnapshotRetention is how long a superseded snapshot stays on the
+	// provider after the pointer has moved before prune even tries. Paused
+	// session sandboxes can pin the template past this window; prune retries
+	// on Conflict. A config that sets a longer sandbox TTL extends this via
+	// snapshotRetentionFor.
+	skillSnapshotRetention = 24 * time.Hour
+
+	// skillSnapshotTTLMargin is added on top of a config's own sandbox TTL
+	// so an in-flight create that resolved the previous pointer still has
+	// a template to boot from.
+	skillSnapshotTTLMargin = time.Hour
+)
+
+// TenantSkillService owns the skill image lifecycle for sandbox configs.
+type TenantSkillService struct {
+	skills        repository.TenantSkillRepository
+	configs       repository.TenantSandboxConfigRepository
+	resolver      interfaces.StorageBackendResolver
+	sandboxes     sandbox.TenantSandboxResolver
+	sandboxPolicy WorkspaceSandboxPolicy
+	agents        interfaces.AgentService
+	// installerAgents reads the stored installer record. It is a separate
+	// dependency from agents because GetAgentByID lives on the custom agent
+	// service, not on interfaces.AgentService.
+	installerAgents installerAgentSource
+	sessions        interfaces.SessionService
+	models          interfaces.ModelService
+	redis           *redis.Client
+
+	// streams and messages are the two halves of an install transcript: the
+	// replayable event log the console tails, and the durable rows it falls
+	// back to once the log's TTL has passed.
+	streams  interfaces.StreamManager
+	messages interfaces.MessageRepository
+
+	now func() time.Time
+
+	// sourceHTTP pulls remote skill archives. Nil means the package SSRF-safe
+	// default; tests inject httptest clients.
+	sourceHTTP *http.Client
+
+	// cleanupTimeout bounds one piece of compensating work. Injectable so a
+	// test can let an install outlast it, which every real install does.
+	cleanupTimeout time.Duration
+
+	// snapshotRetention is how long a superseded snapshot is kept on the
+	// provider. Injectable so a prune test can age a row without waiting a day.
+	snapshotRetention time.Duration
+
+	// installHeartbeat is how often a running install restamps its liveness.
+	// Injectable so a test can observe a beat without waiting half a minute.
+	installHeartbeat time.Duration
+
+	// localLocks serialises installs when Redis is absent. It only guards this
+	// process; multi-replica deployments require Redis for cross-process safety.
+	localLocks *keyedMutex
+
+	// bundleCache keeps recently downloaded skill zips so the admin file
+	// browser (list + N reads) does not hit object storage on every click.
+	bundleCache *skillBundleArchiveCache
+	bundleLoad  singleflight.Group
+
+	cron    *cron.Cron
+	cronMu  sync.Mutex
+	started bool
+
+	// runCancels lets StopSkill abort the goroutine that holds this skill's
+	// install. The map is per-process: a restart or another replica has
+	// nothing to cancel, and StopSkill then only rewrites the stuck row.
+	runCancelMu sync.Mutex
+	runCancels  map[string]*skillRunCancel
+}
+
+// NewTenantSkillService wires the repositories and runtimes the install and
+// remove flows share. Redis may be nil; the local lock then serialises one
+// process only.
+func NewTenantSkillService(
+	skillsRepo repository.TenantSkillRepository,
+	configsRepo repository.TenantSandboxConfigRepository,
+	resolver interfaces.StorageBackendResolver,
+	sandboxes sandbox.TenantSandboxResolver,
+	sandboxPolicy WorkspaceSandboxPolicy,
+	agents interfaces.AgentService,
+	customAgents interfaces.CustomAgentService,
+	sessions interfaces.SessionService,
+	models interfaces.ModelService,
+	redisClient *redis.Client,
+	streams interfaces.StreamManager,
+	messages interfaces.MessageRepository,
+) *TenantSkillService {
+	return &TenantSkillService{
+		skills:            skillsRepo,
+		configs:           configsRepo,
+		resolver:          resolver,
+		sandboxes:         sandboxes,
+		sandboxPolicy:     sandboxPolicy,
+		agents:            agents,
+		installerAgents:   customAgents,
+		sessions:          sessions,
+		models:            models,
+		redis:             redisClient,
+		streams:           streams,
+		messages:          messages,
+		now:               time.Now,
+		cleanupTimeout:    installCleanupTimeout,
+		snapshotRetention: skillSnapshotRetention,
+		installHeartbeat:  skillInstallHeartbeatInterval,
+		localLocks:        newKeyedMutex(),
+		runCancels:        map[string]*skillRunCancel{},
+		bundleCache:       newSkillBundleArchiveCache(),
+		cron: cron.New(cron.WithSeconds(), cron.WithChain(
+			cron.Recover(cron.DefaultLogger),
+		)),
+	}
+}
+
+// withConfigLock serialises every mutation of one config's skill image.
+//
+// This is not defensive locking: a new snapshot is the OLD snapshot plus this
+// run's changes, and the config holds exactly one pointer. Two concurrent
+// installs would each snapshot a base that lacks the other's work, and whoever
+// wrote the pointer last would silently discard the other install.
+//
+// With Redis this is a 30s renewable lease, so every replica of the same
+// workspace contends on one key. Without Redis it is a process-local mutex
+// and two replicas can write the pointer independently.
+func (s *TenantSkillService) withConfigLock(
+	ctx context.Context, tenantID uint64, configID string, fn func(context.Context) error,
+) error {
+	key := skillImageLockKey(tenantID, configID)
+	if s.redis == nil {
+		release, err := s.localLocks.lock(ctx, key)
+		if err != nil {
+			return err
+		}
+		defer release()
+		return fn(ctx)
+	}
+	return redislock.WithRenewableLock(
+		ctx, s.redis, key, skillImageLockLease, skillImageLockRenew, fn,
+	)
+}
+
+func skillImageLockKey(tenantID uint64, configID string) string {
+	return fmt.Sprintf("weknora-skill-image-lock:%d:%s", tenantID, configID)
+}
+
+// clock is this service's time source. Tests inject one; a service built
+// without NewTenantSkillService still gets a working default.
+func (s *TenantSkillService) clock() func() time.Time {
+	if s != nil && s.now != nil {
+		return s.now
+	}
+	return time.Now
+}
+
+// keyedMutex is the no-Redis fallback for withConfigLock.
+type keyedMutex struct {
+	mu sync.Mutex
+	m  map[string]chan struct{}
+}
+
+func newKeyedMutex() *keyedMutex { return &keyedMutex{m: map[string]chan struct{}{}} }
+
+func (k *keyedMutex) lock(ctx context.Context, key string) (func(), error) {
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+
+	k.mu.Lock()
+	entry, ok := k.m[key]
+	if !ok {
+		entry = make(chan struct{}, 1)
+		k.m[key] = entry
+	}
+	k.mu.Unlock()
+	select {
+	case entry <- struct{}{}:
+		return func() { <-entry }, nil
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	}
+}