From da3617f99d0ed938c66b7885545b59acf204a967 Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Sat, 10 Jan 2026 10:18:42 +0000
Subject: [PATCH] Add TaskIDFromContext and ResourceKeyFromContext helpers

Add public helper functions to extract the current task ID and resource
key from context. This is useful for debugging when executing tasks
and resources within handlers.
---
 oncetask/context_logger.go      | 18 +++++++++++++
 oncetask/context_logger_test.go | 46 +++++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+)

diff --git a/oncetask/context_logger.go b/oncetask/context_logger.go
index 6b9d2b4..f005800 100644
--- a/oncetask/context_logger.go
+++ b/oncetask/context_logger.go
@@ -47,6 +47,24 @@ func withResourceKeyTaskContext[TaskKind ~string](ctx context.Context, tasks []O
 	return withTaskContext(ctx, taskID, tasks[0].ResourceKey)
 }
 
+// TaskIDFromContext returns the task ID stored in the context, or an empty string if not present.
+// This is useful for debugging or when you need to access the current task ID within a handler.
+func TaskIDFromContext(ctx context.Context) string {
+	if taskID, ok := ctx.Value(taskIDContextKey).(string); ok {
+		return taskID
+	}
+	return ""
+}
+
+// ResourceKeyFromContext returns the resource key stored in the context, or an empty string if not present.
+// This is useful for debugging or when you need to access the current resource key within a handler.
+func ResourceKeyFromContext(ctx context.Context) string {
+	if resourceKey, ok := ctx.Value(resourceKeyContextKey).(string); ok {
+		return resourceKey
+	}
+	return ""
+}
+
 // ContextHandler is a slog.Handler that automatically extracts the task ID and resource key from context
 // and adds them as attributes to all log records.
 //
diff --git a/oncetask/context_logger_test.go b/oncetask/context_logger_test.go
index ebc1551..72e9ab0 100644
--- a/oncetask/context_logger_test.go
+++ b/oncetask/context_logger_test.go
@@ -192,3 +192,49 @@ func TestContextHandler_WithAttrs(t *testing.T) {
 		t.Errorf("Expected service to be 'oncetask', got: %v", logEntry["service"])
 	}
 }
+
+func TestTaskIDFromContext(t *testing.T) {
+	t.Run("returns empty string when no task ID in context", func(t *testing.T) {
+		ctx := context.Background()
+		if got := TaskIDFromContext(ctx); got != "" {
+			t.Errorf("Expected empty string, got: %q", got)
+		}
+	})
+
+	t.Run("returns task ID when present in context", func(t *testing.T) {
+		ctx := withTaskContext(context.Background(), "test-task-123", "")
+		if got := TaskIDFromContext(ctx); got != "test-task-123" {
+			t.Errorf("Expected 'test-task-123', got: %q", got)
+		}
+	})
+
+	t.Run("returns task ID when both task ID and resource key are present", func(t *testing.T) {
+		ctx := withTaskContext(context.Background(), "task-456", "resource-789")
+		if got := TaskIDFromContext(ctx); got != "task-456" {
+			t.Errorf("Expected 'task-456', got: %q", got)
+		}
+	})
+}
+
+func TestResourceKeyFromContext(t *testing.T) {
+	t.Run("returns empty string when no resource key in context", func(t *testing.T) {
+		ctx := context.Background()
+		if got := ResourceKeyFromContext(ctx); got != "" {
+			t.Errorf("Expected empty string, got: %q", got)
+		}
+	})
+
+	t.Run("returns resource key when present in context", func(t *testing.T) {
+		ctx := withTaskContext(context.Background(), "", "user-123")
+		if got := ResourceKeyFromContext(ctx); got != "user-123" {
+			t.Errorf("Expected 'user-123', got: %q", got)
+		}
+	})
+
+	t.Run("returns resource key when both task ID and resource key are present", func(t *testing.T) {
+		ctx := withTaskContext(context.Background(), "task-456", "resource-789")
+		if got := ResourceKeyFromContext(ctx); got != "resource-789" {
+			t.Errorf("Expected 'resource-789', got: %q", got)
+		}
+	})
+}