diff --git a/internal/config/config.go b/internal/config/config.go index fc22225..1487058 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -48,6 +48,15 @@ type Config struct { OllamaURL string OllamaTimeout time.Duration + // 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 + // Beacon API (CL node) BeaconAPIURL string @@ -166,6 +175,8 @@ func Load() (*Config, error) { RedisURL: envOrDefault("REDIS_URL", "redis://localhost:6379"), RedisPassword: os.Getenv("REDIS_PASSWORD"), OllamaURL: envOrDefault("OLLAMA_URL", "http://localhost:11434"), + TavilyAPIKey: os.Getenv("TAVILY_API_KEY"), + TavilyURL: envOrDefault("TAVILY_URL", "https://api.tavily.com"), BeaconAPIURL: envOrDefault("BEACON_API_URL", "http://localhost:3500"), SessionKeyFile: envOrDefault("SESSION_KEY_FILE", "data/session-keys.enc"), WorkerGatewayURL: os.Getenv("WORKER_GATEWAY_URL"), @@ -275,6 +286,9 @@ func Load() (*Config, error) { errs = append(errs, fmt.Sprintf("LIGHTCHAIN_DRAIN_SLACK: must be >= 0, got %s", cfg.DrainSlack)) } cfg.OllamaTimeout = parseDuration("OLLAMA_TIMEOUT", "120s", &errs) + cfg.SearchEnabled = parseBool("SEARCH_ENABLED", false, &errs) + cfg.SearchTimeout = parseDuration("SEARCH_TIMEOUT", "10s", &errs) + cfg.SearchMaxResults = parseInt("SEARCH_MAX_RESULTS", 5, &errs) cfg.AckTxTimeout = parseDuration("ACK_TX_TIMEOUT", "15s", &errs) // BlobTxTimeout bounds stage 8 (submit_blob): slot wait + SendTransaction // + WaitMined. Default 90s = ~15 blocks at 6s block time, generous buffer @@ -367,6 +381,17 @@ func (c *Config) Validate() []string { if c.OllamaTimeout <= 0 { errs = append(errs, "OLLAMA_TIMEOUT must be positive") } + if c.SearchEnabled { + if c.TavilyAPIKey == "" { + errs = append(errs, "TAVILY_API_KEY is required when SEARCH_ENABLED=true") + } + if c.SearchTimeout <= 0 { + errs = append(errs, "SEARCH_TIMEOUT must be positive") + } + if c.SearchMaxResults <= 0 { + errs = append(errs, "SEARCH_MAX_RESULTS must be positive") + } + } if c.AckTxTimeout <= 0 { errs = append(errs, "ACK_TX_TIMEOUT must be positive") } diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 68a808a..ea396a2 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -490,6 +490,31 @@ func TestLoad_JobExecutionOverrides(t *testing.T) { assert.Equal(t, 5*time.Second, cfg.ReceiptPollInterval) } +func TestLoad_SearchDefaults(t *testing.T) { + validEnv(t) + t.Setenv("WORKER_KEYSTORE_PATH", "/tmp/k") + t.Setenv("WORKER_KEYSTORE_PASSWORD", "p") + // SEARCH_* unset → search disabled, sane defaults. + cfg, err := Load() + require.NoError(t, err) + assert.False(t, cfg.SearchEnabled) + assert.Equal(t, "https://api.tavily.com", cfg.TavilyURL) + assert.Equal(t, 10*time.Second, cfg.SearchTimeout) + assert.Equal(t, 5, cfg.SearchMaxResults) +} + +func TestLoad_SearchEnabled(t *testing.T) { + validEnv(t) + t.Setenv("WORKER_KEYSTORE_PATH", "/tmp/k") + t.Setenv("WORKER_KEYSTORE_PASSWORD", "p") + t.Setenv("SEARCH_ENABLED", "true") + t.Setenv("TAVILY_API_KEY", "tvly-xxx") + cfg, err := Load() + require.NoError(t, err) + assert.True(t, cfg.SearchEnabled) + assert.Equal(t, "tvly-xxx", cfg.TavilyAPIKey) +} + func TestValidate_JobRegistryZeroAddress(t *testing.T) { validEnv(t) t.Setenv("WORKER_KEYSTORE_PATH", "/tmp/keystore.json") diff --git a/internal/heartbeat/monitor.go b/internal/heartbeat/monitor.go index a3c87ea..324f6a7 100644 --- a/internal/heartbeat/monitor.go +++ b/internal/heartbeat/monitor.go @@ -48,15 +48,16 @@ type MonitorConfig struct { // Monitor publishes periodic heartbeat payloads to Redis. type Monitor struct { - redisClient *redis.Client - cfg MonitorConfig - workerAddr string // EIP-55 checksummed hex, no 0x prefix - modelIDs []string // 0x-prefixed lowercase hex bytes32 - startedAt time.Time - httpClient *http.Client - logger *slog.Logger - jobCounter *atomic.Int32 - maxJobs int + redisClient *redis.Client + cfg MonitorConfig + workerAddr string // EIP-55 checksummed hex, no 0x prefix + modelIDs []string // 0x-prefixed lowercase hex bytes32 + capabilities []string // advertised capability tokens, e.g. ["search"] + startedAt time.Time + httpClient *http.Client + logger *slog.Logger + jobCounter *atomic.Int32 + maxJobs int // metrics is optional — when nil the monitor still writes to Redis but // does not update OllamaUp / HeartbeatLastEmit. Production constructs // always pass non-nil; some tests pass nil to keep them focused. @@ -78,23 +79,25 @@ func NewMonitor( cfg MonitorConfig, workerAddr string, modelIDs []string, + capabilities []string, jobCounter *atomic.Int32, maxJobs int, logger *slog.Logger, metricsCollector *metrics.Metrics, ) *Monitor { return &Monitor{ - redisClient: redisClient, - cfg: cfg, - workerAddr: workerAddr, - modelIDs: modelIDs, - startedAt: time.Now(), - httpClient: &http.Client{Timeout: 2 * time.Second}, - logger: logger, - jobCounter: jobCounter, - maxJobs: maxJobs, - metrics: metricsCollector, - done: make(chan struct{}), + redisClient: redisClient, + cfg: cfg, + workerAddr: workerAddr, + modelIDs: modelIDs, + capabilities: capabilities, + startedAt: time.Now(), + httpClient: &http.Client{Timeout: 2 * time.Second}, + logger: logger, + jobCounter: jobCounter, + maxJobs: maxJobs, + metrics: metricsCollector, + done: make(chan struct{}), } } @@ -162,20 +165,31 @@ func (m *Monitor) emit(ctx context.Context) error { return fmt.Errorf("marshal model IDs: %w", err) } + caps := m.capabilities + if caps == nil { + caps = []string{} + } + capsJSON, err := json.Marshal(caps) + if err != nil { + return fmt.Errorf("marshal capabilities: %w", err) + } + ttl := 3 * m.cfg.Interval key := pkgtypes.HeartbeatRedisKey(m.workerAddr) pipe := m.redisClient.TxPipeline() pipe.HSet(ctx, key, map[string]interface{}{ - pkgtypes.HBFieldLastHeartbeat: time.Now().Unix(), - pkgtypes.HBFieldActiveJobs: activeJobs, - pkgtypes.HBFieldMaxJobs: m.maxJobs, - pkgtypes.HBFieldLatencyMs: 0, - pkgtypes.HBFieldGPUUtil: strconv.FormatFloat(0, 'f', -1, 64), - pkgtypes.HBFieldStatus: pkgtypes.HeartbeatStatusActive, - pkgtypes.HBFieldModels: string(modelsJSON), - pkgtypes.HBFieldOllamaStatus: ollamaStatus, - pkgtypes.HBFieldUptime: int64(time.Since(m.startedAt).Seconds()), + pkgtypes.HBFieldLastHeartbeat: time.Now().Unix(), + pkgtypes.HBFieldActiveJobs: activeJobs, + pkgtypes.HBFieldMaxJobs: m.maxJobs, + pkgtypes.HBFieldLatencyMs: 0, + pkgtypes.HBFieldGPUUtil: strconv.FormatFloat(0, 'f', -1, 64), + pkgtypes.HBFieldStatus: pkgtypes.HeartbeatStatusActive, + pkgtypes.HBFieldModels: string(modelsJSON), + pkgtypes.HBFieldOllamaStatus: ollamaStatus, + pkgtypes.HBFieldUptime: int64(time.Since(m.startedAt).Seconds()), + pkgtypes.HBFieldCapabilities: string(capsJSON), + pkgtypes.HBFieldProtocolVersion: pkgtypes.WorkerProtocolVersion, }) pipe.PExpire(ctx, key, ttl) diff --git a/internal/heartbeat/monitor_test.go b/internal/heartbeat/monitor_test.go index 277ccb0..9a63c89 100644 --- a/internal/heartbeat/monitor_test.go +++ b/internal/heartbeat/monitor_test.go @@ -5,6 +5,7 @@ import ( "log/slog" "net/http" "net/http/httptest" + "strconv" "sync/atomic" "testing" "time" @@ -34,7 +35,7 @@ func newTestMonitor(t *testing.T, mr *miniredis.Miniredis, ollamaURL string) (*M logger := slog.New(slog.NewTextHandler(io.Discard, &slog.HandlerOptions{Level: slog.LevelError})) counter := &atomic.Int32{} - m := NewMonitor(redisClient, cfg, testWorkerAddr, []string{"0xmodel1"}, counter, 0, logger, nil) + m := NewMonitor(redisClient, cfg, testWorkerAddr, []string{"0xmodel1"}, nil, counter, 0, logger, nil) return m, redisClient } @@ -240,6 +241,41 @@ func TestEmitOnce_RedisUnreachable(t *testing.T) { require.Error(t, err, "EmitOnce must return error when Redis is stopped") } +func TestEmit_WritesCapabilities(t *testing.T) { + t.Parallel() + mr := miniredis.RunT(t) + rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()}) + t.Cleanup(func() { _ = rdb.Close() }) + + logger := slog.New(slog.NewTextHandler(io.Discard, &slog.HandlerOptions{Level: slog.LevelError})) + counter := &atomic.Int32{} + m := NewMonitor(rdb, MonitorConfig{Interval: time.Second}, testWorkerAddr, + []string{"0xmodel"}, []string{"search"}, counter, 4, logger, nil) + require.NoError(t, m.EmitOnce(t.Context())) + + key := pkgtypes.HeartbeatRedisKey(testWorkerAddr) + got, err := rdb.HGet(t.Context(), key, pkgtypes.HBFieldCapabilities).Result() + require.NoError(t, err) + assert.JSONEq(t, `["search"]`, got) +} + +func TestEmitWritesProtocolVersion(t *testing.T) { + t.Parallel() + mr := miniredis.RunT(t) + rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()}) + t.Cleanup(func() { _ = rdb.Close() }) + + logger := slog.New(slog.NewTextHandler(io.Discard, &slog.HandlerOptions{Level: slog.LevelError})) + counter := &atomic.Int32{} + m := NewMonitor(rdb, MonitorConfig{Interval: time.Second}, testWorkerAddr, + []string{"0xmodel"}, nil, counter, 0, logger, nil) + require.NoError(t, m.EmitOnce(t.Context())) + + got, err := rdb.HGet(t.Context(), pkgtypes.HeartbeatRedisKey(m.workerAddr), pkgtypes.HBFieldProtocolVersion).Result() + require.NoError(t, err) + require.Equal(t, strconv.Itoa(pkgtypes.WorkerProtocolVersion), got) +} + func TestMonitor_emit_DynamicJobCounts(t *testing.T) { t.Parallel() mr := miniredis.RunT(t) @@ -259,7 +295,7 @@ func TestMonitor_emit_DynamicJobCounts(t *testing.T) { OllamaURL: ollamaSrv.URL, } logger := slog.New(slog.NewTextHandler(io.Discard, &slog.HandlerOptions{Level: slog.LevelError})) - m := NewMonitor(redisClient, cfg, testWorkerAddr, []string{"0xmodel1"}, counter, 5, logger, nil) + m := NewMonitor(redisClient, cfg, testWorkerAddr, []string{"0xmodel1"}, nil, counter, 5, logger, nil) ctx := t.Context() require.NoError(t, m.emit(ctx)) diff --git a/internal/ollama/client.go b/internal/ollama/client.go index 2068dba..9c4d60b 100644 --- a/internal/ollama/client.go +++ b/internal/ollama/client.go @@ -2,6 +2,7 @@ package ollama import ( + "bufio" "bytes" "context" "encoding/json" @@ -12,11 +13,28 @@ import ( "time" ) +// deterministicSeed is the fixed random seed sent to Ollama alongside temperature=0. +// Both the worker (original execution) and the disputer (re-execution) use the same +// seed so that greedy decoding is maximally reproducible across both services. +const deterministicSeed = 42 + +// Options controls Ollama model sampling behaviour. +// Temperature 0 selects the highest-probability token at every step (greedy decoding), +// making inference deterministic. Seed pins any residual nondeterminism. +// NOTE: Temperature has NO omitempty tag — the zero value must be serialised as +// "temperature":0, not omitted (an omitted field lets Ollama fall back to its default +// of 0.8, which reintroduces randomness). +type Options struct { + Temperature float64 `json:"temperature"` + Seed int `json:"seed"` +} + // GenerateRequest is the JSON body sent to POST /api/generate. type GenerateRequest struct { - Model string `json:"model"` - Prompt string `json:"prompt"` - Stream bool `json:"stream"` + Model string `json:"model"` + Prompt string `json:"prompt"` + Stream bool `json:"stream"` + Options Options `json:"options"` } // GenerateResponse is the JSON body returned from POST /api/generate (stream=false). @@ -37,6 +55,7 @@ type ChatRequest struct { Model string `json:"model"` Messages []ChatMessage `json:"messages"` Stream bool `json:"stream"` + Options Options `json:"options"` } // ChatResponse is the JSON body returned from POST /api/chat (stream=false). @@ -76,9 +95,10 @@ func NewOllamaClient(baseURL string, timeout time.Duration) *OllamaClient { // Uses stream=false mode for batch inference (DD-W-1). func (c *OllamaClient) Generate(ctx context.Context, model, prompt string) (string, error) { reqBody := GenerateRequest{ - Model: model, - Prompt: prompt, - Stream: false, + Model: model, + Prompt: prompt, + Stream: false, + Options: Options{Temperature: 0, Seed: deterministicSeed}, } data, err := json.Marshal(reqBody) @@ -122,6 +142,7 @@ func (c *OllamaClient) Chat(ctx context.Context, model string, messages []ChatMe Model: model, Messages: messages, Stream: false, + Options: Options{Temperature: 0, Seed: deterministicSeed}, } data, err := json.Marshal(reqBody) @@ -158,6 +179,122 @@ func (c *OllamaClient) Chat(ctx context.Context, model string, messages []ChatMe return chatResp.Message.Content, nil } +// GenerateStream sends a prompt to the Ollama server with stream=true, calling +// onDelta for each non-empty token chunk. Returns the full accumulated response. +func (c *OllamaClient) GenerateStream(ctx context.Context, model, prompt string, onDelta func(string)) (string, error) { + reqBody := GenerateRequest{ + Model: model, + Prompt: prompt, + Stream: true, + Options: Options{Temperature: 0, Seed: deterministicSeed}, + } + + data, err := json.Marshal(reqBody) + if err != nil { + return "", fmt.Errorf("marshal generate stream request: %w", err) + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/api/generate", bytes.NewReader(data)) + if err != nil { + return "", fmt.Errorf("create generate stream request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + + resp, err := c.httpClient.Do(req) + if err != nil { + return "", fmt.Errorf("ollama generate stream request: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(resp.Body) + return "", fmt.Errorf("ollama generate stream returned status %d: %s", resp.StatusCode, string(body)) + } + + var full strings.Builder + scanner := bufio.NewScanner(resp.Body) + for scanner.Scan() { + line := scanner.Bytes() + if len(line) == 0 { + continue + } + var chunk GenerateResponse + if err := json.Unmarshal(line, &chunk); err != nil { + return "", fmt.Errorf("unmarshal generate stream chunk: %w", err) + } + if chunk.Response != "" { + onDelta(chunk.Response) + full.WriteString(chunk.Response) + } + if chunk.Done { + break + } + } + if err := scanner.Err(); err != nil { + return "", fmt.Errorf("read generate stream: %w", err) + } + + return full.String(), nil +} + +// ChatStream sends a multi-turn conversation to the Ollama server with stream=true, +// calling onDelta for each non-empty token chunk. Returns the full accumulated response. +func (c *OllamaClient) ChatStream(ctx context.Context, model string, messages []ChatMessage, onDelta func(string)) (string, error) { + reqBody := ChatRequest{ + Model: model, + Messages: messages, + Stream: true, + Options: Options{Temperature: 0, Seed: deterministicSeed}, + } + + data, err := json.Marshal(reqBody) + if err != nil { + return "", fmt.Errorf("marshal chat stream request: %w", err) + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/api/chat", bytes.NewReader(data)) + if err != nil { + return "", fmt.Errorf("create chat stream request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + + resp, err := c.httpClient.Do(req) + if err != nil { + return "", fmt.Errorf("ollama chat stream request: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(resp.Body) + return "", fmt.Errorf("ollama chat stream returned status %d: %s", resp.StatusCode, string(body)) + } + + var full strings.Builder + scanner := bufio.NewScanner(resp.Body) + for scanner.Scan() { + line := scanner.Bytes() + if len(line) == 0 { + continue + } + var chunk ChatResponse + if err := json.Unmarshal(line, &chunk); err != nil { + return "", fmt.Errorf("unmarshal chat stream chunk: %w", err) + } + if chunk.Message.Content != "" { + onDelta(chunk.Message.Content) + full.WriteString(chunk.Message.Content) + } + if chunk.Done { + break + } + } + if err := scanner.Err(); err != nil { + return "", fmt.Errorf("read chat stream: %w", err) + } + + return full.String(), nil +} + // VerifyModels checks that all required models are loaded on the Ollama server. // Returns an error listing any missing models. func (c *OllamaClient) VerifyModels(ctx context.Context, models []string) error { diff --git a/internal/ollama/client_test.go b/internal/ollama/client_test.go index 5500f44..a62407b 100644 --- a/internal/ollama/client_test.go +++ b/internal/ollama/client_test.go @@ -3,6 +3,7 @@ package ollama import ( "context" "encoding/json" + "io" "net/http" "net/http/httptest" "testing" @@ -75,6 +76,146 @@ func TestGenerate_ContextCancelled(t *testing.T) { require.Error(t, err) } +func TestGenerateStream(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, http.MethodPost, r.Method) + assert.Equal(t, "/api/generate", r.URL.Path) + + var req GenerateRequest + require.NoError(t, json.NewDecoder(r.Body).Decode(&req)) + assert.True(t, req.Stream) + + w.Header().Set("Content-Type", "application/x-ndjson") + io.WriteString(w, `{"response":"Hel","done":false}`+"\n") + io.WriteString(w, `{"response":"lo","done":false}`+"\n") + io.WriteString(w, `{"response":"","done":true}`+"\n") + })) + defer srv.Close() + + c := NewOllamaClient(srv.URL, 5*time.Second) + var deltas []string + full, err := c.GenerateStream(context.Background(), "llama3-8b", "hi", func(d string) { deltas = append(deltas, d) }) + require.NoError(t, err) + assert.Equal(t, []string{"Hel", "lo"}, deltas) + assert.Equal(t, "Hello", full) +} + +func TestChatStream(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, http.MethodPost, r.Method) + assert.Equal(t, "/api/chat", r.URL.Path) + + var req ChatRequest + require.NoError(t, json.NewDecoder(r.Body).Decode(&req)) + assert.True(t, req.Stream) + + w.Header().Set("Content-Type", "application/x-ndjson") + io.WriteString(w, `{"message":{"role":"assistant","content":"Hel"},"done":false}`+"\n") + io.WriteString(w, `{"message":{"role":"assistant","content":"lo"},"done":false}`+"\n") + io.WriteString(w, `{"message":{"role":"assistant","content":""},"done":true}`+"\n") + })) + defer srv.Close() + + c := NewOllamaClient(srv.URL, 5*time.Second) + var deltas []string + full, err := c.ChatStream(context.Background(), "llama3-8b", []ChatMessage{{Role: "user", Content: "hi"}}, func(d string) { deltas = append(deltas, d) }) + require.NoError(t, err) + assert.Equal(t, []string{"Hel", "lo"}, deltas) + assert.Equal(t, "Hello", full) +} + +// TestDeterministicOptions verifies that every request path serialises +// "temperature":0 (not omitted) and "seed":42 so Ollama uses greedy decoding +// and the worker's output is reproducible by the disputer. +func TestDeterministicOptions(t *testing.T) { + t.Parallel() + + assertOptions := func(t *testing.T, body []byte) { + t.Helper() + bodyStr := string(body) + assert.Contains(t, bodyStr, `"temperature":0`, "temperature must be serialised as 0, not omitted") + assert.Contains(t, bodyStr, `"seed":42`, "seed must be 42") + assert.Contains(t, bodyStr, `"options"`, "options object must be present") + } + + t.Run("Generate", func(t *testing.T) { + t.Parallel() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, err := io.ReadAll(r.Body) + require.NoError(t, err) + assertOptions(t, body) + json.NewEncoder(w).Encode(GenerateResponse{Response: "ok", Done: true}) + })) + defer srv.Close() + c := NewOllamaClient(srv.URL, 5*time.Second) + _, err := c.Generate(context.Background(), "llama3-8b", "hi") + require.NoError(t, err) + }) + + t.Run("Chat", func(t *testing.T) { + t.Parallel() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, err := io.ReadAll(r.Body) + require.NoError(t, err) + assertOptions(t, body) + json.NewEncoder(w).Encode(ChatResponse{Message: ChatMessage{Role: "assistant", Content: "ok"}, Done: true}) + })) + defer srv.Close() + c := NewOllamaClient(srv.URL, 5*time.Second) + _, err := c.Chat(context.Background(), "llama3-8b", []ChatMessage{{Role: "user", Content: "hi"}}) + require.NoError(t, err) + }) + + t.Run("GenerateStream", func(t *testing.T) { + t.Parallel() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, err := io.ReadAll(r.Body) + require.NoError(t, err) + assertOptions(t, body) + w.Header().Set("Content-Type", "application/x-ndjson") + io.WriteString(w, `{"response":"ok","done":true}`+"\n") + })) + defer srv.Close() + c := NewOllamaClient(srv.URL, 5*time.Second) + _, err := c.GenerateStream(context.Background(), "llama3-8b", "hi", func(string) {}) + require.NoError(t, err) + }) + + t.Run("ChatStream", func(t *testing.T) { + t.Parallel() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, err := io.ReadAll(r.Body) + require.NoError(t, err) + assertOptions(t, body) + w.Header().Set("Content-Type", "application/x-ndjson") + io.WriteString(w, `{"message":{"role":"assistant","content":"ok"},"done":true}`+"\n") + })) + defer srv.Close() + c := NewOllamaClient(srv.URL, 5*time.Second) + _, err := c.ChatStream(context.Background(), "llama3-8b", []ChatMessage{{Role: "user", Content: "hi"}}, func(string) {}) + require.NoError(t, err) + }) + + t.Run("OptionsStructMarshal", func(t *testing.T) { + t.Parallel() + req := GenerateRequest{ + Model: "llama3-8b", + Prompt: "test", + Stream: false, + Options: Options{Temperature: 0, Seed: deterministicSeed}, + } + b, err := json.Marshal(req) + require.NoError(t, err) + bodyStr := string(b) + assert.Contains(t, bodyStr, `"temperature":0`, "temperature 0 must not be omitted") + assert.Contains(t, bodyStr, `"seed":42`) + }) +} + func TestVerifyModels_AllPresent(t *testing.T) { t.Parallel() diff --git a/internal/pipeline/handler.go b/internal/pipeline/handler.go index a47a67b..de2f335 100644 --- a/internal/pipeline/handler.go +++ b/internal/pipeline/handler.go @@ -23,10 +23,12 @@ import ( "golang.org/x/sync/singleflight" pkgcrypto "github.com/lightchain/pkg/crypto" + "github.com/lightchain/pkg/searchaug" pkgtypes "github.com/lightchain/pkg/types" "github.com/lightchain/worker/internal/metrics" "github.com/lightchain/worker/internal/ollama" + "github.com/lightchain/worker/internal/search" ) // SessionKeyGetter retrieves and stores session keys. @@ -39,6 +41,8 @@ type SessionKeyGetter interface { type InferenceClient interface { Generate(ctx context.Context, model, prompt string) (string, error) Chat(ctx context.Context, model string, messages []ollama.ChatMessage) (string, error) + GenerateStream(ctx context.Context, model, prompt string, onDelta func(string)) (string, error) + ChatStream(ctx context.Context, model string, messages []ollama.ChatMessage, onDelta func(string)) (string, error) } // JobExecutionClient submits job lifecycle transactions on-chain. @@ -47,7 +51,7 @@ type JobExecutionClient interface { // CompleteJob submits a completeJob TX with a single bytes32 response blob hash. // The contract enforces `blobhash(0) == responseBlobHash`, so the caller must // have submitted exactly one blob in the same (blob-carrying) transaction. - CompleteJob(ctx context.Context, jobID uint64, responseBlobHash [32]byte, responseCiphertextHash [32]byte) error + CompleteJob(ctx context.Context, jobID uint64, responseBlobHash, responseCiphertextHash [32]byte) error HasJobAcknowledged(ctx context.Context, jobID uint64) (bool, error) HasJobCompleted(ctx context.Context, jobID uint64) (bool, error) // GetSessionEncWorkerKey retrieves the current encrypted worker key for a session @@ -57,7 +61,7 @@ type JobExecutionClient interface { // GetJobBlobInfo returns the single prompt and response blob hashes for a // completed job along with the blocks they were submitted in. Used to build // conversation history. - GetJobBlobInfo(ctx context.Context, jobID uint64) (promptHash common.Hash, responseHash common.Hash, submitBlock uint64, completionBlock uint64, err error) + GetJobBlobInfo(ctx context.Context, jobID uint64) (promptHash, responseHash common.Hash, submitBlock, completionBlock uint64, err error) } // BlobFetcher fetches EIP-4844 blob data from the consensus layer. @@ -78,13 +82,17 @@ type HandlerConfig struct { ModelIDToName map[string]string ChainID *big.Int JobRegistryAddr common.Address + SearchMaxResults int + SearchTimeout time.Duration } // ResponsePublisher publishes encrypted responses for real-time delivery. // Implementations: RedisResponsePublisher (direct Redis PUBLISH) and // gateway-based publisher (POST to worker-gateway). type ResponsePublisher interface { - PublishResponse(ctx context.Context, jobID, sessionID uint64, correlationID string, signature string, ciphertext []byte) + PublishResponse(ctx context.Context, jobID, sessionID uint64, correlationID, signature string, ciphertext []byte) + PublishMetadata(ctx context.Context, jobID, sessionID uint64, correlationID string, payload []byte) + PublishChunk(ctx context.Context, jobID, sessionID uint64, correlationID string, sequence uint32, payload []byte) } // ReleaseTracker records that a job has just been completed and is now @@ -103,6 +111,7 @@ type JobHandler struct { blobSubmitter BlobSubmitter keyStore SessionKeyGetter ollamaClient InferenceClient + searcher search.Searcher // nil ⇒ web search disabled redisClient *redis.Client responsePublisher ResponsePublisher signingKey *ecdsa.PrivateKey @@ -164,6 +173,7 @@ func NewJobHandler( checkpoints *CheckpointStore, metricsCollector *metrics.Metrics, delivery string, + searcher search.Searcher, ) *JobHandler { modelIDToName := make(map[string]string, len(cfg.ModelIDToName)) for k, v := range cfg.ModelIDToName { @@ -185,6 +195,7 @@ func NewJobHandler( blobSubmitter: blobSubmitter, keyStore: keyStore, ollamaClient: ollamaClient, + searcher: searcher, redisClient: redisClient, responsePublisher: publisher, signingKey: signingKey, @@ -274,13 +285,86 @@ func (p *RedisResponsePublisher) PublishResponse( } } +// PublishMetadata fans out an encrypted metadata frame (e.g. web-search +// sources). Non-fatal — citations are best-effort UX, not the authoritative +// response. No signature: the relay only signature-checks complete frames. +func (p *RedisResponsePublisher) PublishMetadata( + ctx context.Context, + jobID, sessionID uint64, + correlationID string, + payload []byte, +) { + msg := pkgtypes.PubSubMessage{ + Type: pkgtypes.MessageTypeMetadata, + JobID: pkgtypes.JobID(jobID), + SessionID: pkgtypes.SessionID(sessionID), + Sequence: 0, + TotalChunks: 1, + Payload: payload, + CorrelationID: correlationID, + Timestamp: time.Now().Unix(), + } + data, err := json.Marshal(msg) + if err != nil { + p.logger.Warn("failed to marshal metadata frame", "jobID", jobID, "error", err) + return + } + pubCtx := ctx + if p.publishTimeout > 0 { + var cancel context.CancelFunc + pubCtx, cancel = context.WithTimeout(ctx, p.publishTimeout) + defer cancel() + } + channel := fmt.Sprintf("session:%d:responses", sessionID) + if err := p.client.Publish(pubCtx, channel, data).Err(); err != nil { + p.logger.Warn("failed to publish metadata frame", "jobID", jobID, "channel", channel, "error", err) + } +} + +// PublishChunk fans out an encrypted streaming chunk frame (best-effort UX). +// Non-fatal — chunk delivery failures are logged and skipped. +func (p *RedisResponsePublisher) PublishChunk( + ctx context.Context, + jobID, sessionID uint64, + correlationID string, + sequence uint32, + payload []byte, +) { + msg := pkgtypes.PubSubMessage{ + Type: pkgtypes.MessageTypeChunk, + JobID: pkgtypes.JobID(jobID), + SessionID: pkgtypes.SessionID(sessionID), + Sequence: sequence, + TotalChunks: 0, + Payload: payload, + CorrelationID: correlationID, + Timestamp: time.Now().Unix(), + } + data, err := json.Marshal(msg) + if err != nil { + p.logger.Warn("failed to marshal chunk frame", "jobID", jobID, "seq", sequence, "error", err) + return + } + pubCtx := ctx + if p.publishTimeout > 0 { + var cancel context.CancelFunc + pubCtx, cancel = context.WithTimeout(ctx, p.publishTimeout) + defer cancel() + } + channel := fmt.Sprintf("session:%d:responses", sessionID) + if err := p.client.Publish(pubCtx, channel, data).Err(); err != nil { + p.logger.Warn("failed to publish chunk frame", "jobID", jobID, "seq", sequence, "channel", channel, "error", err) + } +} + // HandleJobPayload processes a job from a raw JobPayload (used in gateway mode // where jobs come via HTTP polling instead of Asynq). func (h *JobHandler) HandleJobPayload(ctx context.Context, payload JobPayload) error { h.jobCounter.Add(1) defer h.jobCounter.Add(-1) - h.logger.Info("processing job", + h.logger.Info( + "processing job", "jobID", payload.JobID, "sessionID", payload.SessionID, "model", payload.ModelID, @@ -324,7 +408,8 @@ func (h *JobHandler) HandleTask(ctx context.Context, task *asynq.Task) error { budget = time.Until(deadline) } - h.logger.Info("processing job", + h.logger.Info( + "processing job", "jobID", payload.JobID, "sessionID", payload.SessionID, "model", payload.ModelID, @@ -345,7 +430,8 @@ func (h *JobHandler) HandleTask(ctx context.Context, task *asynq.Task) error { if hasDeadline { minBudget := h.cfg.AckTxTimeout + h.cfg.BlobTxTimeout + 10*time.Second if budget < minBudget { - h.logger.Warn("task budget appears too small for the pipeline", + h.logger.Warn( + "task budget appears too small for the pipeline", "jobID", payload.JobID, "taskBudgetMs", budget.Milliseconds(), "minRecommendedMs", minBudget.Milliseconds(), @@ -357,7 +443,8 @@ func (h *JobHandler) HandleTask(ctx context.Context, task *asynq.Task) error { } if err := h.processJob(ctx, payload); err != nil { - h.logger.Error("job failed", + h.logger.Error( + "job failed", "jobID", payload.JobID, "attempt", retryCount+1, "maxRetry", maxRetry, @@ -367,7 +454,8 @@ func (h *JobHandler) HandleTask(ctx context.Context, task *asynq.Task) error { return err } - h.logger.Info("job completed", + h.logger.Info( + "job completed", "jobID", payload.JobID, "attempt", retryCount+1, ) @@ -429,7 +517,8 @@ func (h *JobHandler) processJob(ctx context.Context, p JobPayload) (err error) { return err } d := rec.End(metrics.OutcomeOK, metrics.CacheNone) - logger.Info("stage 1 complete", + logger.Info( + "stage 1 complete", "stage", "ack", "durationMs", d.Milliseconds(), ) @@ -445,7 +534,8 @@ func (h *JobHandler) processJob(ctx context.Context, p JobPayload) (err error) { if ckpt.HasCiphertext() { cacheState = metrics.CacheHit h.metrics.CheckpointEvents.WithLabelValues(metrics.CheckpointEventHitInference).Inc() - logger.Info("checkpoint hit, skipping stages 2-6", + logger.Info( + "checkpoint hit, skipping stages 2-6", "stage", "checkpoint", "reason", "inference_cached", "ciphertextBytes", len(ckpt.Ciphertext), @@ -461,7 +551,8 @@ func (h *JobHandler) processJob(ctx context.Context, p JobPayload) (err error) { return err } rec = h.metrics.StartStage(metrics.StageFetchBlob, model, delivery) - logger.Info("stage 2 starting", + logger.Info( + "stage 2 starting", "stage", "fetch_blob", "promptBlobHash", p.PromptBlobHash.Hex(), "blockNumber", p.BlockNumber, @@ -473,7 +564,8 @@ func (h *JobHandler) processJob(ctx context.Context, p JobPayload) (err error) { return err } d = rec.End(metrics.OutcomeOK, metrics.CacheMiss) - logger.Info("stage 2 complete", + logger.Info( + "stage 2 complete", "stage", "fetch_blob", "blobBytes", len(blobData), "durationMs", d.Milliseconds(), @@ -497,20 +589,23 @@ func (h *JobHandler) processJob(ctx context.Context, p JobPayload) (err error) { canonical, wasSet, cErr := h.checkpoints.SetCiphertextIfAbsent(ctx, p.JobID, ciphertext) switch { case cErr != nil: - logger.Warn("failed to persist ciphertext to checkpoint", + logger.Warn( + "failed to persist ciphertext to checkpoint", "stage", "checkpoint", "error", cErr, ) case !wasSet: h.metrics.CheckpointEvents.WithLabelValues(metrics.CheckpointEventRaceLost).Inc() - logger.Warn("checkpoint race lost, using canonical ciphertext", + logger.Warn( + "checkpoint race lost, using canonical ciphertext", "stage", "checkpoint", "localBytes", len(ciphertext), "canonicalBytes", len(canonical), ) ciphertext = canonical default: - logger.Debug("ciphertext persisted to checkpoint", + logger.Debug( + "ciphertext persisted to checkpoint", "stage", "checkpoint", "ciphertextBytes", len(ciphertext), ) @@ -523,21 +618,33 @@ func (h *JobHandler) processJob(ctx context.Context, p JobPayload) (err error) { // a duplicate complete frame. if ckpt.Delivered { h.metrics.CheckpointEvents.WithLabelValues(metrics.CheckpointEventHitDelivered).Inc() - logger.Info("checkpoint hit, skipping stage 7", + logger.Info( + "checkpoint hit, skipping stage 7", "stage", "checkpoint", "reason", "delivered", ) } else { rec = h.metrics.StartStage(metrics.StageRedisPublish, model, delivery) - h.publishToRedis(ctx, logger, p.JobID, p.SessionID, p.CorrelationID, ciphertext) + // completePayload and ciphertext intentionally differ for v2 search jobs: + // the relay complete frame carries the PLAIN answer (consumer contract), + // while ciphertext stays the {answer,searchContext} envelope used by the + // blob + on-chain responseCiphertextHash (the disputer reads it). Do NOT + // unify these — see relayCompleteCiphertext. + completePayload := ciphertext + if sk, skErr := h.getOrDeriveSessionKey(ctx, logger, p.SessionID); skErr == nil { + completePayload = h.relayCompleteCiphertext(sk, ciphertext) + } // on key error, fall back to ciphertext (non-fatal; matches existing publish best-effort posture) + h.publishToRedis(ctx, logger, p.JobID, p.SessionID, p.CorrelationID, completePayload) d = rec.End(metrics.OutcomeOK, metrics.CacheMiss) - logger.Info("stage 7 complete", + logger.Info( + "stage 7 complete", "stage", "redis_publish", "durationMs", d.Milliseconds(), ) if h.checkpoints != nil { if mErr := h.checkpoints.MarkDelivered(ctx, p.JobID); mErr != nil { - logger.Warn("failed to mark delivered on checkpoint", + logger.Warn( + "failed to mark delivered on checkpoint", "stage", "checkpoint", "error", mErr, ) @@ -555,7 +662,8 @@ func (h *JobHandler) processJob(ctx context.Context, p JobPayload) (err error) { // Stage 8b: Complete job on-chain. rec = h.metrics.StartStage(metrics.StageCompleteJob, model, delivery) responseCiphertextHash := crypto.Keccak256Hash(ciphertext) - logger.Info("stage 8b starting", + logger.Info( + "stage 8b starting", "stage", "complete_job", "versionedHash", versionedHash.Hex(), "ciphertextHash", responseCiphertextHash.Hex(), @@ -566,7 +674,8 @@ func (h *JobHandler) processJob(ctx context.Context, p JobPayload) (err error) { return err } d = rec.End(metrics.OutcomeOK, metrics.CacheNone) - logger.Info("stage 8b complete", + logger.Info( + "stage 8b complete", "stage", "complete_job", "durationMs", d.Milliseconds(), ) @@ -581,7 +690,8 @@ func (h *JobHandler) processJob(ctx context.Context, p JobPayload) (err error) { // an already-completed on-chain job). Reconciler backs us up. if h.releaseTracker != nil { if mErr := h.releaseTracker.MarkEligible(ctx, p.JobID, time.Now().Unix()); mErr != nil { - logger.Warn("failed to mark job eligible for release; reconciler will backfill", + logger.Warn( + "failed to mark job eligible for release; reconciler will backfill", "stage", "release_tracker", "error", mErr, ) @@ -593,7 +703,8 @@ func (h *JobHandler) processJob(ctx context.Context, p JobPayload) (err error) { // disappears; outright Delete would race concurrent retries. if h.checkpoints != nil { if tErr := h.checkpoints.Tombstone(ctx, p.JobID); tErr != nil { - logger.Warn("failed to tombstone checkpoint after completion", + logger.Warn( + "failed to tombstone checkpoint after completion", "stage", "checkpoint", "error", tErr, ) @@ -614,7 +725,8 @@ func (h *JobHandler) readCheckpoint(ctx context.Context, logger *slog.Logger, jo } ckpt, err := h.checkpoints.Get(ctx, jobID) if err != nil { - logger.Warn("failed to read checkpoint; treating as cache miss", + logger.Warn( + "failed to read checkpoint; treating as cache miss", "stage", "checkpoint", "error", err, ) @@ -643,7 +755,8 @@ func (h *JobHandler) runInferencePipeline( return nil, fmt.Errorf("stage 3 (session key): %w", err) } d := rec.End(metrics.OutcomeOK, metrics.CacheMiss) - logger.Info("stage 3 complete", + logger.Info( + "stage 3 complete", "stage", "session_key", "durationMs", d.Milliseconds(), ) @@ -657,7 +770,8 @@ func (h *JobHandler) runInferencePipeline( rec = h.metrics.StartStage(metrics.StageDecrypt, model, delivery) prompt, err := pkgcrypto.Decrypt(sessionKey, blobData) if err != nil { - logger.Warn("stage 4: initial decrypt failed, refreshing session key", + logger.Warn( + "stage 4: initial decrypt failed, refreshing session key", "stage", "decrypt", "error", err, ) @@ -674,12 +788,45 @@ func (h *JobHandler) runInferencePipeline( } } d = rec.End(metrics.OutcomeOK, metrics.CacheMiss) - logger.Info("stage 4 complete", + logger.Info( + "stage 4 complete", "stage", "decrypt", "promptBytes", len(prompt), "durationMs", d.Milliseconds(), ) + // Stage 4.5: optional web-search augmentation (one-shot Tavily). Fail-open: + // any search failure proceeds with the original prompt and emits no sources. + // Sources are stashed here and published AFTER inference (see end of function) + // so the UI renders "Sources" beneath the answer, not above it. + promptText, searchEnabled, decErr := searchaug.DecodePrompt(prompt) + if decErr != nil { + return nil, fmt.Errorf("stage 4 (decode prompt envelope): %w", decErr) + } + var searchSources []searchaug.Source + if searchEnabled && h.searcher != nil { + searchCtx := ctx + if h.cfg.SearchTimeout > 0 { + var cancel context.CancelFunc + searchCtx, cancel = context.WithTimeout(ctx, h.cfg.SearchTimeout) + defer cancel() + } + maxResults := h.cfg.SearchMaxResults + if maxResults <= 0 { + maxResults = 5 + } + sources, sErr := h.searcher.Search(searchCtx, promptText, maxResults) + if sErr != nil { + logger.Warn("stage 4.5: web search failed, proceeding without context", + "stage", "search", "error", sErr) + } else if len(sources) > 0 { + augSources := toSearchaugSources(sources) + promptText = searchaug.BuildAugmentedPrompt(searchaug.CurrentTemplateVersion, promptText, augSources) + searchSources = augSources // stash; publish after inference + logger.Info("stage 4.5 complete", "stage", "search", "sources", len(sources)) + } + } + // Stage 5: AI inference (with conversation history if prior jobs exist) modelName, err := h.resolveModelName(p.ModelID) if err != nil { @@ -692,7 +839,8 @@ func (h *JobHandler) runInferencePipeline( var hErr error history, hErr = h.buildConversationHistory(ctx, p.PriorJobIDs, sessionKey) if hErr != nil { - logger.Warn("failed to build conversation history, falling back to single prompt", + logger.Warn( + "failed to build conversation history, falling back to single prompt", "stage", "inference", "error", hErr, ) @@ -701,48 +849,97 @@ func (h *JobHandler) runInferencePipeline( } rec = h.metrics.StartStage(metrics.StageInference, model, delivery) - logger.Info("stage 5 starting", + logger.Info( + "stage 5 starting", "stage", "inference", "model", modelName, "promptBytes", len(prompt), "historyTurns", len(history), ) + + // Streaming chunk batching: flush when buffer reaches 24 bytes or at stream end. + // seq is monotonically increasing from 1. Failures are non-fatal (best-effort UX). + var seq uint32 + var chunkBuf strings.Builder + flushChunk := func() { + if chunkBuf.Len() == 0 { + return + } + seq++ + if h.responsePublisher != nil { + if enc, encErr := pkgcrypto.Encrypt(sessionKey, []byte(chunkBuf.String())); encErr == nil { + h.responsePublisher.PublishChunk(ctx, p.JobID, p.SessionID, p.CorrelationID, seq, enc) + } else { + logger.Warn("chunk encrypt failed, skipping chunk", "seq", seq, "error", encErr) + } + } + chunkBuf.Reset() + } + onDelta := func(d string) { + chunkBuf.WriteString(d) + if chunkBuf.Len() >= 24 { + flushChunk() + } + } + if len(history) > 0 { - messages := append(history, ollama.ChatMessage{Role: "user", Content: string(prompt)}) - response, err = h.ollamaClient.Chat(ctx, modelName, messages) + messages := append(history, ollama.ChatMessage{Role: "user", Content: promptText}) + response, err = h.ollamaClient.ChatStream(ctx, modelName, messages, onDelta) if err != nil { rec.End(metrics.OutcomeError, metrics.CacheMiss) return nil, fmt.Errorf("stage 5 (chat inference): %w", err) } } else { - response, err = h.ollamaClient.Generate(ctx, modelName, string(prompt)) + response, err = h.ollamaClient.GenerateStream(ctx, modelName, promptText, onDelta) if err != nil { rec.End(metrics.OutcomeError, metrics.CacheMiss) return nil, fmt.Errorf("stage 5 (inference): %w", err) } } + flushChunk() // flush any remaining buffered delta d = rec.End(metrics.OutcomeOK, metrics.CacheMiss) - logger.Info("stage 5 complete", + logger.Info( + "stage 5 complete", "stage", "inference", "model", modelName, "responseBytes", len(response), "durationMs", d.Milliseconds(), ) - // Stage 6: Encrypt response + // Stage 6: Encode and encrypt response. For search jobs the response is + // wrapped in a v2 JSON envelope that captures the search context so the + // 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. rec = h.metrics.StartStage(metrics.StageEncrypt, model, delivery) - ciphertext, err := pkgcrypto.Encrypt(sessionKey, []byte(response)) + respBytes, encErr := searchaug.EncodeResponse(response, searchSources) + if encErr != nil { + rec.End(metrics.OutcomeError, metrics.CacheMiss) + return nil, 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, fmt.Errorf("stage 6 (encrypt response): %w", err) } d = rec.End(metrics.OutcomeOK, metrics.CacheMiss) - logger.Info("stage 6 complete", + logger.Info( + "stage 6 complete", "stage", "encrypt", "ciphertextBytes", len(ciphertext), "durationMs", d.Milliseconds(), ) + // Emit citations AFTER the answer is ready so the UI renders Sources + // beneath the response (best-effort, non-fatal). + if len(searchSources) > 0 && h.responsePublisher != nil { + if metaPayload, encErr := pkgcrypto.Encrypt(sessionKey, sourcesMetadataJSON(searchSources)); encErr != nil { + logger.Warn("post-inference: encrypt sources failed", "stage", "search", "error", encErr) + } else { + h.responsePublisher.PublishMetadata(ctx, p.JobID, p.SessionID, p.CorrelationID, metaPayload) + } + } + return ciphertext, nil } @@ -752,14 +949,16 @@ func (h *JobHandler) ensureAcknowledged(ctx context.Context, logger *slog.Logger return fmt.Errorf("check acknowledged state: %w", err) } if acknowledged { - logger.Info("stage 1: job already acknowledged on-chain, skipping ack tx", + logger.Info( + "stage 1: job already acknowledged on-chain, skipping ack tx", "stage", "ack", "path", "already_acked", ) return nil } - logger.Info("stage 1: sending ack tx", + logger.Info( + "stage 1: sending ack tx", "stage", "ack", "path", "sending_ack", "timeout", h.cfg.AckTxTimeout.String(), @@ -774,7 +973,8 @@ func (h *JobHandler) ensureAcknowledged(ctx context.Context, logger *slog.Logger return fmt.Errorf("acknowledge job: %w (recheck failed: %v)", err, checkErr) } if acknowledged { - logger.Warn("stage 1: ack tx returned error but job is acknowledged on-chain", + logger.Warn( + "stage 1: ack tx returned error but job is acknowledged on-chain", "stage", "ack", "error", err, ) @@ -809,7 +1009,8 @@ func (h *JobHandler) ensureBlobSubmitted( rec := h.metrics.StartStage(metrics.StageSubmitBlob, model, delivery) rec.End(metrics.OutcomeSkipped, metrics.CacheHit) h.metrics.CheckpointEvents.WithLabelValues(metrics.CheckpointEventHitBlob).Inc() - logger.Info("checkpoint hit, skipping stage 8a", + logger.Info( + "checkpoint hit, skipping stage 8a", "stage", "checkpoint", "reason", "blob_submitted", "versionedHash", ckpt.VersionedHash.Hex(), @@ -818,7 +1019,8 @@ func (h *JobHandler) ensureBlobSubmitted( } rec := h.metrics.StartStage(metrics.StageSubmitBlob, model, delivery) - logger.Info("stage 8a starting", + logger.Info( + "stage 8a starting", "stage", "submit_blob", "ciphertextBytes", len(ciphertext), "timeout", h.cfg.BlobTxTimeout.String(), @@ -836,14 +1038,16 @@ func (h *JobHandler) ensureBlobSubmitted( } versionedHash := common.Hash(blobHashes[0]) d := rec.End(metrics.OutcomeOK, metrics.CacheMiss) - logger.Info("stage 8a complete", + logger.Info( + "stage 8a complete", "stage", "submit_blob", "versionedHash", versionedHash.Hex(), "durationMs", d.Milliseconds(), ) if h.checkpoints != nil { if err := h.checkpoints.SetVersionedHash(ctx, jobID, versionedHash); err != nil { - logger.Warn("failed to persist versionedHash on checkpoint", + logger.Warn( + "failed to persist versionedHash on checkpoint", "stage", "checkpoint", "error", err, ) @@ -865,7 +1069,8 @@ func (h *JobHandler) completeJob( return fmt.Errorf("complete job: %w (recheck failed: %v)", err, checkErr) } if completed { - logger.Warn("stage 8b: completeJob tx returned error but job is completed on-chain", + logger.Warn( + "stage 8b: completeJob tx returned error but job is completed on-chain", "stage", "complete_job", "error", err, ) @@ -887,7 +1092,8 @@ func (h *JobHandler) completeJob( func (h *JobHandler) getOrDeriveSessionKey(ctx context.Context, logger *slog.Logger, sessionID uint64) ([]byte, error) { if key, err := h.keyStore.GetKey(sessionID); err == nil { h.metrics.SessionKeyEvents.WithLabelValues(metrics.SessionKeyPathCacheHit).Inc() - logger.Info("stage 3: session key cache hit", + logger.Info( + "stage 3: session key cache hit", "stage", "session_key", "path", "cache_hit", ) @@ -901,14 +1107,16 @@ func (h *JobHandler) getOrDeriveSessionKey(ctx context.Context, logger *slog.Log // just finished between our outer miss and entering Do. if k, err := h.keyStore.GetKey(sessionID); err == nil { h.metrics.SessionKeyEvents.WithLabelValues(metrics.SessionKeyPathCacheHit).Inc() - logger.Info("stage 3: session key cache hit (inside flight)", + logger.Info( + "stage 3: session key cache hit (inside flight)", "stage", "session_key", "path", "cache_hit", ) return k, nil } h.metrics.SessionKeyEvents.WithLabelValues(metrics.SessionKeyPathChainDerive).Inc() - logger.Info("stage 3: session key cache miss, deriving from chain", + logger.Info( + "stage 3: session key cache miss, deriving from chain", "stage", "session_key", "path", "chain_derive", ) @@ -930,7 +1138,8 @@ func (h *JobHandler) getOrDeriveSessionKey(ctx context.Context, logger *slog.Log // session key. func (h *JobHandler) refreshSessionKey(ctx context.Context, logger *slog.Logger, sessionID uint64) ([]byte, error) { h.metrics.SessionKeyEvents.WithLabelValues(metrics.SessionKeyPathRefresh).Inc() - logger.Info("stage 4: refreshing session key from chain (possible rotation)", + logger.Info( + "stage 4: refreshing session key from chain (possible rotation)", "stage", "session_key", "path", "refresh", ) @@ -955,7 +1164,8 @@ func (h *JobHandler) deriveAndStoreSessionKey(ctx context.Context, logger *slog. if err := h.keyStore.StoreKey(sessionID, sessionKey); err != nil { // Log but don't fail — key is in memory for this job - logger.Warn("failed to persist session key", + logger.Warn( + "failed to persist session key", "stage", "session_key", "error", err, ) @@ -990,6 +1200,28 @@ func init() { } } +// relayCompleteCiphertext returns the ciphertext to deliver on the relay +// `complete` frame. For a v2 search envelope it re-encrypts just the plain +// answer (the v1.1 consumer contract — the envelope with searchContext stays in +// the blob for the disputer). For a legacy/plain ciphertext it returns it +// unchanged. Best-effort: on any decrypt/decode/encrypt error it falls back to +// the original ciphertext (never fails the job). +func (h *JobHandler) relayCompleteCiphertext(sessionKey, blobCiphertext []byte) []byte { + plain, err := pkgcrypto.Decrypt(sessionKey, blobCiphertext) + if err != nil { + return blobCiphertext + } + env := searchaug.DecodeResponse(plain) + if env.V != searchaug.ResponseEnvelopeVersion { + return blobCiphertext // legacy/non-search: already the plain answer + } + ac, err := pkgcrypto.Encrypt(sessionKey, []byte(env.Answer)) + if err != nil { + return blobCiphertext + } + return ac +} + // publishToRedis signs and publishes the response via the configured ResponsePublisher. // Errors are logged but non-fatal — the on-chain blob is the authoritative response. func (h *JobHandler) publishToRedis( @@ -1005,7 +1237,8 @@ func (h *JobHandler) publishToRedis( sig, err := signMismatchEvidence(h.cfg.ChainID, h.cfg.JobRegistryAddr, jobID, sessionID, ciphertext, h.signingKey) if err != nil { - h.logger.Warn("failed to sign response", + h.logger.Warn( + "failed to sign response", "stage", "redis_publish", "jobID", jobID, "error", err, @@ -1104,7 +1337,7 @@ func (h *JobHandler) buildConversationHistory( if err != nil { return nil, fmt.Errorf("decrypt response for job %d: %w", jobID, err) } - messages = append(messages, ollama.ChatMessage{Role: "assistant", Content: string(responseText)}) + messages = append(messages, ollama.ChatMessage{Role: "assistant", Content: searchaug.DecodeResponse(responseText).Answer}) } } @@ -1127,6 +1360,43 @@ func (h *JobHandler) resolveModelName(modelID string) (string, error) { return modelID, nil } +// toSearchaugSources converts search.Source (the Tavily/search-package type) +// to searchaug.Source (the shared pkg type). Fields are identical; this +// converter avoids a dependency between the two packages. +func toSearchaugSources(sources []search.Source) []searchaug.Source { + out := make([]searchaug.Source, len(sources)) + for i, s := range sources { + out[i] = searchaug.Source{ + Position: s.Position, + Title: s.Title, + URL: s.URL, + Snippet: s.Snippet, + } + } + return out +} + +// sourcesMetadataJSON renders the citation payload in the exact shape the +// frontend's parseWebSearchSources expects. Accepts []searchaug.Source so it +// can be called after the search.Source → searchaug.Source conversion. +func sourcesMetadataJSON(sources []searchaug.Source) []byte { + type wire struct { + Position int `json:"position"` + Title string `json:"title"` + URL string `json:"url"` + Snippet string `json:"snippet"` + } + out := struct { + Type string `json:"type"` + Sources []wire `json:"sources"` + }{Type: "webSearchSources"} + for _, s := range sources { + out.Sources = append(out.Sources, wire{s.Position, s.Title, s.URL, s.Snippet}) + } + data, _ := json.Marshal(out) + return data +} + func normalizeModelLookupKey(modelID string) string { return strings.ToLower(strings.TrimPrefix(strings.TrimPrefix(modelID, "0x"), "0X")) } diff --git a/internal/pipeline/handler_search_test.go b/internal/pipeline/handler_search_test.go new file mode 100644 index 0000000..9138996 --- /dev/null +++ b/internal/pipeline/handler_search_test.go @@ -0,0 +1,638 @@ +package pipeline + +import ( + "context" + "encoding/json" + "fmt" + "io" + "log/slog" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/ethereum/go-ethereum/common" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + pkgcrypto "github.com/lightchain/pkg/crypto" + "github.com/lightchain/pkg/searchaug" + "github.com/lightchain/worker/internal/metrics" + "github.com/lightchain/worker/internal/ollama" + "github.com/lightchain/worker/internal/search" +) + +// recordingPublisher implements ResponsePublisher and records the wall-clock +// time at which PublishMetadata is called. This lets ordering tests assert +// that metadata is published after inference completes. It also records chunk +// sequences for TestStage5_StreamsChunks. +type recordingPublisher struct { + mu sync.Mutex + metadataCalled bool + metadataAt time.Time + responseCalled bool + responseAt time.Time + chunkSeqs []uint32 +} + +func (r *recordingPublisher) PublishMetadata(_ context.Context, _, _ uint64, _ string, _ []byte) { + r.mu.Lock() + defer r.mu.Unlock() + r.metadataCalled = true + r.metadataAt = time.Now() +} + +func (r *recordingPublisher) PublishResponse(_ context.Context, _, _ uint64, _, _ string, _ []byte) { + r.mu.Lock() + defer r.mu.Unlock() + r.responseCalled = true + r.responseAt = time.Now() +} + +func (r *recordingPublisher) PublishChunk(_ context.Context, _, _ uint64, _ string, seq uint32, _ []byte) { + r.mu.Lock() + defer r.mu.Unlock() + r.chunkSeqs = append(r.chunkSeqs, seq) +} + +// trackingOllama wraps mockOllama and records when Generate/Chat complete. +type trackingOllama struct { + inner *mockOllama + mu sync.Mutex + completedAt time.Time +} + +func (t *trackingOllama) Generate(ctx context.Context, model, prompt string) (string, error) { + resp, err := t.inner.Generate(ctx, model, prompt) + t.mu.Lock() + t.completedAt = time.Now() + t.mu.Unlock() + return resp, err +} + +func (t *trackingOllama) Chat(ctx context.Context, model string, messages []ollama.ChatMessage) (string, error) { + resp, err := t.inner.Chat(ctx, model, messages) + t.mu.Lock() + t.completedAt = time.Now() + t.mu.Unlock() + return resp, err +} + +func (t *trackingOllama) GenerateStream(ctx context.Context, model, prompt string, onDelta func(string)) (string, error) { + resp, err := t.inner.GenerateStream(ctx, model, prompt, onDelta) + t.mu.Lock() + t.completedAt = time.Now() + t.mu.Unlock() + return resp, err +} + +func (t *trackingOllama) ChatStream(ctx context.Context, model string, messages []ollama.ChatMessage, onDelta func(string)) (string, error) { + resp, err := t.inner.ChatStream(ctx, model, messages, onDelta) + t.mu.Lock() + t.completedAt = time.Now() + t.mu.Unlock() + return resp, err +} + +// fakeSearcher satisfies search.Searcher and returns a fixed set of sources. +type fakeSearcher struct { + sources []search.Source + err error +} + +func (f *fakeSearcher) Search(_ context.Context, _ string, _ int) ([]search.Source, error) { + return f.sources, f.err +} + +// fakeStreamInference implements InferenceClient with configurable per-delta output +// so TestStage5_StreamsChunks can control exactly which deltas are emitted. +type fakeStreamInference struct { + deltas []string +} + +func (f *fakeStreamInference) Generate(_ context.Context, _, _ string) (string, error) { + return strings.Join(f.deltas, ""), nil +} + +func (f *fakeStreamInference) Chat(_ context.Context, _ string, _ []ollama.ChatMessage) (string, error) { + return strings.Join(f.deltas, ""), nil +} + +func (f *fakeStreamInference) GenerateStream(_ context.Context, _, _ string, onDelta func(string)) (string, error) { + var full strings.Builder + for _, d := range f.deltas { + onDelta(d) + full.WriteString(d) + } + return full.String(), nil +} + +func (f *fakeStreamInference) ChatStream(_ context.Context, _ string, _ []ollama.ChatMessage, onDelta func(string)) (string, error) { + return f.GenerateStream(context.Background(), "", "", onDelta) +} + +func twoSources() []search.Source { + return []search.Source{ + {Position: 1, Title: "Result One", URL: "https://example.com/1", Snippet: "first snippet"}, + {Position: 2, Title: "Result Two", URL: "https://example.com/2", Snippet: "second snippet"}, + } +} + +// TestStage45_MetadataEmittedAfterInference asserts that the web-search +// citation frame (PublishMetadata) is sent AFTER inference completes, not +// before. This is the ordering the UI requires so Sources appear beneath +// the answer. +func TestStage45_MetadataEmittedAfterInference(t *testing.T) { + t.Parallel() + + sessionKey := testSessionKey(t) + ecdhKey := testECDHKey(t) + encSessionKey := encryptSessionKeyForWorker(t, sessionKey, ecdhKey) + // Search flag now comes from the blob envelope; the fetcher must return the + // same envelope the direct blobData carries so both paths exercise search. + promptCiphertext, err := pkgcrypto.Encrypt(sessionKey, searchaug.EncodePrompt("What is the capital of France?", true)) + require.NoError(t, err) + + chain := &mockChainClient{ + ackJobFn: func(_ context.Context, _ uint64) error { return nil }, + completeJobFn: func(_ context.Context, _ uint64, _, _ [32]byte) error { 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, _ []byte) ([][32]byte, error) { + return [][32]byte{{0x01}}, nil + }} + + inner := &mockOllama{generateFn: func(_ context.Context, _, _ string) (string, error) { + return "Paris", nil + }} + tracker := &trackingOllama{inner: inner} + rec := &recordingPublisher{} + + counter := &atomic.Int32{} + logger := slog.New(slog.NewTextHandler(io.Discard, &slog.HandlerOptions{Level: slog.LevelError})) + + handler := NewJobHandler( + chain, fetcher, submitter, newMockKeyStore(), tracker, + nil, // no real Redis — publisher injected directly + testSigningKey(t), ecdhKey, counter, logger, + HandlerConfig{ + AckTxTimeout: 5 * time.Second, + BlobTxTimeout: 60 * time.Second, + SearchMaxResults: 3, + }, + rec, // injected recording publisher + nil, // no checkpoints + testMetrics(t), + metrics.DeliveryAsynq, + &fakeSearcher{sources: twoSources()}, + ) + + payload := testPayload(t) + + // Encrypt the session key into the keystore so stage 3 succeeds without chain + ks := newMockKeyStore() + ks.keys[payload.SessionID] = sessionKey + handler.keyStore = ks + + // Search flag now comes from the blob envelope, not the payload field. + blobData, err := pkgcrypto.Encrypt(sessionKey, searchaug.EncodePrompt("What is the capital of France?", true)) + require.NoError(t, err) + + _, err = handler.runInferencePipeline( + context.Background(), + logger, + payload, + blobData, + "llama3-8b", + metrics.DeliveryAsynq, + ) + require.NoError(t, err) + + require.True(t, rec.metadataCalled, "sources must be published") + require.False(t, tracker.completedAt.IsZero(), "inference must have been called") + + assert.True( + t, + !rec.metadataAt.Before(tracker.completedAt), + "PublishMetadata (at %v) must not happen before inference completed (at %v)", + rec.metadataAt, tracker.completedAt, + ) +} + +func TestSourcesMetadataPayload_ShapeMatchesFrontend(t *testing.T) { + sources := []searchaug.Source{{Position: 1, Title: "T1", URL: "https://a", Snippet: "s1"}} + payload := sourcesMetadataJSON(sources) + var decoded struct { + Type string `json:"type"` + Sources []struct { + Position int `json:"position"` + Title string `json:"title"` + URL string `json:"url"` + Snippet string `json:"snippet"` + } `json:"sources"` + } + require.NoError(t, json.Unmarshal(payload, &decoded)) + assert.Equal(t, "webSearchSources", decoded.Type) + require.Len(t, decoded.Sources, 1) + assert.Equal(t, "https://a", decoded.Sources[0].URL) +} + +// TestStage5_StreamsChunks asserts that stage 5 publishes chunk frames in monotonically +// increasing sequence order and that the returned ciphertext decrypts to the full text. +func TestStage5_StreamsChunks(t *testing.T) { + t.Parallel() + + sessionKey := testSessionKey(t) + ecdhKey := testECDHKey(t) + + // deltas chosen so exactly two flushes happen: + // delta 1: "Hello, streaming world!!" = 24 bytes → buffer reaches threshold (>=24) → mid-stream flush (seq=1) + // delta 2: "Done" = 4 bytes → buffer = 4, below threshold, not flushed mid-stream + // after loop: flushChunk() fires on the remaining 4-byte remainder → final flush (seq=2) + // This exercises both the threshold flush path AND the final-remainder flush path. + deltas := []string{"Hello, streaming world!!", "Done"} + gen := &fakeStreamInference{deltas: deltas} + rec := &recordingPublisher{} + + chain := &mockChainClient{ + ackJobFn: func(_ context.Context, _ uint64) error { return nil }, + completeJobFn: func(_ context.Context, _ uint64, _, _ [32]byte) error { return nil }, + getEncWorkerKeyFn: func(_ context.Context, _ uint64) ([]byte, error) { + return encryptSessionKeyForWorker(t, sessionKey, ecdhKey), nil + }, + } + fetcher := &mockBlobFetcher{fetchFn: func(_ context.Context, _ common.Hash, _ uint64) ([]byte, error) { + return encryptBlob(t, sessionKey, "test prompt"), nil + }} + submitter := &mockBlobSubmitter{submitFn: func(_ context.Context, _ []byte) ([][32]byte, error) { + return [][32]byte{{0x01}}, nil + }} + + counter := &atomic.Int32{} + logger := slog.New(slog.NewTextHandler(io.Discard, &slog.HandlerOptions{Level: slog.LevelError})) + + handler := NewJobHandler( + chain, fetcher, submitter, newMockKeyStore(), gen, + nil, + testSigningKey(t), ecdhKey, counter, logger, + HandlerConfig{ + AckTxTimeout: 5 * time.Second, + BlobTxTimeout: 60 * time.Second, + }, + rec, + nil, + testMetrics(t), + metrics.DeliveryAsynq, + nil, + ) + + payload := testPayload(t) + ks := newMockKeyStore() + ks.keys[payload.SessionID] = sessionKey + handler.keyStore = ks + + blobData, err := pkgcrypto.Encrypt(sessionKey, []byte("test prompt")) + require.NoError(t, err) + + ciphertext, err := handler.runInferencePipeline( + context.Background(), + logger, + payload, + blobData, + "llama3-8b", + metrics.DeliveryAsynq, + ) + require.NoError(t, err) + + // Must produce exactly two chunk frames: seq=1 from the mid-stream threshold flush + // (buffer reached 24 bytes after delta 1) and seq=2 from the final-remainder flush + // (the leftover 4 bytes of delta 2 flushed after GenerateStream returns). + rec.mu.Lock() + seqs := rec.chunkSeqs + rec.mu.Unlock() + assert.Equal(t, []uint32{1, 2}, seqs, "expected exactly two monotonic chunk seqs: threshold flush then final-remainder flush") + + // Full response ciphertext must decrypt to the concatenated deltas. + // Note: non-search job → EncodeResponse returns raw answer bytes → ciphertext + // decrypts back to the plain text (legacy format unchanged). + plaintext, decErr := pkgcrypto.Decrypt(sessionKey, ciphertext) + require.NoError(t, decErr) + assert.Equal(t, strings.Join(deltas, ""), string(plaintext)) +} + +// TestStage6_EmitsEnvelopeForSearchJob asserts that when web search returns +// sources, stage 6 encrypts a v2 JSON envelope (not raw answer bytes) so the +// disputer can replay the search context deterministically. +func TestStage6_EmitsEnvelopeForSearchJob(t *testing.T) { + t.Parallel() + + sessionKey := testSessionKey(t) + ecdhKey := testECDHKey(t) + + gen := &mockOllama{generateFn: func(_ context.Context, _, _ string) (string, error) { + return "the answer", nil + }} + rec := &recordingPublisher{} + + chain := &mockChainClient{ + ackJobFn: func(_ context.Context, _ uint64) error { return nil }, + completeJobFn: func(_ context.Context, _ uint64, _, _ [32]byte) error { return nil }, + getEncWorkerKeyFn: func(_ context.Context, _ uint64) ([]byte, error) { + return encryptSessionKeyForWorker(t, sessionKey, ecdhKey), nil + }, + } + // Fetcher returns the same search envelope the direct blobData carries so + // both paths exercise the search-enabled branch. + fetcher := &mockBlobFetcher{fetchFn: func(_ context.Context, _ common.Hash, _ uint64) ([]byte, error) { + return pkgcrypto.Encrypt(sessionKey, searchaug.EncodePrompt("what year is it?", true)) + }} + submitter := &mockBlobSubmitter{submitFn: func(_ context.Context, _ []byte) ([][32]byte, error) { + return [][32]byte{{0x01}}, nil + }} + + counter := &atomic.Int32{} + logger := slog.New(slog.NewTextHandler(io.Discard, &slog.HandlerOptions{Level: slog.LevelError})) + + handler := NewJobHandler( + chain, fetcher, submitter, newMockKeyStore(), gen, + nil, + testSigningKey(t), ecdhKey, counter, logger, + HandlerConfig{ + AckTxTimeout: 5 * time.Second, + BlobTxTimeout: 60 * time.Second, + SearchMaxResults: 3, + }, + rec, + nil, + testMetrics(t), + metrics.DeliveryAsynq, + &fakeSearcher{sources: twoSources()}, + ) + + payload := testPayload(t) + ks := newMockKeyStore() + ks.keys[payload.SessionID] = sessionKey + handler.keyStore = ks + + // Search flag now comes from the blob envelope, not the payload field. + blobData, err := pkgcrypto.Encrypt(sessionKey, searchaug.EncodePrompt("what year is it?", true)) + require.NoError(t, err) + + ct, err := handler.runInferencePipeline( + context.Background(), + logger, + payload, + blobData, + "llama3-8b", + metrics.DeliveryAsynq, + ) + require.NoError(t, err) + + // Decrypt the ciphertext and decode the v2 envelope. + plaintext, decErr := pkgcrypto.Decrypt(sessionKey, ct) + require.NoError(t, decErr) + env := searchaug.DecodeResponse(plaintext) + + assert.Equal(t, searchaug.ResponseEnvelopeVersion, env.V) + assert.Equal(t, "the answer", env.Answer) + assert.Len(t, env.SearchContext, 2) +} + +// TestStage6_RawAnswerForNonSearchJob asserts that when no web search is +// performed, stage 6 encrypts the raw answer bytes (legacy format) — the +// v2 envelope is NOT emitted. +func TestStage6_RawAnswerForNonSearchJob(t *testing.T) { + t.Parallel() + + sessionKey := testSessionKey(t) + ecdhKey := testECDHKey(t) + + gen := &mockOllama{generateFn: func(_ context.Context, _, _ string) (string, error) { + return "plain", nil + }} + rec := &recordingPublisher{} + + chain := &mockChainClient{ + ackJobFn: func(_ context.Context, _ uint64) error { return nil }, + completeJobFn: func(_ context.Context, _ uint64, _, _ [32]byte) error { return nil }, + getEncWorkerKeyFn: func(_ context.Context, _ uint64) ([]byte, error) { + return encryptSessionKeyForWorker(t, sessionKey, ecdhKey), nil + }, + } + fetcher := &mockBlobFetcher{fetchFn: func(_ context.Context, _ common.Hash, _ uint64) ([]byte, error) { + return encryptBlob(t, sessionKey, "hello"), nil + }} + submitter := &mockBlobSubmitter{submitFn: func(_ context.Context, _ []byte) ([][32]byte, error) { + return [][32]byte{{0x01}}, nil + }} + + counter := &atomic.Int32{} + logger := slog.New(slog.NewTextHandler(io.Discard, &slog.HandlerOptions{Level: slog.LevelError})) + + handler := NewJobHandler( + chain, fetcher, submitter, newMockKeyStore(), gen, + nil, + testSigningKey(t), ecdhKey, counter, logger, + HandlerConfig{ + AckTxTimeout: 5 * time.Second, + BlobTxTimeout: 60 * time.Second, + }, + rec, + nil, + testMetrics(t), + metrics.DeliveryAsynq, + nil, // no searcher → SearchEnabled is ignored, searchSources stays nil + ) + + payload := testPayload(t) + // SearchEnabled = false (default), no searcher wired + ks := newMockKeyStore() + ks.keys[payload.SessionID] = sessionKey + handler.keyStore = ks + + blobData, err := pkgcrypto.Encrypt(sessionKey, []byte("hello")) + require.NoError(t, err) + + ct, err := handler.runInferencePipeline( + context.Background(), + logger, + payload, + blobData, + "llama3-8b", + metrics.DeliveryAsynq, + ) + require.NoError(t, err) + + // Decrypt and decode; for non-search jobs EncodeResponse returns raw bytes + // so DecodeResponse falls back to treating the bytes as a raw answer. + plaintext, decErr := pkgcrypto.Decrypt(sessionKey, ct) + require.NoError(t, decErr) + env := searchaug.DecodeResponse(plaintext) + + assert.Equal(t, "plain", env.Answer) + assert.Nil(t, env.SearchContext, "non-search job must not emit a SearchContext") +} + +// TestBuildConversationHistory_DecodesV2EnvelopeToAnswer asserts that when a +// prior job's response blob is a v2 search envelope, buildConversationHistory +// assembles an assistant message whose Content is the decoded answer — NOT the +// raw envelope JSON. For legacy/non-search blobs DecodeResponse falls back to +// treating the bytes as a raw answer, so behavior for those is unchanged. +func TestBuildConversationHistory_DecodesV2EnvelopeToAnswer(t *testing.T) { + t.Parallel() + + sessionKey := testSessionKey(t) + ecdhKey := testECDHKey(t) + + const priorAnswer = "prior answer from search" + priorSources := []searchaug.Source{ + {Position: 1, Title: "T1", URL: "https://example.com", Snippet: "snip"}, + } + + // Build a v2 envelope and encrypt it as the worker would store it. + envelope, err := searchaug.EncodeResponse(priorAnswer, priorSources) + require.NoError(t, err) + encEnvelope, err := pkgcrypto.Encrypt(sessionKey, envelope) + require.NoError(t, err) + + // Build a plain prompt blob for the same prior job. + encPrompt := encryptBlob(t, sessionKey, "prior user question") + + // Sentinel hashes to distinguish prompt from response fetches. + promptHash := common.HexToHash("0x0101010101010101010101010101010101010101010101010101010101010101") + responseHash := common.HexToHash("0x0202020202020202020202020202020202020202020202020202020202020202") + + chain := &mockChainClient{ + ackJobFn: func(_ context.Context, _ uint64) error { return nil }, + completeJobFn: func(_ context.Context, _ uint64, _, _ [32]byte) error { return nil }, + getEncWorkerKeyFn: func(_ context.Context, _ uint64) ([]byte, error) { + return encryptSessionKeyForWorker(t, sessionKey, ecdhKey), nil + }, + getJobBlobInfoFn: func(_ context.Context, _ uint64) (common.Hash, common.Hash, uint64, uint64, error) { + return promptHash, responseHash, 10, 11, nil + }, + } + + fetcher := &mockBlobFetcher{fetchFn: func(_ context.Context, hash common.Hash, _ uint64) ([]byte, error) { + switch hash { + case promptHash: + return encPrompt, nil + case responseHash: + return encEnvelope, nil + default: + return nil, fmt.Errorf("unexpected hash %s", hash) + } + }} + submitter := &mockBlobSubmitter{submitFn: func(_ context.Context, _ []byte) ([][32]byte, error) { + return [][32]byte{{0x03}}, nil + }} + + counter := &atomic.Int32{} + logger := slog.New(slog.NewTextHandler(io.Discard, &slog.HandlerOptions{Level: slog.LevelError})) + + handler := NewJobHandler( + chain, fetcher, submitter, newMockKeyStore(), &mockOllama{generateFn: func(_ context.Context, _, _ string) (string, error) { + return "", nil + }}, + nil, + testSigningKey(t), ecdhKey, counter, logger, + HandlerConfig{ + AckTxTimeout: 5 * time.Second, + BlobTxTimeout: 60 * time.Second, + }, + &recordingPublisher{}, + nil, + testMetrics(t), + metrics.DeliveryAsynq, + nil, + ) + + msgs, err := handler.buildConversationHistory(context.Background(), []uint64{99}, sessionKey) + require.NoError(t, err) + require.Len(t, msgs, 2, "expected one user + one assistant message") + + userMsg := msgs[0] + assert.Equal(t, "user", userMsg.Role) + assert.Equal(t, "prior user question", userMsg.Content) + + assistantMsg := msgs[1] + assert.Equal(t, "assistant", assistantMsg.Role) + // Must be the decoded answer, not the raw envelope JSON. + assert.Equal(t, priorAnswer, assistantMsg.Content, + "assistant Content must be the decoded answer, not the raw v2 envelope JSON") +} + +// TestDecodePromptDrivesSearch verifies the searchaug.DecodePrompt contract +// that the handler now depends on: a search envelope enables search and unwraps +// the text, while raw bytes return (text, false, nil). +func TestDecodePromptDrivesSearch(t *testing.T) { + // search envelope → searcher invoked, prompt unwrapped + p, s, err := searchaug.DecodePrompt(searchaug.EncodePrompt("find news", true)) + if err != nil || p != "find news" || !s { + t.Fatalf("got (%q,%v,%v)", p, s, err) + } + // raw text → no search + p2, s2, _ := searchaug.DecodePrompt([]byte("plain")) + if p2 != "plain" || s2 { + t.Fatalf("raw text must not enable search, got (%q,%v)", p2, s2) + } + // 0x00-sentinel but invalid JSON → hard error (the handler returns a wrapped + // stage-4 failure rather than feeding NUL+garbage to the model). + if _, _, decErr := searchaug.DecodePrompt([]byte{0x00, '{', 'x'}); decErr == nil { + t.Fatalf("malformed envelope must return a non-nil error") + } +} + +// TestRelayCompleteCiphertext verifies that relayCompleteCiphertext strips the +// v2 search envelope before the relay complete frame so the frontend receives +// the plain answer, not the JSON blob. +func TestRelayCompleteCiphertext(t *testing.T) { + t.Parallel() + + // Minimal JobHandler — relayCompleteCiphertext only uses pkgcrypto + searchaug, + // no other dependencies. + h := &JobHandler{} + + t.Run("v2 envelope: complete frame decrypts to plain answer", func(t *testing.T) { + t.Parallel() + + sk := testSessionKey(t) + + // Build a v2 blob ciphertext as stage 6 would produce it. + envBytes, err := searchaug.EncodeResponse("the answer", toSearchaugSources(twoSources())) + require.NoError(t, err) + blobCt, err := pkgcrypto.Encrypt(sk, envBytes) + require.NoError(t, err) + + out := h.relayCompleteCiphertext(sk, blobCt) + + // The relay frame must decrypt to the plain answer — NOT the envelope JSON. + plain, decErr := pkgcrypto.Decrypt(sk, out) + require.NoError(t, decErr) + assert.Equal(t, "the answer", string(plain), + "relay complete frame must carry the plain answer, not the v2 envelope JSON") + }) + + t.Run("legacy/plain: complete frame passes through unchanged (decrypts to plain answer)", func(t *testing.T) { + t.Parallel() + + sk := testSessionKey(t) + + // Non-search job: EncodeResponse with nil sources returns raw answer bytes. + rawBytes, err := searchaug.EncodeResponse("plain answer", nil) + require.NoError(t, err) + blobCt, err := pkgcrypto.Encrypt(sk, rawBytes) + require.NoError(t, err) + + out := h.relayCompleteCiphertext(sk, blobCt) + + plain, decErr := pkgcrypto.Decrypt(sk, out) + require.NoError(t, decErr) + assert.Equal(t, "plain answer", string(plain), + "relay complete frame for non-search job must decrypt to the plain answer") + }) +} diff --git a/internal/pipeline/handler_test.go b/internal/pipeline/handler_test.go index 46e89b5..de5edb4 100644 --- a/internal/pipeline/handler_test.go +++ b/internal/pipeline/handler_test.go @@ -53,6 +53,7 @@ type mockChainClient struct { hasJobAcknowledgedFn func(ctx context.Context, jobID uint64) (bool, error) hasJobCompletedFn func(ctx context.Context, jobID uint64) (bool, error) getEncWorkerKeyFn func(ctx context.Context, sessionID uint64) ([]byte, error) + getJobBlobInfoFn func(ctx context.Context, jobID uint64) (common.Hash, common.Hash, uint64, uint64, error) } func (m *mockChainClient) AcknowledgeJob(ctx context.Context, jobID uint64) error { @@ -76,7 +77,10 @@ func (m *mockChainClient) HasJobCompleted(ctx context.Context, jobID uint64) (bo func (m *mockChainClient) GetSessionEncWorkerKey(ctx context.Context, sessionID uint64) ([]byte, error) { return m.getEncWorkerKeyFn(ctx, sessionID) } -func (m *mockChainClient) GetJobBlobInfo(_ context.Context, _ uint64) (common.Hash, common.Hash, uint64, uint64, error) { +func (m *mockChainClient) GetJobBlobInfo(ctx context.Context, jobID uint64) (common.Hash, common.Hash, uint64, uint64, error) { + if m.getJobBlobInfoFn != nil { + return m.getJobBlobInfoFn(ctx, jobID) + } return common.Hash{}, common.Hash{}, 0, 0, nil } @@ -138,6 +142,22 @@ func (m *mockOllama) Chat(ctx context.Context, model string, messages []ollama.C return m.generateFn(ctx, model, messages[len(messages)-1].Content) } +func (m *mockOllama) GenerateStream(ctx context.Context, model, prompt string, onDelta func(string)) (string, error) { + resp, err := m.Generate(ctx, model, prompt) + if err == nil && resp != "" { + onDelta(resp) + } + return resp, err +} + +func (m *mockOllama) ChatStream(ctx context.Context, model string, messages []ollama.ChatMessage, onDelta func(string)) (string, error) { + resp, err := m.Chat(ctx, model, messages) + if err == nil && resp != "" { + onDelta(resp) + } + return resp, err +} + // --- Helpers --- func testSigningKey(t *testing.T) *ecdsa.PrivateKey { @@ -228,6 +248,7 @@ func newTestHandlerWithConfig( nil, // checkpoints — disabled in most handler tests; see TestHandleTask_Checkpoint* for coverage testMetrics(t), metrics.DeliveryAsynq, + nil, ) } @@ -307,6 +328,7 @@ func TestHandleTask_FullPipelineSuccess(t *testing.T) { nil, // checkpoints disabled for this test testMetrics(t), metrics.DeliveryAsynq, + nil, ) payload := testPayload(t) @@ -419,7 +441,7 @@ func TestHandleTask_SessionKeyCacheHit(t *testing.T) { logger := slog.New(slog.NewTextHandler(io.Discard, &slog.HandlerOptions{Level: slog.LevelError})) handler := NewJobHandler(chain, fetcher, submitter, ks, ollama, rc, testSigningKey(t), ecdhKey, counter, logger, - HandlerConfig{AckTxTimeout: 5 * time.Second, BlobTxTimeout: 60 * time.Second, RedisPublishTimeout: 5 * time.Second}, nil, nil, testMetrics(t), metrics.DeliveryAsynq) + HandlerConfig{AckTxTimeout: 5 * time.Second, BlobTxTimeout: 60 * time.Second, RedisPublishTimeout: 5 * time.Second}, nil, nil, testMetrics(t), metrics.DeliveryAsynq, nil) payload := testPayload(t) data, _ := json.Marshal(payload) @@ -465,7 +487,7 @@ func TestHandleTask_SessionKeyCacheMiss_FetchAndStore(t *testing.T) { logger := slog.New(slog.NewTextHandler(io.Discard, &slog.HandlerOptions{Level: slog.LevelError})) handler := NewJobHandler(chain, fetcher, submitter, ks, ollama, rc, testSigningKey(t), ecdhKey, counter, logger, - HandlerConfig{AckTxTimeout: 5 * time.Second, BlobTxTimeout: 60 * time.Second, RedisPublishTimeout: 5 * time.Second}, nil, nil, testMetrics(t), metrics.DeliveryAsynq) + HandlerConfig{AckTxTimeout: 5 * time.Second, BlobTxTimeout: 60 * time.Second, RedisPublishTimeout: 5 * time.Second}, nil, nil, testMetrics(t), metrics.DeliveryAsynq, nil) payload := testPayload(t) data, _ := json.Marshal(payload) @@ -533,7 +555,7 @@ func TestHandleTask_SessionKeyRotated_Refreshes(t *testing.T) { logger := slog.New(slog.NewTextHandler(io.Discard, &slog.HandlerOptions{Level: slog.LevelError})) handler := NewJobHandler(chain, fetcher, submitter, ks, ollama, rc, testSigningKey(t), ecdhKey, counter, logger, - HandlerConfig{AckTxTimeout: 5 * time.Second, BlobTxTimeout: 60 * time.Second, RedisPublishTimeout: 5 * time.Second}, nil, nil, testMetrics(t), metrics.DeliveryAsynq) + HandlerConfig{AckTxTimeout: 5 * time.Second, BlobTxTimeout: 60 * time.Second, RedisPublishTimeout: 5 * time.Second}, nil, nil, testMetrics(t), metrics.DeliveryAsynq, nil) payload := testPayload(t) data, _ := json.Marshal(payload) @@ -597,7 +619,7 @@ func TestGetOrDeriveSessionKey_CoalescesConcurrentMisses(t *testing.T) { chain, &mockBlobFetcher{}, &mockBlobSubmitter{}, ks, &mockOllama{}, rc, testSigningKey(t), ecdhKey, counter, logger, HandlerConfig{AckTxTimeout: 5 * time.Second, BlobTxTimeout: 60 * time.Second, RedisPublishTimeout: 5 * time.Second}, - nil, nil, testMetrics(t), metrics.DeliveryAsynq, + nil, nil, testMetrics(t), metrics.DeliveryAsynq, nil, ) var wg sync.WaitGroup @@ -657,7 +679,7 @@ func TestHandleTask_CompleteJobFailure(t *testing.T) { logger := slog.New(slog.NewTextHandler(io.Discard, &slog.HandlerOptions{Level: slog.LevelError})) handler := NewJobHandler(chain, fetcher, submitter, ks, ollama, rc, testSigningKey(t), ecdhKey, counter, logger, - HandlerConfig{AckTxTimeout: 5 * time.Second, BlobTxTimeout: 60 * time.Second, RedisPublishTimeout: 5 * time.Second}, nil, nil, testMetrics(t), metrics.DeliveryAsynq) + HandlerConfig{AckTxTimeout: 5 * time.Second, BlobTxTimeout: 60 * time.Second, RedisPublishTimeout: 5 * time.Second}, nil, nil, testMetrics(t), metrics.DeliveryAsynq, nil) payload := testPayload(t) data, _ := json.Marshal(payload) @@ -692,7 +714,7 @@ func TestHandleTask_JobCounterIncrementDecrement(t *testing.T) { logger := slog.New(slog.NewTextHandler(io.Discard, &slog.HandlerOptions{Level: slog.LevelError})) handler := NewJobHandler(chain, fetcher, submitter, newMockKeyStore(), ollama, rc, testSigningKey(t), testECDHKey(t), counter, logger, - HandlerConfig{AckTxTimeout: 5 * time.Second, BlobTxTimeout: 60 * time.Second, RedisPublishTimeout: 5 * time.Second}, nil, nil, testMetrics(t), metrics.DeliveryAsynq) + HandlerConfig{AckTxTimeout: 5 * time.Second, BlobTxTimeout: 60 * time.Second, RedisPublishTimeout: 5 * time.Second}, nil, nil, testMetrics(t), metrics.DeliveryAsynq, nil) payload := testPayload(t) data, _ := json.Marshal(payload) @@ -738,7 +760,7 @@ func TestHandleTask_RedisPublishFailure_NonFatal(t *testing.T) { logger := slog.New(slog.NewTextHandler(io.Discard, &slog.HandlerOptions{Level: slog.LevelError})) handler := NewJobHandler(chain, fetcher, submitter, ks, ollama, rc, testSigningKey(t), ecdhKey, counter, logger, - HandlerConfig{AckTxTimeout: 5 * time.Second, BlobTxTimeout: 60 * time.Second, RedisPublishTimeout: 5 * time.Second}, nil, nil, testMetrics(t), metrics.DeliveryAsynq) + HandlerConfig{AckTxTimeout: 5 * time.Second, BlobTxTimeout: 60 * time.Second, RedisPublishTimeout: 5 * time.Second}, nil, nil, testMetrics(t), metrics.DeliveryAsynq, nil) payload := testPayload(t) data, _ := json.Marshal(payload) @@ -1002,6 +1024,7 @@ func TestHandleTask_CheckpointHit_SkipsInference(t *testing.T) { store, testMetrics(t), metrics.DeliveryAsynq, + nil, ) payload := testPayload(t) @@ -1076,6 +1099,7 @@ func TestHandleTask_CheckpointRaceLoserUsesCanonical(t *testing.T) { store, testMetrics(t), metrics.DeliveryAsynq, + nil, ) payload := testPayload(t) @@ -1160,6 +1184,7 @@ func TestHandleTask_CheckpointTombstonedAfterCompletion(t *testing.T) { store, testMetrics(t), metrics.DeliveryAsynq, + nil, ) payload := testPayload(t) @@ -1222,6 +1247,7 @@ func TestHandleTask_CheckpointRetained_WhenStage8bFails(t *testing.T) { store, testMetrics(t), metrics.DeliveryAsynq, + nil, ) payload := testPayload(t) diff --git a/internal/pipeline/release_tracker_test.go b/internal/pipeline/release_tracker_test.go index 3cbf9af..1d8ecfa 100644 --- a/internal/pipeline/release_tracker_test.go +++ b/internal/pipeline/release_tracker_test.go @@ -102,7 +102,7 @@ func runFullPipelineWithTracker(t *testing.T, tracker ReleaseTracker) (handlerEr RedisPublishTimeout: 5 * time.Second, ModelIDToName: map[string]string{expectedModelID: "llama3-8b"}, }, - nil, nil, testMetrics(t), metrics.DeliveryAsynq, + nil, nil, testMetrics(t), metrics.DeliveryAsynq, nil, ) handler.SetReleaseTracker(tracker) diff --git a/internal/search/client.go b/internal/search/client.go new file mode 100644 index 0000000..fe8a6bf --- /dev/null +++ b/internal/search/client.go @@ -0,0 +1,100 @@ +// Package search performs web-search augmentation for inference jobs. +// v1 uses Tavily with a single, non-streaming query (the raw prompt). +package search + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "net/http" + "time" +) + +// Source is one search result, in display/citation order (Position is 1-based). +type Source struct { + Position int `json:"position"` + Title string `json:"title"` + URL string `json:"url"` + Snippet string `json:"snippet"` +} + +// Searcher turns a query into ranked web results. Implemented by TavilyClient; +// the interface is the seam for tests and future providers. +type Searcher interface { + Search(ctx context.Context, query string, maxResults int) ([]Source, error) +} + +// TavilyClient calls the Tavily Search API (POST {baseURL}/search). +type TavilyClient struct { + baseURL string + apiKey string + httpClient *http.Client +} + +func NewTavilyClient(baseURL, apiKey string, timeout time.Duration) *TavilyClient { + return &TavilyClient{ + baseURL: baseURL, + apiKey: apiKey, + httpClient: &http.Client{Timeout: timeout}, + } +} + +type tavilyRequest struct { + APIKey string `json:"api_key"` + Query string `json:"query"` + MaxResults int `json:"max_results"` + SearchDepth string `json:"search_depth"` +} + +type tavilyResponse struct { + Results []struct { + Title string `json:"title"` + URL string `json:"url"` + Content string `json:"content"` + } `json:"results"` +} + +func (c *TavilyClient) Search(ctx context.Context, query string, maxResults int) ([]Source, error) { + body, err := json.Marshal(tavilyRequest{ + APIKey: c.apiKey, + Query: query, + MaxResults: maxResults, + SearchDepth: "basic", + }) + if err != nil { + return nil, fmt.Errorf("marshal tavily request: %w", err) + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/search", bytes.NewReader(body)) + if err != nil { + return nil, fmt.Errorf("build tavily request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + + resp, err := c.httpClient.Do(req) + if err != nil { + return nil, fmt.Errorf("tavily request: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("tavily returned status %d", resp.StatusCode) + } + + var parsed tavilyResponse + if err := json.NewDecoder(resp.Body).Decode(&parsed); err != nil { + return nil, fmt.Errorf("decode tavily response: %w", err) + } + + sources := make([]Source, 0, len(parsed.Results)) + for i, r := range parsed.Results { + sources = append(sources, Source{ + Position: i + 1, + Title: r.Title, + URL: r.URL, + Snippet: r.Content, + }) + } + return sources, nil +} diff --git a/internal/search/client_test.go b/internal/search/client_test.go new file mode 100644 index 0000000..31b765c --- /dev/null +++ b/internal/search/client_test.go @@ -0,0 +1,44 @@ +package search + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestTavilyClient_Search(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, "/search", r.URL.Path) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"results":[ + {"title":"T1","url":"https://a.example","content":"snippet one"}, + {"title":"T2","url":"https://b.example","content":"snippet two"} + ]}`)) + })) + defer srv.Close() + + c := NewTavilyClient(srv.URL, "tvly-test", 5*time.Second) + got, err := c.Search(context.Background(), "what is lightchain", 2) + require.NoError(t, err) + require.Len(t, got, 2) + assert.Equal(t, 1, got[0].Position) + assert.Equal(t, "T1", got[0].Title) + assert.Equal(t, "https://a.example", got[0].URL) + assert.Equal(t, "snippet one", got[0].Snippet) + assert.Equal(t, 2, got[1].Position) +} + +func TestTavilyClient_HTTPError(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + })) + defer srv.Close() + c := NewTavilyClient(srv.URL, "tvly-test", 5*time.Second) + _, err := c.Search(context.Background(), "q", 3) + require.Error(t, err) +} diff --git a/internal/service/service.go b/internal/service/service.go index 9af8949..80ba304 100644 --- a/internal/service/service.go +++ b/internal/service/service.go @@ -36,6 +36,7 @@ import ( "github.com/lightchain/worker/internal/pipeline" "github.com/lightchain/worker/internal/registration" "github.com/lightchain/worker/internal/release" + "github.com/lightchain/worker/internal/search" ) // startupHeartbeatTimeout is the maximum time allowed for the initial heartbeat @@ -406,6 +407,11 @@ func New(cfg *config.Config) (*Service, error) { } // Job pipeline handler + var searcher search.Searcher + if cfg.SearchEnabled { + searcher = search.NewTavilyClient(cfg.TavilyURL, cfg.TavilyAPIKey, cfg.SearchTimeout) + } + handler := pipeline.NewJobHandler( chainClient, blobFetcher, @@ -424,11 +430,14 @@ func New(cfg *config.Config) (*Service, error) { ModelIDToName: modelIDToName, ChainID: big.NewInt(cfg.ChainID), JobRegistryAddr: cfg.JobRegistryAddress, + SearchMaxResults: cfg.SearchMaxResults, + SearchTimeout: cfg.SearchTimeout, }, nil, // publisher — fallback wires RedisResponsePublisher from redisClient checkpoints, metricsCollector, metrics.DeliveryAsynq, + searcher, ) handler.SetReleaseTracker(releaseTracker) @@ -465,16 +474,19 @@ func New(cfg *config.Config) (*Service, error) { jobCounter, logger, pipeline.HandlerConfig{ - AckTxTimeout: cfg.AckTxTimeout, - BlobTxTimeout: cfg.BlobTxTimeout, - ModelIDToName: modelIDToName, - ChainID: big.NewInt(cfg.ChainID), - JobRegistryAddr: cfg.JobRegistryAddress, + AckTxTimeout: cfg.AckTxTimeout, + BlobTxTimeout: cfg.BlobTxTimeout, + ModelIDToName: modelIDToName, + ChainID: big.NewInt(cfg.ChainID), + JobRegistryAddr: cfg.JobRegistryAddress, + SearchMaxResults: cfg.SearchMaxResults, + SearchTimeout: cfg.SearchTimeout, }, gwPublisher, checkpoints, metricsCollector, metrics.DeliveryGateway, + searcher, ) gwHandler.SetReleaseTracker(releaseTracker) @@ -541,8 +553,13 @@ func New(cfg *config.Config) (*Service, error) { Interval: cfg.HeartbeatInterval, OllamaURL: cfg.OllamaURL, } + capabilities := []string{} + if cfg.SearchEnabled && cfg.TavilyAPIKey != "" { + capabilities = append(capabilities, "search") + } + addrHex := checksumHexNoPrefix(workerAddr) - monitor := heartbeat.NewMonitor(redisClient, monitorCfg, addrHex, modelHexStrings, jobCounter, cfg.MaxConcurrentJobs, logger, metricsCollector) + monitor := heartbeat.NewMonitor(redisClient, monitorCfg, addrHex, modelHexStrings, capabilities, jobCounter, cfg.MaxConcurrentJobs, logger, metricsCollector) // Gate startup on a real heartbeat write startCtx, startCancel := context.WithTimeout(context.Background(), startupHeartbeatTimeout) @@ -658,6 +675,29 @@ func (p *gatewayResponsePublisher) PublishResponse( } } +func (p *gatewayResponsePublisher) PublishMetadata( + _ context.Context, + jobID, _ uint64, + _ string, + _ []byte, +) { + // v1 no-op: gateway mode routes metadata via the relay's Redis pub/sub path, + // not HTTP. Citations are best-effort; missing them does not break job delivery. + p.logger.Debug("gateway: metadata publish skipped (v1 no-op)", "jobID", jobID) +} + +func (p *gatewayResponsePublisher) PublishChunk( + _ context.Context, + jobID, _ uint64, + _ string, + seq uint32, + _ []byte, +) { + // v1 no-op: gateway mode does not forward chunk frames over HTTP. + // Streaming is best-effort UX; missing chunks do not break job delivery. + p.logger.Debug("gateway: chunk publish skipped (v1 no-op)", "jobID", jobID, "seq", seq) +} + // Run starts the heartbeat goroutine and Asynq server (or gateway poll loop), // waits for SIGINT/SIGTERM, then gracefully shuts down. func (s *Service) Run(ctx context.Context) error { diff --git a/workertest/harness.go b/workertest/harness.go index ac51bc8..ad24c2f 100644 --- a/workertest/harness.go +++ b/workertest/harness.go @@ -61,6 +61,8 @@ type ChatMessage = ollama.ChatMessage type InferenceClient interface { Generate(ctx context.Context, model, prompt string) (string, error) Chat(ctx context.Context, model string, messages []ChatMessage) (string, error) + GenerateStream(ctx context.Context, model, prompt string, onDelta func(string)) (string, error) + ChatStream(ctx context.Context, model string, messages []ChatMessage, onDelta func(string)) (string, error) } // Options configures the in-process worker harness. @@ -200,6 +202,7 @@ func New(t testing.TB, opts Options) *Harness { // values so NormalizeModel returns the actual tag, not "unknown". metrics.New(modelTagsFromMap(opts.ModelIDToName)), metrics.DeliveryAsynq, + nil, // searcher — web search disabled in test harness ) redisConnOpt := asynq.RedisClientOpt{