fix(selfhosted): align heartbeat ownership recovery

arkruntime/environment_work.go

@@ -108,6 +108,9 @@ func (c *Client) HeartbeatWork(
 	if expectedLastHeartbeat, ok := body.ExpectedLastHeartbeat.Get(); ok && expectedLastHeartbeat != "" {
 		q.Set("expected_last_heartbeat", expectedLastHeartbeat)
 	}
+	if workerID, ok := body.WorkerID.Get(); ok && workerID != "" {
+		setters = append(setters, WithCustomHeader(environmentWorkWorkerIDHeader, workerID))
+	}
 	u := c.fullURL(fmt.Sprintf("%s/%s/work/%s/heartbeat",
 		environmentsPrefix,
 		environment.PathEscape(body.EnvironmentID),

arkruntime/lib/environments/heartbeat.go

@@ -41,6 +41,7 @@ func (w *EnvironmentWorker) heartbeatLoop(ctx context.Context, work claimedWork,
 		resp, err := api.HeartbeatWork(beatCtx, selfhosted.HeartbeatWorkRequest{
 			EnvironmentID:         work.EnvironmentID,
 			WorkID:                work.ID,
+			WorkerID:              environment.NewOptString(w.opts.WorkerID),
 			ExpectedLastHeartbeat: environment.NewOptString(last),
 			DesiredTTLSeconds:     environment.NewOptInt64(int64(ttl / time.Second)),
 		})

arkruntime/lib/environments/integration_test.go

@@ -450,6 +450,9 @@ func TestEnvironmentWorkerHeartbeatLeaseLostOnPreconditionFailed(t *testing.T) {
 	if value, ok := got.ExpectedLastHeartbeat.Get(); !ok || value != selfhosted.ExpectedLastHeartbeatNoHeartbeat {
 		t.Fatalf("expected_last_heartbeat=%+v", got.ExpectedLastHeartbeat)
 	}
+	if value, ok := got.WorkerID.Get(); !ok || value != "worker_local" {
+		t.Fatalf("worker_id=%+v", got.WorkerID)
+	}
 	if value, ok := got.DesiredTTLSeconds.Get(); !ok || value != int64(heartbeatDefault/time.Second) {
 		t.Fatalf("desired_ttl_seconds=%+v", got.DesiredTTLSeconds)
 	}

arkruntime/model/environment/work_shim.go

@@ -45,6 +45,7 @@ type AckWorkRequest struct {
 type HeartbeatWorkRequest struct {
 	EnvironmentID         string    `json:"environment_id"`
 	WorkID                string    `json:"work_id"`
+	WorkerID              OptString `json:"worker_id,omitempty"`
 	ExpectedLastHeartbeat OptString `json:"expected_last_heartbeat,omitempty"`
 	DesiredTTLSeconds     OptInt64  `json:"desired_ttl_seconds,omitempty"`
 }

arkruntime/self_hosted_client_test.go

@@ -17,7 +17,10 @@ import (
 	"github.com/volcengine/ark-runtime-go/arkruntime/model/session"
 )
 
-const testBearerToken = "Bearer test-api-key"
+const (
+	testBearerToken = "Bearer test-api-key"
+	testWorkerID    = "worker-1"
+)
 
 func TestEnvironmentWorkRequests(t *testing.T) {
 	var seen []string
@@ -33,7 +36,7 @@ func TestEnvironmentWorkRequests(t *testing.T) {
 
 		switch r.Method + " " + r.URL.Path {
 		case "GET /environments/env-1/work/poll":
-			if got := r.Header.Get(environmentWorkWorkerIDHeader); got != "worker-1" {
+			if got := r.Header.Get(environmentWorkWorkerIDHeader); got != testWorkerID {
 				t.Fatalf("Ark-Worker-ID = %q", got)
 			}
 			if got := r.URL.Query().Get("block_ms"); got != "999" {
@@ -53,20 +56,22 @@ func TestEnvironmentWorkRequests(t *testing.T) {
 				"latest_heartbeat_at": "2026-08-10T00:00:00Z",
 				"state":               "queued",
 				"stop_reason":         "worker_abnormal",
-				"recovery_count":      5,
 				"type":                "work",
 			}
 			if err := json.NewEncoder(w).Encode(pollResponse); err != nil {
 				t.Errorf("encode poll response: %v", err)
 			}
 		case "POST /environments/env-1/work/work-1/ack":
 			assertNoBody(t, r)
-			if got := r.Header.Get(environmentWorkWorkerIDHeader); got != "worker-1" {
+			if got := r.Header.Get(environmentWorkWorkerIDHeader); got != testWorkerID {
 				t.Fatalf("Ark-Worker-ID = %q", got)
 			}
 			_, _ = w.Write([]byte(`{"id":"work-1","created_at":"2026-08-10T00:00:00Z","data":{"id":"sess-1","type":"session"},"environment_id":"env-1","state":"starting","type":"work"}`))
 		case "POST /environments/env-1/work/work-1/heartbeat":
 			assertNoBody(t, r)
+			if got := r.Header.Get(environmentWorkWorkerIDHeader); got != testWorkerID {
+				t.Fatalf("Ark-Worker-ID = %q", got)
+			}
 			if got := r.URL.Query().Get("desired_ttl_seconds"); got != "30" {
 				t.Fatalf("desired_ttl_seconds = %q", got)
 			}
@@ -95,7 +100,7 @@ func TestEnvironmentWorkRequests(t *testing.T) {
 
 	work, err := client.PollWork(ctx, &environment.PollWorkRequest{
 		EnvironmentID: "env-1",
-		WorkerID:      "worker-1",
+		WorkerID:      testWorkerID,
 		BlockMS:       999,
 	})
 	if err != nil {
@@ -107,19 +112,17 @@ func TestEnvironmentWorkRequests(t *testing.T) {
 	if got := work.LatestHeartbeatValue(); got != "2026-08-10T00:00:00Z" {
 		t.Fatalf("PollWork().LatestHeartbeatValue() = %q", got)
 	}
-	if got, ok := work.RecoveryCount.Get(); !ok || got != 5 {
-		t.Fatalf("PollWork().RecoveryCount = %d, set=%v", got, ok)
-	}
 	if err := client.AckWork(ctx, &environment.AckWorkRequest{
 		EnvironmentID: "env-1",
 		WorkID:        "work-1",
-		WorkerID:      environment.NewOptString("worker-1"),
+		WorkerID:      environment.NewOptString(testWorkerID),
 	}); err != nil {
 		t.Fatalf("AckWork() error = %v", err)
 	}
 	heartbeat, err := client.HeartbeatWork(ctx, &environment.HeartbeatWorkRequest{
 		EnvironmentID:         "env-1",
 		WorkID:                "work-1",
+		WorkerID:              environment.NewOptString(testWorkerID),
 		DesiredTTLSeconds:     environment.NewOptInt64(30),
 		ExpectedLastHeartbeat: environment.NewOptString("2026-08-10T00:00:00Z"),
 	})