Add task ID and resource key helpers to context

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.
 //

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)
+		}
+	})
+}