APP-4344 Propagate Oz lifecycle hook metadata through cloud workers

README.md

@@ -123,7 +123,7 @@ backend:
 Config keys:
 
 - `dispatch_command` (required): shell command (run via `/bin/sh -c`) invoked once per task to dispatch it.
-- `cancel_command` (optional): shell command invoked best-effort when a dispatched task is cancelled. If unset, the worker relies on agent-side cancellation.
+- `cancel_command` (optional for ordinary tasks, required for hook-enabled tasks): shell command invoked best-effort when a dispatched task is cancelled. If unset, the worker relies on agent-side cancellation and rejects hook-enabled task assignments before claiming them.
 - `dispatch_timeout` (optional): how long the dispatch command may run before it is considered failed (humantime format, e.g. `60s`). Defaults to `60s`.
 - `environment`: extra environment variables exposed to the dispatch/cancel commands (same `name`/`value` semantics as the other backends; omit `value` to inherit from the host).
 
@@ -135,7 +135,7 @@ The dispatch contract:
 
   ```json
   {
-    "version": 1,
+    "version": 2,
     "run_id": "...",
     "execution_id": "...",
     "server_root_url": "https://app.warp.dev",
@@ -144,11 +144,19 @@ The dispatch contract:
     "base_args": ["agent", "run", "--task-id", "...", "--server-root-url", "..."],
     "env": { "GITHUB_ACCESS_TOKEN": "...", "...": "..." },
     "sidecars": [ { "image": "...", "mount_path": "/agent", "read_write": false } ],
-    "task": { "id": "...", "title": "...", "task_definition": { "prompt": "..." } }
+    "task": { "id": "...", "title": "...", "task_definition": { "prompt": "..." } },
+    "oz_lifecycle_hooks": {
+      "required": true,
+      "supported_payload_schema_versions": ["warp.oz_hook.v1"],
+      "project_trust": [
+        { "git_root": "/workspace", "config_path": "/workspace/.warp/hooks.json", "sha256": "..." }
+      ]
+    }
   }
   ```
 
   `base_args` is the `oz agent run …` argument vector your runtime should launch the agent with, inside an environment built from `docker_image` and `sidecars`.
+  Hook-enabled payloads include `oz_lifecycle_hooks` and the same context in `base_args` under `--oz-lifecycle-hooks-context`; runtimes must preserve both unchanged. The serialized context is limited to 64 KiB so it remains safely below Linux's per-argument limit. The worker accepts these payloads only when `cancel_command` is configured, so the remote runtime and any hook subprocesses can be contained by task cancellation.
 - Exit code `0` means the task was dispatched successfully; the worker will not finalize it (the remote agent reports terminal state to Warp itself). A non-zero exit or a dispatch that exceeds `dispatch_timeout` marks the task failed.
 - The cancel command (when configured) receives `OZ_RUN_ID`, `OZ_EXECUTION_ID`, and `OZ_WORKER_BACKEND=command` in its environment, each with its `WARP_`-prefixed alias.
 

internal/metrics/metrics.go

@@ -89,6 +89,7 @@ var activeInstruments atomic.Pointer[instruments]
 // and alerts can query them by name even before the worker has handled a task.
 const (
 	RejectReasonAtCapacity       = "at_capacity"
+	RejectReasonIncompatibleTask = "incompatible_task"
 	WSReconnectReasonDialFailed  = "dial_failed"
 	WSReconnectReasonRemoteClose = "remote_close"
 )
@@ -302,9 +303,11 @@ func primeInstruments(ctx context.Context, set *instruments) {
 	set.tasksActive.Add(ctx, 0)
 	set.tasksMaxConcurrent.Record(ctx, 0)
 	set.tasksClaimed.Add(ctx, 0)
-	set.tasksRejected.Add(ctx, 0,
-		metric.WithAttributes(attribute.String("reason", RejectReasonAtCapacity)),
-	)
+	for _, reason := range []string{RejectReasonAtCapacity, RejectReasonIncompatibleTask} {
+		set.tasksRejected.Add(ctx, 0,
+			metric.WithAttributes(attribute.String("reason", reason)),
+		)
+	}
 	for _, r := range taskResults {
 		set.tasksCompleted.Add(ctx, 0,
 			metric.WithAttributes(attribute.String("result", string(r))),

internal/types/messages.go

@@ -62,6 +62,10 @@ type TaskAssignmentMessage struct {
 	// backends size the task container/pod from it; omitted when the run has no explicit
 	// runner instance shape, in which case the worker keeps its default sizing.
 	InstanceShape *InstanceShape `json:"instance_shape,omitempty"`
+	// OzLifecycleHooks carries the authenticated, non-secret context required by
+	// the embedded first-party Oz runtime.
+	OzLifecycleHooks      *OzLifecycleHooksContext `json:"oz_lifecycle_hooks,omitempty"`
+	ozLifecycleHooksError error
 }
 
 // TaskClaimedMessage is sent from worker to server after successfully claiming a task

internal/types/oz_lifecycle_hooks.go

@@ -0,0 +1,148 @@
+package types
+
+import (
+	"bytes"
+	"encoding/hex"
+	"encoding/json"
+	"fmt"
+	"io"
+	"path/filepath"
+	"strings"
+)
+
+const (
+	OzHookPayloadSchemaV1          = "warp.oz_hook.v1"
+	MaxOzLifecycleHooksContextSize = 64 << 10
+	MaxOzLifecycleHookTrustRecords = 64
+)
+
+func (m *TaskAssignmentMessage) UnmarshalJSON(data []byte) error {
+	type assignmentAlias TaskAssignmentMessage
+	var decoded struct {
+		*assignmentAlias
+		OzLifecycleHooks json.RawMessage `json:"oz_lifecycle_hooks"`
+	}
+	decoded.assignmentAlias = (*assignmentAlias)(m)
+	if err := json.Unmarshal(data, &decoded); err != nil {
+		return err
+	}
+
+	m.OzLifecycleHooks = nil
+	m.ozLifecycleHooksError = nil
+	if len(decoded.OzLifecycleHooks) == 0 || bytes.Equal(decoded.OzLifecycleHooks, []byte("null")) {
+		return nil
+	}
+	var context OzLifecycleHooksContext
+	if err := json.Unmarshal(decoded.OzLifecycleHooks, &context); err != nil {
+		m.ozLifecycleHooksError = err
+		return nil
+	}
+	m.OzLifecycleHooks = &context
+	return nil
+}
+
+func (m *TaskAssignmentMessage) OzLifecycleHooksValidationError() error {
+	if m == nil {
+		return nil
+	}
+	return m.ozLifecycleHooksError
+}
+
+type OzLifecycleHooksContext struct {
+	Required                       bool                         `json:"required"`
+	SupportedPayloadSchemaVersions []string                     `json:"supported_payload_schema_versions"`
+	ProjectTrust                   []OzLifecycleHookTrustRecord `json:"project_trust"`
+}
+
+type OzLifecycleHookTrustRecord struct {
+	GitRoot    string `json:"git_root"`
+	ConfigPath string `json:"config_path"`
+	SHA256     string `json:"sha256"`
+}
+
+func (c *OzLifecycleHooksContext) UnmarshalJSON(data []byte) error {
+	if len(data) > MaxOzLifecycleHooksContextSize {
+		return fmt.Errorf("oz lifecycle hooks context exceeds %d bytes", MaxOzLifecycleHooksContextSize)
+	}
+
+	type contextAlias OzLifecycleHooksContext
+	var decoded contextAlias
+	decoder := json.NewDecoder(bytes.NewReader(data))
+	decoder.DisallowUnknownFields()
+	if err := decoder.Decode(&decoded); err != nil {
+		return fmt.Errorf("invalid oz lifecycle hooks context: %w", err)
+	}
+	if err := ensureJSONEOF(decoder); err != nil {
+		return fmt.Errorf("invalid oz lifecycle hooks context: %w", err)
+	}
+
+	*c = OzLifecycleHooksContext(decoded)
+	return c.Validate()
+}
+
+func (c *OzLifecycleHooksContext) Validate() error {
+	if c == nil {
+		return nil
+	}
+	if !c.Required {
+		return fmt.Errorf("oz lifecycle hooks context must set required to true")
+	}
+	if len(c.SupportedPayloadSchemaVersions) == 0 {
+		return fmt.Errorf("oz lifecycle hooks context requires a supported payload schema version")
+	}
+	for _, version := range c.SupportedPayloadSchemaVersions {
+		if version != OzHookPayloadSchemaV1 {
+			return fmt.Errorf("unsupported oz lifecycle hook payload schema version %q", version)
+		}
+	}
+	if len(c.ProjectTrust) > MaxOzLifecycleHookTrustRecords {
+		return fmt.Errorf("oz lifecycle hooks context has %d project trust records; maximum is %d", len(c.ProjectTrust), MaxOzLifecycleHookTrustRecords)
+	}
+	if c.ProjectTrust == nil {
+		return fmt.Errorf("oz lifecycle hooks context requires project_trust")
+	}
+	for i, record := range c.ProjectTrust {
+		if !filepath.IsAbs(record.GitRoot) || filepath.Clean(record.GitRoot) != record.GitRoot {
+			return fmt.Errorf("project trust record %d has a non-canonical git_root", i)
+		}
+		expectedConfigPath := filepath.Join(record.GitRoot, ".warp", "hooks.json")
+		if record.ConfigPath != expectedConfigPath {
+			return fmt.Errorf("project trust record %d has an invalid config_path", i)
+		}
+		hash, err := hex.DecodeString(record.SHA256)
+		if err != nil || len(hash) != 32 || strings.ToLower(record.SHA256) != record.SHA256 {
+			return fmt.Errorf("project trust record %d has an invalid sha256", i)
+		}
+	}
+
+	data, err := json.Marshal(c)
+	if err != nil {
+		return fmt.Errorf("marshal oz lifecycle hooks context: %w", err)
+	}
+	if len(data) > MaxOzLifecycleHooksContextSize {
+		return fmt.Errorf("oz lifecycle hooks context exceeds %d bytes", MaxOzLifecycleHooksContextSize)
+	}
+	return nil
+}
+
+func (c *OzLifecycleHooksContext) MarshalForCLI() (string, error) {
+	if err := c.Validate(); err != nil {
+		return "", err
+	}
+	data, err := json.Marshal(c)
+	if err != nil {
+		return "", fmt.Errorf("marshal oz lifecycle hooks context: %w", err)
+	}
+	return string(data), nil
+}
+
+func ensureJSONEOF(decoder *json.Decoder) error {
+	var extra any
+	if err := decoder.Decode(&extra); err != io.EOF {
+		if err == nil {
+			return fmt.Errorf("multiple JSON values")
+		}
+		return err
+	}
+	return nil
+}

internal/types/oz_lifecycle_hooks_test.go

@@ -0,0 +1,105 @@
+package types
+
+import (
+	"encoding/json"
+	"fmt"
+	"strings"
+	"testing"
+)
+
+func TestOzLifecycleHooksContextStrictUnmarshal(t *testing.T) {
+	valid := fmt.Sprintf(`{
+		"required": true,
+		"supported_payload_schema_versions": ["%s"],
+		"project_trust": [{
+			"git_root": "/workspace/repo",
+			"config_path": "/workspace/repo/.warp/hooks.json",
+			"sha256": "%s"
+		}]
+	}`, OzHookPayloadSchemaV1, strings.Repeat("a", 64))
+
+	var context OzLifecycleHooksContext
+	if err := json.Unmarshal([]byte(valid), &context); err != nil {
+		t.Fatalf("valid context rejected: %v", err)
+	}
+	if !context.Required || len(context.ProjectTrust) != 1 {
+		t.Fatalf("unexpected decoded context: %+v", context)
+	}
+
+	tests := []struct {
+		name string
+		json string
+	}{
+		{
+			name: "unknown context field",
+			json: strings.Replace(valid, `"required": true`, `"required": true, "unknown": true`, 1),
+		},
+		{
+			name: "unknown trust field",
+			json: strings.Replace(valid, `"git_root":`, `"unknown": true, "git_root":`, 1),
+		},
+		{
+			name: "required false",
+			json: strings.Replace(valid, `"required": true`, `"required": false`, 1),
+		},
+		{
+			name: "empty schema versions",
+			json: strings.Replace(valid, `["warp.oz_hook.v1"]`, `[]`, 1),
+		},
+		{
+			name: "unsupported schema version",
+			json: strings.Replace(valid, OzHookPayloadSchemaV1, "warp.oz_hook.v2", 1),
+		},
+		{
+			name: "missing project trust",
+			json: fmt.Sprintf(`{"required":true,"supported_payload_schema_versions":["%s"]}`, OzHookPayloadSchemaV1),
+		},
+		{
+			name: "non-canonical git root",
+			json: strings.Replace(valid, `"/workspace/repo"`, `"workspace/repo"`, 1),
+		},
+		{
+			name: "mismatched config path",
+			json: strings.Replace(valid, `"/workspace/repo/.warp/hooks.json"`, `"/workspace/other/.warp/hooks.json"`, 1),
+		},
+		{
+			name: "invalid sha256",
+			json: strings.Replace(valid, strings.Repeat("a", 64), "not-a-hash", 1),
+		},
+		{
+			name: "non-canonical sha256",
+			json: strings.Replace(valid, strings.Repeat("a", 64), strings.Repeat("A", 64), 1),
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			var context OzLifecycleHooksContext
+			if err := json.Unmarshal([]byte(test.json), &context); err == nil {
+				t.Fatal("expected context to be rejected")
+			}
+		})
+	}
+}
+
+func TestOzLifecycleHooksContextLimits(t *testing.T) {
+	context := OzLifecycleHooksContext{
+		Required:                       true,
+		SupportedPayloadSchemaVersions: []string{OzHookPayloadSchemaV1},
+		ProjectTrust:                   make([]OzLifecycleHookTrustRecord, MaxOzLifecycleHookTrustRecords+1),
+	}
+	if err := context.Validate(); err == nil {
+		t.Fatal("expected trust record limit to be enforced")
+	}
+
+	oversized := fmt.Sprintf(
+		`{"required":true,"supported_payload_schema_versions":["%s"],"project_trust":[],"padding":"%s"}`,
+		OzHookPayloadSchemaV1,
+		strings.Repeat("x", MaxOzLifecycleHooksContextSize),
+	)
+	var decoded OzLifecycleHooksContext
+	err := json.Unmarshal([]byte(oversized), &decoded)
+	if err == nil || !strings.Contains(err.Error(), "exceeds") {
+		t.Fatalf("expected size-limit error, got %v", err)
+	}
+}

internal/worker/backend.go

@@ -131,6 +131,9 @@ type TaskParams struct {
 	TaskID      string
 	ExecutionID string
 	Task        *types.Task
+	// OzLifecycleHooks is the authenticated, non-secret hook context forwarded
+	// to the embedded first-party Oz runtime.
+	OzLifecycleHooks *types.OzLifecycleHooksContext
 
 	// EnvVars contains pre-resolved common environment variables (TASK_ID, Git config,
 	// assignment env vars). Backends append their own config-specific env vars.
@@ -174,6 +177,9 @@ type Backend interface {
 	// PreservesTasksOnShutdown reports whether active task execution units can
 	// safely outlive the worker process during shutdown.
 	PreservesTasksOnShutdown() bool
+	// SupportsOzLifecycleHooks reports whether the backend preserves Oz argv,
+	// sandbox placement, and task cancellation for hook-enabled tasks.
+	SupportsOzLifecycleHooks() bool
 	// Shutdown cleans up backend resources.
 	Shutdown(ctx context.Context)
 }

internal/worker/command.go

@@ -150,6 +150,9 @@ func (b *CommandBackend) CancelTask(ctx context.Context, params *CancelParams) e
 // PreservesTasksOnShutdown reports true: dispatched tasks run on the operator's
 // remote runtime, independent of this worker, so worker shutdown must not cancel them.
 func (b *CommandBackend) PreservesTasksOnShutdown() bool { return true }
+func (b *CommandBackend) SupportsOzLifecycleHooks() bool {
+	return b.config.CancelCommand != ""
+}
 
 // Shutdown has nothing to clean up; the backend owns no local resources.
 func (b *CommandBackend) Shutdown(ctx context.Context) {