diff --git a/cmd/cli/main.go b/cmd/cli/main.go index b55cb47..25fc1ec 100644 --- a/cmd/cli/main.go +++ b/cmd/cli/main.go @@ -78,6 +78,8 @@ func main() { runRegister() case "add-models": runAddModels() + case "remove-models": + runRemoveModels() case "reinstate": runReinstate() case "top-up-stake": @@ -316,6 +318,35 @@ func runAddModels() { } } +// runRemoveModels takes SUPPORTED_MODELS as the list of models to REMOVE from +// the worker's on-chain set, e.g. +// SUPPORTED_MODELS=deepseek-r1:32b lightchain-worker remove-models +func runRemoveModels() { + cfg, logger := loadAndValidateCfg() + signingKey, workerAddr := loadSigningKey(cfg, logger) + modelIDs, modelNames := parseModels(cfg, logger) + + chainClient := dialChain(cfg, signingKey, logger) + defer chainClient.Close() + + h := &cli.Handler{ + Client: chainClient, + WorkerAddr: workerAddr, + ModelIDs: modelIDs, + ModelNames: modelNames, + Out: os.Stdout, + Logger: logger, + } + + ctx, cancel := context.WithTimeout(context.Background(), txTimeout) + defer cancel() + + if err := h.RemoveModels(ctx); err != nil { + logger.Error("remove-models failed", "error", err) + os.Exit(1) + } +} + func runDeregister() { cfg, logger := loadAndValidateCfg() signingKey, workerAddr := loadSigningKey(cfg, logger) diff --git a/internal/chain/client.go b/internal/chain/client.go index 4cc618b..9a1b484 100644 --- a/internal/chain/client.go +++ b/internal/chain/client.go @@ -224,6 +224,13 @@ func (c *ChainClient) AddSupportedModel(ctx context.Context, modelID [32]byte) e }) } +// RemoveSupportedModel drops a model from this worker's on-chain list (WorkerRegistry.removeSupportedModel). +func (c *ChainClient) RemoveSupportedModel(ctx context.Context, modelID [32]byte) error { + return c.submitPreparedTx(ctx, "RemoveSupportedModel", nil, func(opts *bind.TransactOpts) (*types.Transaction, error) { + return c.registry.RemoveSupportedModel(opts, modelID) + }) +} + // SetCapabilities declares this worker's full on-chain capability mask (overwrite, not merge).. func (c *ChainClient) SetCapabilities(ctx context.Context, mask *big.Int) error { return c.submitPreparedTx(ctx, "SetCapabilities", nil, func(opts *bind.TransactOpts) (*types.Transaction, error) { @@ -675,6 +682,17 @@ func (c *ChainClient) AcknowledgeJobAsync( return pending.Hash(), pending.Wait, nil } +// AbandonJob hands a job the worker cannot finish back to the registry +// (JobRegistry.abandonJob): the consumer's fee is refunded and the job leaves +// this worker's in-flight count at once instead of after deadline + +// resolutionTimeout. Without it a worker that fails a few jobs an hour fills +// the maxConcurrentJobs cap and the chain refuses its session claims. +func (c *ChainClient) AbandonJob(ctx context.Context, jobID uint64) error { + return c.submitPreparedTx(ctx, "AbandonJob", nil, func(opts *bind.TransactOpts) (*types.Transaction, error) { + return c.jobRegistry.AbandonJob(opts, new(big.Int).SetUint64(jobID)) + }) +} + // CompleteJob submits a completeJob transaction with a single bytes32 response // blob hash. Post-audit the contract takes one bytes32 and enforces // `blobhash(0) == responseBlobHash`, so the blob-carrying TX must contain diff --git a/internal/chain/client_test.go b/internal/chain/client_test.go index e25675f..7705f2d 100644 --- a/internal/chain/client_test.go +++ b/internal/chain/client_test.go @@ -29,6 +29,7 @@ type MockRegistrationClient struct { IsWorkerRegisteredFn func(ctx context.Context, worker common.Address) (bool, error) RegisterWorkerFn func(ctx context.Context, encryptionPubKey []byte, stake *big.Int) error AddSupportedModelFn func(ctx context.Context, modelID [32]byte) error + RemoveSupportedModelFn func(ctx context.Context, modelID [32]byte) error DeregisterWorkerFn func(ctx context.Context) error GetMinWorkerStakeFn func(ctx context.Context) (*big.Int, error) GetWorkerEncryptionKeyFn func(ctx context.Context, worker common.Address) ([]byte, error) @@ -60,6 +61,12 @@ func (m *MockRegistrationClient) RegisterWorker(ctx context.Context, encryptionP func (m *MockRegistrationClient) AddSupportedModel(ctx context.Context, modelID [32]byte) error { return m.AddSupportedModelFn(ctx, modelID) } +func (m *MockRegistrationClient) RemoveSupportedModel(ctx context.Context, modelID [32]byte) error { + if m.RemoveSupportedModelFn == nil { + return nil + } + return m.RemoveSupportedModelFn(ctx, modelID) +} func (m *MockRegistrationClient) DeregisterWorker(ctx context.Context) error { return m.DeregisterWorkerFn(ctx) diff --git a/internal/chain/interfaces.go b/internal/chain/interfaces.go index 1d47b79..ba4707b 100644 --- a/internal/chain/interfaces.go +++ b/internal/chain/interfaces.go @@ -89,6 +89,10 @@ type RegistrationClient interface { // Per Amendment A-2, this is non-payable — no additional ETH is required per model. AddSupportedModel(ctx context.Context, modelID [32]byte) error + // RemoveSupportedModel drops a model ID from the worker's supported model set on-chain, + // so sortition stops routing that model's session requests to this worker. + RemoveSupportedModel(ctx context.Context, modelID [32]byte) error + // DeregisterWorker removes the worker from the registry and withdraws all stake. DeregisterWorker(ctx context.Context) error diff --git a/internal/cli/handler.go b/internal/cli/handler.go index 199f2bf..48f3050 100644 --- a/internal/cli/handler.go +++ b/internal/cli/handler.go @@ -155,6 +155,39 @@ func (h *Handler) AddModels(ctx context.Context) error { return nil } +// RemoveModels removes models from an already-registered worker's on-chain +// list. SUPPORTED_MODELS in the environment only decides what the worker +// serves; sortition routes by the registry, so a model the box can no longer +// serve has to leave the registry too or its session requests keep landing +// here and timing out. +func (h *Handler) RemoveModels(ctx context.Context) error { + if len(h.ModelIDs) == 0 { + return fmt.Errorf("SUPPORTED_MODELS is required for remove-models: list the models to remove") + } + if len(h.ModelNames) != len(h.ModelIDs) { + return fmt.Errorf("ModelNames length %d does not match ModelIDs length %d", len(h.ModelNames), len(h.ModelIDs)) + } + + var failed []string + removed := 0 + for i, modelID := range h.ModelIDs { + if err := h.Client.RemoveSupportedModel(ctx, modelID); err != nil { + h.Logger.Error("remove model failed", "model", h.ModelNames[i], "index", i, "error", err) + failed = append(failed, fmt.Sprintf("%q: %v", h.ModelNames[i], err)) + continue + } + removed++ + h.Logger.Info("model removed", "model", h.ModelNames[i]) + } + + _, _ = fmt.Fprintf(h.Out, "Removed %d of %d models for worker %s\n", removed, len(h.ModelIDs), h.WorkerAddr.Hex()) + if len(failed) > 0 { + return fmt.Errorf("%d model(s) not removed (a model the worker never supported reverts): %s", + len(failed), strings.Join(failed, "; ")) + } + return nil +} + // Deregister removes the worker from the registry and withdraws stake. func (h *Handler) Deregister(ctx context.Context) error { mgr := registration.NewManager(h.Client, h.WorkerAddr, nil, h.Logger) diff --git a/internal/cli/handler_test.go b/internal/cli/handler_test.go index ea07a7f..066dddd 100644 --- a/internal/cli/handler_test.go +++ b/internal/cli/handler_test.go @@ -26,6 +26,7 @@ type mockRegClient struct { isWorkerRegisteredFn func(ctx context.Context, worker common.Address) (bool, error) registerWorkerFn func(ctx context.Context, encryptionPubKey []byte, stake *big.Int) error addSupportedModelFn func(ctx context.Context, modelID [32]byte) error + removeSupportedModelFn func(ctx context.Context, modelID [32]byte) error deregisterWorkerFn func(ctx context.Context) error getMinWorkerStakeFn func(ctx context.Context) (*big.Int, error) getWorkerEncryptionKeyFn func(ctx context.Context, worker common.Address) ([]byte, error) @@ -55,6 +56,12 @@ func (m *mockRegClient) RegisterWorker(ctx context.Context, encryptionPubKey []b func (m *mockRegClient) AddSupportedModel(ctx context.Context, modelID [32]byte) error { return m.addSupportedModelFn(ctx, modelID) } +func (m *mockRegClient) RemoveSupportedModel(ctx context.Context, modelID [32]byte) error { + if m.removeSupportedModelFn == nil { + return nil + } + return m.removeSupportedModelFn(ctx, modelID) +} func (m *mockRegClient) DeregisterWorker(ctx context.Context) error { return m.deregisterWorkerFn(ctx) } diff --git a/internal/cli/init_test.go b/internal/cli/init_test.go index 97352cd..3c9cb3a 100644 --- a/internal/cli/init_test.go +++ b/internal/cli/init_test.go @@ -198,6 +198,8 @@ func (c *initChain) RegisterWorker(_ context.Context, encKey []byte, stake *big. return nil } +func (c *initChain) RemoveSupportedModel(context.Context, [32]byte) error { return nil } + func (c *initChain) AddSupportedModel(_ context.Context, id [32]byte) error { c.addCalls = append(c.addCalls, id) if c.addErr != nil { diff --git a/internal/config/config.go b/internal/config/config.go index 09ee041..beb10ea 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -57,11 +57,18 @@ type Config struct { // Web search (Tavily). SearchEnabled gates advertisement of the "search" // capability and the per-job search stage. When false the worker behaves // exactly as before. - SearchEnabled bool - TavilyAPIKey string - TavilyURL string - SearchTimeout time.Duration - SearchMaxResults int + SearchEnabled bool + // PublicOutputEnabled serves envelope v4 public-output jobs and + // advertises the "public-output" capability. Default on; turn it off + // on an engine without JSON-schema `format` support. ExactProfile + // additionally advertises "exact-profile": the operator asserts the + // pinned deterministic profile for exact-match action jobs. + PublicOutputEnabled bool + ExactProfile bool + TavilyAPIKey string + TavilyURL string + SearchTimeout time.Duration + SearchMaxResults int // OllamaStream selects stream=true inference. When true the worker // publishes incremental `chunk` frames as tokens arrive, collapsing // time-to-first-token from full generation time to ~1s. The terminal @@ -312,6 +319,11 @@ type Config struct { // Sortition mode (Phase 3 — worker self-claims sessions, no dispatcher). SortitionEnabled bool + // AbandonFailedJobs makes the worker call JobRegistry.abandonJob for a job + // it failed, refunding the consumer and freeing its in-flight slot at once + // instead of after deadline + resolutionTimeout. Needs the abandonJob + // contract upgrade; off by default until it is live (ABANDON_FAILED_JOBS). + AbandonFailedJobs bool SessionManagerAddress common.Address SortitionStateDir string SortitionChunkSize uint64 @@ -326,6 +338,12 @@ type Config struct { // (1800 blocks at 2 s slots, 600 at 6 s); raise it where callers open // requests with longer expiries. Defaults to 2000; 0 disables. SortitionSessionLookbackBlocks uint64 + // SortitionMaxRequestAgeBlocks drops a pending session request once it is + // this many blocks old without a claim. The consumer-api stops waiting 60 s + // after it requested the session and never provides keys for it, so a + // claim after that only spends gas on a session nobody uses. Defaults to + // 150 (5 min at 2 s slots, 15 min at 6 s); 0 disables. + SortitionMaxRequestAgeBlocks uint64 // SortitionSessionRetryLimit bounds how many passes the JobWatcher retries // a job whose session is not Active before giving up and skipping it, so // one session that never returns to Active cannot wedge the watcher and @@ -461,6 +479,12 @@ func Load() (*Config, error) { } cfg.OllamaTimeout = parseDuration("OLLAMA_TIMEOUT", "120s", &errs) cfg.SearchEnabled = parseBool("SEARCH_ENABLED", false, &errs) + cfg.AbandonFailedJobs = parseBool("ABANDON_FAILED_JOBS", false, &errs) + cfg.PublicOutputEnabled = parseBool("PUBLIC_OUTPUT", true, &errs) + cfg.ExactProfile = parseBool("EXACT_PROFILE", false, &errs) + if cfg.ExactProfile && !cfg.PublicOutputEnabled { + errs = append(errs, "EXACT_PROFILE requires PUBLIC_OUTPUT") + } cfg.SearchTimeout = parseDuration("SEARCH_TIMEOUT", "10s", &errs) cfg.SearchMaxResults = parseInt("SEARCH_MAX_RESULTS", 5, &errs) cfg.AckTxTimeout = parseDuration("ACK_TX_TIMEOUT", "15s", &errs) @@ -564,6 +588,7 @@ func Load() (*Config, error) { cfg.SortitionPollInterval = parseDuration("SORTITION_POLL_INTERVAL", "1s", &errs) cfg.SortitionConfirmations = parseUint64("SORTITION_CONFIRMATIONS", 0, &errs) cfg.SortitionSessionLookbackBlocks = parseUint64("SORTITION_SESSION_LOOKBACK_BLOCKS", 2000, &errs) + cfg.SortitionMaxRequestAgeBlocks = parseUint64("SORTITION_MAX_REQUEST_AGE_BLOCKS", 150, &errs) cfg.SortitionSessionRetryLimit = parseInt("SORTITION_SESSION_RETRY_LIMIT", 10, &errs) if len(errs) > 0 { diff --git a/internal/ollama/client.go b/internal/ollama/client.go index 1bb1fbe..f22d271 100644 --- a/internal/ollama/client.go +++ b/internal/ollama/client.go @@ -31,6 +31,10 @@ const ( // (re-execution) rely on greedy decoding being reproducible across services. const deterministicSeed = 42 +// DeterministicSeed is deterministicSeed for callers that record the seed a +// job ran under (the public response R echoes it). +const DeterministicSeed = deterministicSeed + // Options mirrors the Ollama request "options" object. Only the knobs the // worker needs are modelled. Every field is omitempty so an unset option // leaves the server-side default untouched. diff --git a/internal/pipeline/handler.go b/internal/pipeline/handler.go index d99de3b..cde4ad7 100644 --- a/internal/pipeline/handler.go +++ b/internal/pipeline/handler.go @@ -163,6 +163,20 @@ type HandlerConfig struct { // an empty map reproduces pre-override request bodies byte for byte. ModelOptions map[string]ollama.ClientOptions + // PublicOutput serves envelope v4 public-output jobs (config + // PUBLIC_OUTPUT). Such a job runs on a pinned copy of the inference + // client (the envelope's numCtx and numPredict, temperature 0, the + // deterministic seed, thinking off, the envelope's format) and commits + // the plaintext public response R instead of a ciphertext. When off, + // a public job is refused without retry rather than silently + // encrypted: an encrypted answer can never be delivered to the + // contract that asked for it, and refusing lets the fee be refunded. + // PublicMaxNumCtx and PublicMaxNumPredict bound what a job may ask + // for (0 or negative: no bound); a job over them is refused the same way. + PublicOutput bool + PublicMaxNumCtx int + PublicMaxNumPredict int + // DeadlineGuardEnabled turns on the on-chain deadline guard: the // pipeline reads job.deadline at pickup, before inference, and before // the stage-8a blob tx, and aborts with a no-retry error when the job @@ -197,16 +211,16 @@ type HandlerConfig struct { // answer is unaffected). TTSTimeout bounds one synthesis call; zero // falls back to 20 s. It also feeds the deadline budget check together // with a fixed stream margin. - STTEnabled bool - TTSEnabled bool - TTSVoice string + STTEnabled bool + TTSEnabled bool + TTSVoice string // SpeechModelName is the model whose jobs settle audio rather than // text. Empty leaves the speech path inert. - SpeechModelName string - TTSMaxChars int - TTSTimeout time.Duration - SearchMaxResults int - SearchTimeout time.Duration + SpeechModelName string + TTSMaxChars int + TTSTimeout time.Duration + SearchMaxResults int + SearchTimeout time.Duration } // ResponsePublisher publishes encrypted responses for real-time delivery. @@ -273,6 +287,7 @@ type ReleaseTracker interface { // the gateway via HandleJobPayload. type JobHandler struct { chainClient JobExecutionClient + abandonFailedJobs bool blobFetcher BlobFetcher blobSubmitter BlobSubmitter keyStore SessionKeyGetter @@ -739,14 +754,67 @@ const errorFramePublishTimeout = 5 * time.Second // failed. It stays silent when the answer was already delivered, because an // error frame then would replace a delivered answer with an error. func (h *JobHandler) reportFailure(ctx context.Context, p JobPayload, err error) { - if h.responsePublisher == nil || errors.Is(err, errAnswerDelivered) { + if errors.Is(err, errAnswerDelivered) { return } - // A cancelled job context (timeout, shutdown) is often why the job - // failed; the consumer still needs to hear about it. - pubCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), errorFramePublishTimeout) - defer cancel() - h.responsePublisher.PublishError(pubCtx, p.JobID, p.SessionID, p.CorrelationID) + if h.responsePublisher != nil { + // A cancelled job context (timeout, shutdown) is often why the job + // failed; the consumer still needs to hear about it. + pubCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), errorFramePublishTimeout) + defer cancel() + h.responsePublisher.PublishError(pubCtx, p.JobID, p.SessionID, p.CorrelationID) + } + h.abandonOnChain(ctx, p, err) +} + +// jobAbandoner is implemented by chain clients that can hand a failed job back +// to the registry (JobRegistry.abandonJob). Optional so test doubles and older +// deployments without the method keep working unchanged. +type jobAbandoner interface { + AbandonJob(ctx context.Context, jobID uint64) error +} + +// abandonTxTimeout bounds the best-effort abandon tx so a dead RPC cannot +// hold the handler after the job already failed. +const abandonTxTimeout = 90 * time.Second + +// SetAbandonFailedJobs turns on the on-chain abandon of failed jobs. Off by +// default: JobRegistry.abandonJob only exists from the abandonJob upgrade on, +// and before it every call would just revert. +func (h *JobHandler) SetAbandonFailedJobs(on bool) { h.abandonFailedJobs = on } + +// abandonOnChain settles a failed job on chain so it stops counting toward this +// worker's in-flight cap. Until the job settles the registry keeps it +// Acknowledged for deadline + resolutionTimeout (48 h on mainnet), and a worker +// with maxConcurrentJobs of those is refused by SessionManager.eligibleNow: +// its models look available and every request for them times out. Best effort: +// a failure here is logged and nothing else, the keeper's no-fault settle still +// clears the job later. Runs in the background so the error frame the consumer +// is waiting on is not delayed by a transaction. +func (h *JobHandler) abandonOnChain(ctx context.Context, p JobPayload, cause error) { + if !h.abandonFailedJobs || h.chainClient == nil { + return + } + ab, ok := h.chainClient.(jobAbandoner) + if !ok { + return + } + // A job whose answer is already on chain must not be abandoned. + if errors.Is(cause, errAnswerDelivered) { + return + } + go func() { + txCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), abandonTxTimeout) + defer cancel() + if err := ab.AbandonJob(txCtx, p.JobID); err != nil { + // JobNotInState: it completed, timed out or was already settled meanwhile; nothing to do. + h.logger.Warn("abandonJob failed; the keeper's no-fault settle will clear it later", + "jobID", p.JobID, "sessionID", p.SessionID, "error", err) + return + } + h.logger.Info("job abandoned on chain: fee refunded, in-flight slot released", + "jobID", p.JobID, "sessionID", p.SessionID, "cause", cause.Error()) + }() } // processJob executes the 8-stage inference pipeline. @@ -1208,12 +1276,24 @@ func (h *JobHandler) runInferencePipeline( envelope, err := promptenv.Decode([]byte(promptText)) if err != nil { err = fmt.Errorf("stage 5 (decode prompt): %w", err) - // A malformed self-contained job fails the same way on every attempt. - if errors.Is(err, promptenv.ErrInvalidSelfContained) { + // A malformed self-contained or public job fails the same way on + // every attempt. + if errors.Is(err, promptenv.ErrInvalidSelfContained) || errors.Is(err, promptenv.ErrInvalidPublic) { return nil, 0, noRetry(err) } return nil, 0, err } + // A public-output job replaces the operator's generation knobs with the + // envelope's and commits its answer in the clear. Anything that stops + // that from being replayable is a refusal, not a best effort. + var pub *publicJob + if envelope.Public() { + pub, err = h.preparePublicJob(&envelope, prompt, searchEnabled, modelName) + if err != nil { + return nil, 0, noRetry(fmt.Errorf("stage 5 (public output): %w", err)) + } + infClient, effectiveOpts, overrideApplied = pub.client, pub.client.Options(), true + } // The speech model reads one text aloud and has no chat call for a // conversation, so a self-contained job for it can never be served. if envelope.SelfContained() && h.cfg.SpeechModelName != "" && modelName == h.cfg.SpeechModelName { @@ -1357,21 +1437,35 @@ func (h *JobHandler) runInferencePipeline( // disputer can reproduce the exact augmented prompt on replay. For plain // (non-search) jobs EncodeResponse returns raw answer bytes — byte-identical // to the legacy format. + // + // A public-output job commits the canonical public response R in + // plaintext instead: the terminal frame, the blob and completeJob all + // carry R itself, and keccak256(R) is what the agent contract checks. rec = h.metrics.StartStage(metrics.StageEncrypt, model, delivery) - respBytes, encErr := searchaug.EncodeResponse(response, searchSources) - if encErr != nil { - rec.End(metrics.OutcomeError, metrics.CacheMiss) - return nil, 0, fmt.Errorf("stage 6 (encode response): %w", encErr) - } - ciphertext, err := pkgcrypto.Encrypt(sessionKey, respBytes) - if err != nil { - rec.End(metrics.OutcomeError, metrics.CacheMiss) - return nil, streamer.frames(), fmt.Errorf("stage 6 (encrypt response): %w", err) + var ciphertext []byte + if pub != nil { + ciphertext, err = pub.encode(response) + if err != nil { + rec.End(metrics.OutcomeError, metrics.CacheMiss) + return nil, streamer.frames(), noRetry(fmt.Errorf("stage 6 (encode public response): %w", err)) + } + } else { + respBytes, encErr := searchaug.EncodeResponse(response, searchSources) + if encErr != nil { + rec.End(metrics.OutcomeError, metrics.CacheMiss) + return nil, 0, fmt.Errorf("stage 6 (encode response): %w", encErr) + } + ciphertext, err = pkgcrypto.Encrypt(sessionKey, respBytes) + if err != nil { + rec.End(metrics.OutcomeError, metrics.CacheMiss) + return nil, streamer.frames(), fmt.Errorf("stage 6 (encrypt response): %w", err) + } } d = rec.End(metrics.OutcomeOK, metrics.CacheMiss) logger.Info( "stage 6 complete", "stage", "encrypt", + "public", pub != nil, "ciphertextBytes", len(ciphertext), "durationMs", d.Milliseconds(), ) diff --git a/internal/pipeline/public.go b/internal/pipeline/public.go new file mode 100644 index 0000000..46cc351 --- /dev/null +++ b/internal/pipeline/public.go @@ -0,0 +1,114 @@ +package pipeline + +import ( + "encoding/json" + "errors" + "fmt" + + "github.com/lightchain/pkg/promptenv" + "github.com/lightchain/pkg/publicresp" + + "github.com/lightchain/worker/internal/ollama" +) + +// ErrPublicOutputDisabled is returned for a public-output job on a worker +// configured with PUBLIC_OUTPUT=false. The job is refused without retry: an +// encrypted answer would never match what the requester can deliver. +var ErrPublicOutputDisabled = errors.New("public output is disabled on this worker") + +// publicJob is what stage 5 and stage 6 need for an envelope v4 public-output +// job: the pinned inference client and the hashes that go into R. +type publicJob struct { + client *ollama.OllamaClient + modelName string + params promptenv.Params + hasFormat bool + promptHash string + formatHash string +} + +// preparePublicJob checks a public-output envelope against this worker's +// limits and builds the pinned client the job must run on. plainEnvelope is +// the decrypted prompt exactly as the blob carried it: R.promptHash commits +// to those bytes so the disputer can tell a worker that answered a different +// prompt. +func (h *JobHandler) preparePublicJob( + env *promptenv.Envelope, + plainEnvelope []byte, + searchEnabled bool, + modelName string, +) (*publicJob, error) { + switch { + case !h.cfg.PublicOutput: + return nil, ErrPublicOutputDisabled + case searchEnabled: + return nil, errors.New("public output is not supported for search jobs") + case env.Params == nil: + // Decode guarantees this; kept so the method is safe on its own. + return nil, errors.New("public output needs params") + case h.cfg.PublicMaxNumCtx > 0 && env.Params.NumCtx > h.cfg.PublicMaxNumCtx: + return nil, fmt.Errorf("numCtx %d exceeds this worker's limit %d", env.Params.NumCtx, h.cfg.PublicMaxNumCtx) + case h.cfg.PublicMaxNumPredict > 0 && env.Params.NumPredict > h.cfg.PublicMaxNumPredict: + return nil, fmt.Errorf("numPredict %d exceeds this worker's limit %d", env.Params.NumPredict, h.cfg.PublicMaxNumPredict) + } + oc, ok := h.ollamaClient.(overridingInferenceClient) + if !ok { + return nil, errors.New("the inference client cannot pin generation parameters") + } + format, err := env.CanonicalFormat() + if err != nil { + return nil, fmt.Errorf("format: %w", err) + } + formatHash, err := publicresp.FormatHash(format) + if err != nil { + return nil, err + } + + params := *env.Params + client := oc.WithOverrides(func(o *ollama.ClientOptions) { + // Exactly the profile R echoes (draft 04 §4.2): the envelope's + // sizes, greedy decoding, the fixed seed, no thinking. Operator + // MODEL_OPTIONS do not apply to a public job. + o.NumCtx = params.NumCtx + o.NumPredict = params.NumPredict + greedy := 0.0 + o.Temperature = &greedy + o.Think = ollama.ThinkSetting(false) + o.ThinkLevel = "" + if format != nil { + o.Format = json.RawMessage(format) + } else { + o.Format = nil + } + }) + return &publicJob{ + client: client, + modelName: modelName, + params: params, + hasFormat: format != nil, + promptHash: publicresp.PromptHash(plainEnvelope), + formatHash: formatHash, + }, nil +} + +// encode builds R from the engine's text. The bytes returned are what the +// job commits: the terminal frame carries them, the blob anchors them and +// completeJob hashes them. +func (p *publicJob) encode(text string) ([]byte, error) { + output, answer, valid := publicresp.FromEngineText(text, p.hasFormat) + return publicresp.Encode(publicresp.Response{ + V: publicresp.Version, + PromptHash: p.promptHash, + FormatHash: p.formatHash, + Params: publicresp.Params{ + Model: p.modelName, + NumCtx: p.params.NumCtx, + NumPredict: p.params.NumPredict, + Seed: ollama.DeterministicSeed, + Temperature: 0, + }, + Valid: valid, + Output: output, + Answer: answer, + }) +} diff --git a/internal/pipeline/public_test.go b/internal/pipeline/public_test.go new file mode 100644 index 0000000..65a5304 --- /dev/null +++ b/internal/pipeline/public_test.go @@ -0,0 +1,249 @@ +package pipeline + +import ( + "context" + "encoding/json" + "errors" + "io" + "log/slog" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" + + "github.com/alicebob/miniredis/v2" + "github.com/ethereum/go-ethereum/common" + "github.com/ethereum/go-ethereum/crypto" + "github.com/hibiken/asynq" + "github.com/redis/go-redis/v9" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/lightchain/pkg/publicresp" + "github.com/lightchain/worker/internal/metrics" + "github.com/lightchain/worker/internal/ollama" +) + +const publicSalt = "0x0102030405060708091011121314151617181920212223242526272829303132" + +// publicEnvelope is a v4 public-output prompt as the blob carries it. +func publicEnvelope(format string, numCtx, numPredict int) string { + f := "" + if format != "" { + f = `,"format":` + format + } + return `{"v":4,"text":"What is 2+2?","output":"public"` + f + + `,"params":{"numCtx":` + itoa(numCtx) + `,"numPredict":` + itoa(numPredict) + `},"salt":"` + publicSalt + `"}` +} + +func itoa(n int) string { + b, _ := json.Marshal(n) + return string(b) +} + +// ollamaGenerateServer answers POST /api/generate with one batch response +// and hands every request body to capture. +func ollamaGenerateServer(t *testing.T, response string, capture func(body map[string]any)) *httptest.Server { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + require.Equal(t, "/api/generate", r.URL.Path) + var body map[string]any + require.NoError(t, json.NewDecoder(r.Body).Decode(&body)) + capture(body) + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]any{"model": body["model"], "response": response, "done": true}) + })) + t.Cleanup(srv.Close) + return srv +} + +type publicRun struct { + plaintext string + blob []byte + completedHash [32]byte + completeCalled bool + requests []map[string]any +} + +// runPublicJob drives HandleTask for a public envelope against a real +// *ollama.OllamaClient (the pinned-client path needs WithOverrides, which +// the handler mocks do not have) talking to an httptest Ollama. +func runPublicJob(t *testing.T, plaintext, engineText string, cfg func(*HandlerConfig)) (*publicRun, error) { + t.Helper() + mr := miniredis.RunT(t) + rc := redis.NewClient(&redis.Options{Addr: mr.Addr()}) + t.Cleanup(func() { _ = rc.Close() }) + + run := &publicRun{plaintext: plaintext} + sessionKey := testSessionKey(t) + ecdhKey := testECDHKey(t) + encSessionKey := encryptSessionKeyForWorker(t, sessionKey, ecdhKey) + promptCiphertext := encryptBlob(t, sessionKey, plaintext) + modelID := common.Bytes2Hex(crypto.Keccak256Hash([]byte("llama3-8b")).Bytes()) + + chain := &mockChainClient{ + ackJobFn: func(context.Context, uint64) error { return nil }, + completeJobFn: func(_ context.Context, _ uint64, _ [32]byte, responseCiphertextHash [32]byte) error { + run.completeCalled = true + run.completedHash = responseCiphertextHash + return nil + }, + getEncWorkerKeyFn: func(context.Context, uint64) ([]byte, error) { return encSessionKey, nil }, + } + fetcher := &mockBlobFetcher{fetchFn: func(context.Context, common.Hash, uint64) ([]byte, error) { + return promptCiphertext, nil + }} + submitter := &mockBlobSubmitter{submitFn: func(_ context.Context, data []byte) ([][32]byte, error) { + run.blob = append([]byte(nil), data...) + return [][32]byte{{0x01}}, nil + }} + srv := ollamaGenerateServer(t, engineText, func(body map[string]any) { run.requests = append(run.requests, body) }) + // Operator defaults that a public job must override. + client := ollama.NewOllamaClientWithOptions(srv.URL, 5*time.Second, ollama.ClientOptions{ + NumPredict: 1024, NumCtx: 8192, Think: ollama.ThinkSetting(true), ThinkLevel: "high", + }) + + hc := HandlerConfig{ + AckTxTimeout: 5 * time.Second, + BlobTxTimeout: 60 * time.Second, + RedisPublishTimeout: 5 * time.Second, + ModelIDToName: map[string]string{modelID: "llama3-8b"}, + PublicOutput: true, + PublicMaxNumCtx: 8192, + PublicMaxNumPredict: 1024, + } + if cfg != nil { + cfg(&hc) + } + logger := slog.New(slog.NewTextHandler(io.Discard, &slog.HandlerOptions{Level: slog.LevelError})) + handler := NewJobHandler( + chain, fetcher, submitter, newMockKeyStore(), client, + rc, testSigningKey(t), ecdhKey, &atomic.Int32{}, logger, + hc, nil, nil, testMetrics(t), metrics.DeliveryAsynq, nil, + ) + + payload := testPayload(t) + payload.ModelID = modelID + data, err := json.Marshal(payload) + require.NoError(t, err) + return run, handler.HandleTask(context.Background(), asynq.NewTask(TaskTypeJobInference, data)) +} + +func TestHandleTask_PublicOutput_CommitsCanonicalR(t *testing.T) { + t.Parallel() + + schema := `{"type":"object","properties":{"n":{"type":"integer"},"ok":{"type":"boolean"}}}` + plaintext := publicEnvelope(schema, 2048, 64) + run, err := runPublicJob(t, plaintext, "{\"ok\": true, \"n\": 4}", nil) + require.NoError(t, err) + + // The engine ran under the envelope's profile, not the operator's. + require.Len(t, run.requests, 1) + req := run.requests[0] + opts, _ := req["options"].(map[string]any) + assert.Equal(t, float64(2048), opts["num_ctx"]) + assert.Equal(t, float64(64), opts["num_predict"]) + assert.Equal(t, float64(0), opts["temperature"]) + assert.Equal(t, float64(ollama.DeterministicSeed), opts["seed"]) + assert.Equal(t, false, req["think"], "thinking is forced off for a format job") + gotFormat, _ := json.Marshal(req["format"]) + assert.JSONEq(t, schema, string(gotFormat), "the schema reaches the engine as format") + + // The blob is R in plaintext and completeJob committed keccak256(R). + require.True(t, run.completeCalled) + r, err := publicresp.Decode(run.blob) + require.NoError(t, err, "blob is canonical R") + assert.Equal(t, crypto.Keccak256Hash(run.blob), common.Hash(run.completedHash)) + assert.Equal(t, publicresp.Hash(run.blob), common.Hash(run.completedHash).Hex()) + + assert.Equal(t, publicresp.Version, r.V) + assert.Equal(t, publicresp.PromptHash([]byte(plaintext)), r.PromptHash, "promptHash is keccak of the decrypted envelope") + wantFormat, err := publicresp.FormatHash(json.RawMessage(schema)) + require.NoError(t, err) + assert.Equal(t, wantFormat, r.FormatHash) + assert.Equal(t, publicresp.Params{Model: "llama3-8b", NumCtx: 2048, NumPredict: 64, Seed: 42, Temperature: 0}, r.Params) + assert.True(t, r.Valid) + assert.Nil(t, r.Answer) + assert.Equal(t, `{"n":4,"ok":true}`, string(r.Output), "output is canonical inside R") +} + +func TestHandleTask_PublicOutput_WithoutFormatIsRawAnswer(t *testing.T) { + t.Parallel() + + plaintext := publicEnvelope("", 2048, 64) + run, err := runPublicJob(t, plaintext, "4", nil) + require.NoError(t, err) + require.Len(t, run.requests, 1) + _, hasFormat := run.requests[0]["format"] + assert.False(t, hasFormat) + + r, err := publicresp.Decode(run.blob) + require.NoError(t, err) + assert.Equal(t, publicresp.ZeroHash, r.FormatHash) + assert.False(t, r.Valid) + require.NotNil(t, r.Answer) + assert.Equal(t, "4", *r.Answer) + assert.Nil(t, r.Output) +} + +// A generation cut mid-object is committed as it came, valid false, so a +// replay reproduces it; the worker never retries with a new seed. +func TestHandleTask_PublicOutput_InvalidJSONIsCommittedNotRetried(t *testing.T) { + t.Parallel() + + plaintext := publicEnvelope(`{"type":"object"}`, 2048, 8) + run, err := runPublicJob(t, plaintext, `{"summary":"the mod`, nil) + require.NoError(t, err) + require.Len(t, run.requests, 1) + r, err := publicresp.Decode(run.blob) + require.NoError(t, err) + assert.False(t, r.Valid) + require.NotNil(t, r.Answer) + assert.Equal(t, `{"summary":"the mod`, *r.Answer) +} + +func TestHandleTask_PublicOutput_DisabledRefusesWithoutRetry(t *testing.T) { + t.Parallel() + + run, err := runPublicJob(t, publicEnvelope("", 2048, 64), "4", func(c *HandlerConfig) { c.PublicOutput = false }) + require.Error(t, err) + assert.True(t, errors.Is(err, asynq.SkipRetry), "no retry: %v", err) + assert.True(t, errors.Is(err, ErrPublicOutputDisabled), "%v", err) + assert.Empty(t, run.requests, "the engine is never called") + assert.False(t, run.completeCalled) +} + +func TestHandleTask_PublicOutput_ParamsOverLimitRefusedWithoutRetry(t *testing.T) { + t.Parallel() + + run, err := runPublicJob(t, publicEnvelope("", 16384, 64), "4", nil) + require.Error(t, err) + assert.True(t, errors.Is(err, asynq.SkipRetry), "no retry: %v", err) + assert.Empty(t, run.requests) + assert.False(t, run.completeCalled) +} + +// A v4 envelope that is public but malformed (no salt) is refused at decode, +// without retry, before anything is sent to the engine. +func TestHandleTask_PublicOutput_MalformedEnvelopeRefused(t *testing.T) { + t.Parallel() + + plaintext := `{"v":4,"text":"hi","output":"public","params":{"numCtx":2048,"numPredict":64}}` + run, err := runPublicJob(t, plaintext, "4", nil) + require.Error(t, err) + assert.True(t, errors.Is(err, asynq.SkipRetry), "no retry: %v", err) + assert.Empty(t, run.requests) +} + +// A private job on the same handler still encrypts: the public branch does +// not change the default path. +func TestHandleTask_PrivateJobStillEncrypts(t *testing.T) { + t.Parallel() + + run, err := runPublicJob(t, "What is 2+2?", "4", nil) + require.NoError(t, err) + _, decErr := publicresp.Decode(run.blob) + assert.Error(t, decErr, "an encrypted blob is not R") + assert.NotEqual(t, []byte("4"), run.blob) +} diff --git a/internal/registration/manager.go b/internal/registration/manager.go index 91ac8b3..7b22bba 100644 --- a/internal/registration/manager.go +++ b/internal/registration/manager.go @@ -177,10 +177,28 @@ const SelfContainedCapability = "self-contained" // SearchCapability is declared by a worker configured for web search. const SearchCapability = "search" +// PublicOutputCapability is declared by a worker that serves public-output +// jobs: envelope v4 `output: public`, answered as the plaintext public +// response R (github.com/lightchain/pkg/publicresp) instead of ciphertext. +// A session that requires it never lands on a worker that would encrypt. +const PublicOutputCapability = "public-output" + +// ExactProfileCapability is declared by an operator who runs the pinned +// deterministic profile for exact-match jobs (one GPU class, one engine +// build, no batching). It is an operator assertion, measured in shadow mode. +const ExactProfileCapability = "exact-profile" + // DesiredCapabilities is the capability list this binary declares on chain: -// self-contained always, search only when it is configured. -func DesiredCapabilities(search bool) []string { +// self-contained always, search only when it is configured, public-output +// unless the operator turned it off, exact-profile only when asserted. +func DesiredCapabilities(search, publicOutput, exactProfile bool) []string { caps := []string{SelfContainedCapability} + if publicOutput { + caps = append(caps, PublicOutputCapability) + } + if publicOutput && exactProfile { + caps = append(caps, ExactProfileCapability) + } if search { caps = append(caps, SearchCapability) } diff --git a/internal/registration/manager_test.go b/internal/registration/manager_test.go index 4884bf8..dd80c94 100644 --- a/internal/registration/manager_test.go +++ b/internal/registration/manager_test.go @@ -51,6 +51,7 @@ func (m *mockClient) RegisterWorker(ctx context.Context, encryptionPubKey []byte func (m *mockClient) AddSupportedModel(ctx context.Context, modelID [32]byte) error { return m.addSupportedModelFn(ctx, modelID) } +func (m *mockClient) RemoveSupportedModel(context.Context, [32]byte) error { return nil } func (m *mockClient) DeregisterWorker(ctx context.Context) error { return m.deregisterWorkerFn(ctx) @@ -584,7 +585,7 @@ func TestDesiredCapabilities_DeclaresSelfContained(t *testing.T) { }, } - NewManager(mock, testAddr, nil, newLogger()).EnsureCapabilities(context.Background(), DesiredCapabilities(tc.search)) + NewManager(mock, testAddr, nil, newLogger()).EnsureCapabilities(context.Background(), DesiredCapabilities(tc.search, false, false)) require.NotNil(t, setMask) assert.Zero(t, setMask.Cmp(big.NewInt(tc.want)), "declared mask %s", setMask) @@ -609,10 +610,22 @@ func TestDesiredCapabilities_SelfContainedNotRegistered(t *testing.T) { }, } - NewManager(mock, testAddr, nil, logger).EnsureCapabilities(context.Background(), DesiredCapabilities(false)) + NewManager(mock, testAddr, nil, logger).EnsureCapabilities(context.Background(), DesiredCapabilities(false, false, false)) assert.Contains(t, logs.String(), "capability not registered on-chain") assert.Contains(t, logs.String(), SelfContainedCapability) require.NotNil(t, setMask, "the dropped search bit must still be cleared") assert.Zero(t, setMask.Sign()) } + +func TestDesiredCapabilities_PublicOutputAndExactProfile(t *testing.T) { + t.Parallel() + + assert.Equal(t, []string{SelfContainedCapability}, DesiredCapabilities(false, false, false)) + assert.Equal(t, []string{SelfContainedCapability, PublicOutputCapability}, DesiredCapabilities(false, true, false)) + assert.ElementsMatch(t, + []string{SelfContainedCapability, SearchCapability, PublicOutputCapability, ExactProfileCapability}, + DesiredCapabilities(true, true, true)) + assert.Equal(t, []string{SelfContainedCapability}, DesiredCapabilities(false, false, true), + "exact-profile is meaningless without public output") +} diff --git a/internal/service/service.go b/internal/service/service.go index 4ac35ae..06fb246 100644 --- a/internal/service/service.go +++ b/internal/service/service.go @@ -530,6 +530,9 @@ func New(cfg *config.Config) (*Service, error) { StreamChunkTokens: cfg.StreamChunkTokens, StreamChunkInterval: cfg.StreamChunkInterval, ModelOptions: cfg.ModelOptions, + PublicOutput: cfg.PublicOutputEnabled, + PublicMaxNumCtx: cfg.OllamaNumCtx, + PublicMaxNumPredict: cfg.OllamaNumPredict, DeadlineGuardEnabled: cfg.DeadlineGuardEnabled, SettleReserve: cfg.SettleReserve, CompletionReserve: cfg.CompletionReserve, @@ -556,6 +559,7 @@ func New(cfg *config.Config) (*Service, error) { // A chat job's earlier jobs come from the chain on every intake, as the // disputer finds them to re-run the job, never from a dispatcher's index. handler.SetPriorJobLister(chainClient) + handler.SetAbandonFailedJobs(cfg.AbandonFailedJobs) // --- Gateway mode: skip Asynq and direct Redis heartbeat --- // SortitionEnabled takes precedence: if both are set we already warned above @@ -604,7 +608,10 @@ func New(cfg *config.Config) (*Service, error) { AckOverlapEnabled: cfg.AckOverlapEnabled, // MODEL_OPTIONS applies in gateway mode too: tier caps are // generation parameters, independent of the delivery channel. - ModelOptions: cfg.ModelOptions, + ModelOptions: cfg.ModelOptions, + PublicOutput: cfg.PublicOutputEnabled, + PublicMaxNumCtx: cfg.OllamaNumCtx, + PublicMaxNumPredict: cfg.OllamaNumPredict, // The deadline guard matters just as much in gateway mode: // the gateway retries failed jobs with its own backoff, and // every guard abort is cheap (pre-inference), so the @@ -638,6 +645,7 @@ func New(cfg *config.Config) (*Service, error) { gwHandler.SetVoiceEngine(voiceEngine) } gwHandler.SetPriorJobLister(chainClient) + gwHandler.SetAbandonFailedJobs(cfg.AbandonFailedJobs) logger.Info( "worker service initialized (gateway mode)", @@ -784,7 +792,9 @@ func New(cfg *config.Config) (*Service, error) { // Every failure degrades to an empty mask — unconstrained requests // stay claimable and the on-chain claimSession check remains the // guarantee for constrained ones. - desiredCaps := registration.DesiredCapabilities(cfg.SearchEnabled && cfg.TavilyAPIKey != "") + desiredCaps := registration.DesiredCapabilities( + cfg.SearchEnabled && cfg.TavilyAPIKey != "", cfg.PublicOutputEnabled, cfg.ExactProfile, + ) capCtx, capCancel := context.WithTimeout(context.Background(), 30*time.Second) registration.NewManager(chainClient, workerAddr, nil, logger).EnsureCapabilities(capCtx, desiredCaps) capCancel() @@ -813,6 +823,8 @@ func New(cfg *config.Config) (*Service, error) { Logger: logger, OwnCapabilities: ownCaps, LookbackBlocks: cfg.SortitionSessionLookbackBlocks, + + MaxRequestAgeBlocks: cfg.SortitionMaxRequestAgeBlocks, }) jobWatcher = sortition.NewJobWatcher(sortition.JobWatcherOpts{ Client: chainClient, diff --git a/internal/sortition/session_watcher.go b/internal/sortition/session_watcher.go index 93b35d9..3adfc6f 100644 --- a/internal/sortition/session_watcher.go +++ b/internal/sortition/session_watcher.go @@ -55,6 +55,13 @@ type SessionWatcherOpts struct { // the look-back. It is also how far behind the safe head a worker with no // stored cursor starts discovery (see CursorStore.GetOrSeed). LookbackBlocks uint64 + // MaxRequestAgeBlocks drops a request once it is this many blocks behind the + // head without a claim. The consumer-api gives up on a request after its + // claim timeout (60 s) and never provides session keys for it, but the + // request stays Open on chain until its expiry (an hour), so a worker that + // frees up late used to claim dozens of abandoned requests at once: one tx + // each, for sessions nobody is waiting on. Zero disables the check. + MaxRequestAgeBlocks uint64 } // SessionWatcher polls the chain for SessionRequested events and claims @@ -85,6 +92,10 @@ type SessionWatcher struct { // evaluation (nil = not yet fetched). The mask is immutable per request, so a // successful fetch is cached for every later pass. pending map[uint64]*big.Int + // requestBlock remembers each pending request's block so the age check in + // the evaluation phase needs no extra chain read. + requestBlock map[uint64]uint64 + maxAge uint64 } // NewSessionWatcher constructs a SessionWatcher from the given options. @@ -115,7 +126,34 @@ func NewSessionWatcher(o SessionWatcherOpts) *SessionWatcher { ownCaps: ownCaps, lookback: o.LookbackBlocks, pending: make(map[uint64]*big.Int), + + requestBlock: make(map[uint64]uint64), + maxAge: o.MaxRequestAgeBlocks, + } +} + +// discover adds a request to the pending set, remembering its block for the age check. +func (w *SessionWatcher) discover(r chain.SessionRequestedEvent) { + if _, ok := w.pending[r.ReqID]; !ok { + w.pending[r.ReqID] = nil // capability mask fetched lazily at evaluation + } + w.requestBlock[r.ReqID] = r.RequestBlock +} + +// forget drops a request from the pending set and its bookkeeping. +func (w *SessionWatcher) forget(reqID uint64) { + delete(w.pending, reqID) + delete(w.requestBlock, reqID) +} + +// tooOld reports whether a pending request is older than MaxRequestAgeBlocks: +// its consumer gave up on it long ago, so claiming it would only spend a tx. +func (w *SessionWatcher) tooOld(reqID, head uint64) bool { + if w.maxAge == 0 { + return false } + rb, ok := w.requestBlock[reqID] + return ok && rb > 0 && head > rb && head-rb > w.maxAge } // Start runs a ticker loop calling RunOnce on each tick until ctx is cancelled. @@ -213,7 +251,7 @@ func (w *SessionWatcher) RunOnce(ctx context.Context) error { } for _, r := range reqs { - w.pending[r.ReqID] = nil // capability mask fetched lazily at evaluation + w.discover(r) } if err := w.cursor.Set(cursorSessionRequested, hi); err != nil { @@ -227,6 +265,14 @@ func (w *SessionWatcher) RunOnce(ctx context.Context) error { // prior pass when the worker was ineligible is re-evaluated here every // pass until the sortition claim window widens. for reqID := range w.pending { + // Abandoned by its consumer: the consumer-api stops waiting after its + // claim timeout and never provides keys, so a late claim is a wasted tx. + if w.tooOld(reqID, head.Number) { + w.log.Debug("dropping stale session request", "reqId", reqID, "requestBlock", w.requestBlock[reqID], "head", head.Number, "maxAgeBlocks", w.maxAge) + w.forget(reqID) + continue + } + ri, err := w.c.GetRequestInfo(ctx, reqID) if err != nil { w.log.Warn("getRequestInfo failed", "reqId", reqID, "error", err) @@ -234,12 +280,12 @@ func (w *SessionWatcher) RunOnce(ctx context.Context) error { } if ri.Status != 0 { // not Open (Claimed=1, Ready=2, Expired=3) - delete(w.pending, reqID) + w.forget(reqID) continue } if uint64(time.Now().Unix()) > ri.Expiry { - delete(w.pending, reqID) + w.forget(reqID) continue } @@ -261,7 +307,7 @@ func (w *SessionWatcher) RunOnce(ctx context.Context) error { if caps.Sign() != 0 && new(big.Int).And(caps, w.ownCaps).Cmp(caps) != 0 { w.log.Debug("skipping request requiring capabilities this worker lacks", "reqId", reqID, "requiredMask", caps.String()) - delete(w.pending, reqID) + w.forget(reqID) continue } @@ -288,7 +334,7 @@ func (w *SessionWatcher) RunOnce(ctx context.Context) error { // Remove optimistically: if our tx lost the race the request's // status will be non-Open on the next pass and it would be pruned // from pending then anyway. - delete(w.pending, reqID) + w.forget(reqID) } // else: not yet eligible — keep in pending for re-evaluation next pass. } @@ -318,7 +364,7 @@ func (w *SessionWatcher) reseed(ctx context.Context, cursor uint64) error { return err } for _, r := range reqs { - w.pending[r.ReqID] = nil // capability mask fetched lazily at evaluation + w.discover(r) } found += len(reqs) } diff --git a/internal/sortition/session_watcher_test.go b/internal/sortition/session_watcher_test.go index 12b9989..3b220d3 100644 --- a/internal/sortition/session_watcher_test.go +++ b/internal/sortition/session_watcher_test.go @@ -303,6 +303,52 @@ func TestSessionWatcher_DropsExpiredRequest(t *testing.T) { require.NotContains(t, sw.pending, uint64(1), "expired request must be pruned from pending") } +// A request older than MaxRequestAgeBlocks was abandoned by its consumer (the +// consumer-api stops waiting after its claim timeout); claiming it would only +// spend a tx on a session nobody uses. A fresh request on the same pass is +// still claimed. +func TestSessionWatcher_DropsStaleRequest(t *testing.T) { + var counter atomic.Int32 + mc := &mockClaimClient{ + head: chain.HeadInfo{Number: 1000}, + requested: []chain.SessionRequestedEvent{ + {ReqID: 1, BlockNumber: 100, RequestBlock: 100}, // 900 blocks old + {ReqID: 2, BlockNumber: 990, RequestBlock: 990}, // 10 blocks old + }, + eligible: map[uint64]bool{1: true, 2: true}, + } + cs, err := NewCursorStore(t.TempDir()) + require.NoError(t, err) + require.NoError(t, cs.Set(cursorSessionRequested, 0)) + sw := NewSessionWatcher(SessionWatcherOpts{ + Client: mc, + Cursor: cs, + Worker: common.HexToAddress("0x0000000000000000000000000000000000000001"), + JobCounter: &counter, + MaxConcurrent: 4, + ChunkSize: 5000, + Logger: testLogger(t), + MaxRequestAgeBlocks: 150, + }) + require.NoError(t, sw.RunOnce(context.Background())) + require.Equal(t, []uint64{2}, mc.claimed, "only the fresh request is claimed") + require.NotContains(t, sw.pending, uint64(1), "stale request must leave pending") + require.NotContains(t, sw.requestBlock, uint64(1), "stale request bookkeeping must be dropped") +} + +// With the age check off (the default of a zero option) nothing is dropped for age. +func TestSessionWatcher_AgeCheckDisabledByDefault(t *testing.T) { + var counter atomic.Int32 + mc := &mockClaimClient{ + head: chain.HeadInfo{Number: 1000}, + requested: []chain.SessionRequestedEvent{{ReqID: 1, BlockNumber: 100, RequestBlock: 100}}, + eligible: map[uint64]bool{1: true}, + } + sw, _ := newSW(t, mc, &counter) + require.NoError(t, sw.RunOnce(context.Background())) + require.Equal(t, []uint64{1}, mc.claimed) +} + // ────────────────────────────────────────────── // Capability-aware claiming // ────────────────────────────────────────────── diff --git a/workertest/sortition.go b/workertest/sortition.go index 311cd93..f1790a0 100644 --- a/workertest/sortition.go +++ b/workertest/sortition.go @@ -133,7 +133,7 @@ func NewSortitionHarness(t testing.TB, opts SortitionOptions) *SortitionHarness if opts.TavilyURL != "" { searcher = search.NewTavilyClient(opts.TavilyURL, "workertest-key", 10*time.Second) } - desiredCaps := registration.DesiredCapabilities(opts.TavilyURL != "") + desiredCaps := registration.DesiredCapabilities(opts.TavilyURL != "", true, false) capCtx, capCancel := context.WithTimeout(context.Background(), 30*time.Second) registration.NewManager(chainClient, workerAddr, nil, opts.Logger).EnsureCapabilities(capCtx, desiredCaps) ownCaps, err := chainClient.GetWorkerCapabilities(capCtx, workerAddr)