Add persistent job claim and file artifact reads

runtime/worker_test.go

@@ -7,6 +7,8 @@ import (
 	"fmt"
 	"net/http"
 	"net/http/httptest"
+	"os"
+	"path/filepath"
 	"reflect"
 	"runtime"
 	"strings"
@@ -173,6 +175,9 @@ func TestWorkerClaimsAcksProgressAndCompletesJob(t *testing.T) {
 	if len(client.ackRequests) != 1 || len(client.progressRequests) != 1 || len(client.resultRequests) != 1 || len(client.cancelPollRequests) != 1 {
 		t.Fatalf("expected ack/progress/result/cancel calls, got ack=%d progress=%d result=%d cancel=%d", len(client.ackRequests), len(client.progressRequests), len(client.resultRequests), len(client.cancelPollRequests))
 	}
+	if len(client.claimRequests) != 1 || client.claimRequests[0].WaitSeconds != jobClaimWaitSeconds {
+		t.Fatalf("expected persistent claim wait, got %+v", client.claimRequests)
+	}
 	if client.progressRequests[0].Progress.Percent != 10 || client.resultRequests[0].State != "succeeded" || client.resultRequests[0].ResultRef == "" {
 		t.Fatalf("unexpected job channel payloads: progress=%+v result=%+v", client.progressRequests[0], client.resultRequests[0])
 	}
@@ -214,6 +219,56 @@ func TestWorkerDispatchesSelfUpdateJob(t *testing.T) {
 	}
 }
 
+func TestWorkerFileReadLargeFileUploadsArtifact(t *testing.T) {
+	cfg := workerTestConfig(t)
+	client := newFakeWorkerClient()
+	assignment := workerJobAssignment(protocol.RunCapabilityFilesRead)
+	assignment.TargetKey = "logs/big.log"
+	assignment.ExecutionInput.WorkspaceScope = "run-local"
+	assignment.ExecutionInput.MaxReadBytes = 4
+	client.claimJob = assignment
+	client.buildInput = protocol.DistributionBuildInputResponse{JobID: assignment.JobID}
+	scope, err := NewWorkspaceResolver(cfg.WorkspaceRoot).Scope(assignment.ServerInstanceID, assignment.ExecutionInput.WorkspaceScope)
+	if err != nil {
+		t.Fatalf("scope: %v", err)
+	}
+	if err := os.MkdirAll(filepath.Join(scope, "logs"), 0o700); err != nil {
+		t.Fatalf("mkdir fixture: %v", err)
+	}
+	payload := []byte(strings.Repeat("A", fileArtifactChunkSize) + "tail")
+	if err := os.WriteFile(filepath.Join(scope, assignment.TargetKey), payload, 0o600); err != nil {
+		t.Fatalf("write fixture: %v", err)
+	}
+	worker, err := NewWorker(cfg, client)
+	if err != nil {
+		t.Fatalf("new worker: %v", err)
+	}
+	if err := worker.Register(context.Background()); err != nil {
+		t.Fatalf("register: %v", err)
+	}
+
+	handled, err := worker.ClaimAndRunOnce(context.Background())
+	if err != nil || !handled {
+		t.Fatalf("claim/run handled=%v err=%v", handled, err)
+	}
+	if len(client.artifactOpenRequests) != 1 {
+		t.Fatalf("expected artifact transfer, got %+v", client.artifactOpenRequests)
+	}
+	opened := client.artifactOpenRequests[0]
+	if opened.OwnerID != assignment.JobID || opened.SizeBytes != int64(len(payload)) || opened.Checksum != bytesChecksum(payload) || opened.ChunkSizeBytes != fileArtifactChunkSize {
+		t.Fatalf("unexpected artifact open: %+v", opened)
+	}
+	if string(client.artifactPayload) != string(payload) {
+		t.Fatalf("uploaded artifact payload mismatch")
+	}
+	if len(client.resultRequests) != 1 || client.resultRequests[0].State != "succeeded" || !strings.HasPrefix(client.resultRequests[0].ResultRef, "artifact://artifact-job-worker-file-read") {
+		t.Fatalf("expected artifact result ref, got %+v", client.resultRequests)
+	}
+	if client.resultRequests[0].ExecutionResult.Content != "" || client.resultRequests[0].ExecutionResult.SizeBytes != int64(len(payload)) {
+		t.Fatalf("large file result must not inline content: %+v", client.resultRequests[0].ExecutionResult)
+	}
+}
+
 func TestWorkerRegistersPackageIdentity(t *testing.T) {
 	client := newFakeWorkerClient()
 	cfg := workerTestConfig(t)
@@ -608,6 +663,19 @@ func TestSessionArtifactChunkClientOverridesSpooledIdentity(t *testing.T) {
 	}
 }
 
+func TestSessionArtifactChunkClientDropsPlatformMissingTransfer(t *testing.T) {
+	recorder := &recordingDurableArtifactClient{err: api.PlatformRequestError{Status: http.StatusNotFound, Code: "not_found"}}
+	chunk := protocol.ArtifactChunkUploadRequest{RunEndpointID: "old-endpoint", SessionToken: "old-token", TransferID: "transfer-stale", ArtifactID: "artifact-stale", ChunkIndex: 3}
+	client := sessionArtifactChunkClient{client: recorder, runEndpointID: "run-current", sessionToken: "token-current"}
+	response, err := client.UploadArtifactChunk(context.Background(), chunk)
+	if err != nil || !response.Accepted || response.TransferID != chunk.TransferID || response.ChunkIndex != chunk.ChunkIndex {
+		t.Fatalf("expected stale chunk ack for queue cleanup, response=%+v err=%v", response, err)
+	}
+	if recorder.chunk.RunEndpointID != "run-current" || recorder.chunk.SessionToken != "token-current" {
+		t.Fatalf("expected current identity before stale drop, got %+v", recorder.chunk)
+	}
+}
+
 func TestWorkerRetryBackoffIsBounded(t *testing.T) {
 	if got := boundedRetryBackoff(75 * time.Millisecond); got != 75*time.Millisecond {
 		t.Fatalf("expected configured backoff, got %s", got)
@@ -640,7 +708,7 @@ func TestWorkerIntegrationWithPlatformLikeServer(t *testing.T) {
 		case "/api/v1/run/jobs/claim":
 			var request protocol.RunJobClaimRequest
 			decodeWorkerTestJSON(t, r, &request)
-			if request.SessionToken != "session-token" || len(request.Capabilities) == 0 {
+			if request.SessionToken != "session-token" || len(request.Capabilities) == 0 || request.WaitSeconds != jobClaimWaitSeconds {
 				t.Fatalf("unexpected claim: %+v", request)
 			}
 			writeWorkerTestJSON(t, w, protocol.RunJobClaimResponse{Accepted: true, RunEndpointID: request.RunEndpointID, HasJob: true, Job: &assignment, NextPollSeconds: 2, ServerTime: workerTestTime()})
@@ -811,10 +879,14 @@ func (client *recordingDurableLogClient) IngestLogBatch(_ context.Context, batch
 
 type recordingDurableArtifactClient struct {
 	chunk protocol.ArtifactChunkUploadRequest
+	err   error
 }
 
 func (client *recordingDurableArtifactClient) UploadArtifactChunk(_ context.Context, chunk protocol.ArtifactChunkUploadRequest) (protocol.ArtifactChunkUploadResponse, error) {
 	client.chunk = chunk
+	if client.err != nil {
+		return protocol.ArtifactChunkUploadResponse{}, client.err
+	}
 	return protocol.ArtifactChunkUploadResponse{Accepted: true, TransferID: chunk.TransferID, ArtifactID: chunk.ArtifactID, ChunkIndex: chunk.ChunkIndex}, nil
 }