kdeps · jjuliano · Mar 10, 2025 · Mar 10, 2025 · Mar 10, 2025 · Mar 10, 2025
diff --git a/main.go b/main.go
@@ -13,6 +13,7 @@ import (
 	"github.com/kdeps/kdeps/pkg/cfg"
 	"github.com/kdeps/kdeps/pkg/docker"
 	"github.com/kdeps/kdeps/pkg/environment"
+	"github.com/kdeps/kdeps/pkg/ktx"
 	"github.com/kdeps/kdeps/pkg/logging"
 	"github.com/kdeps/kdeps/pkg/resolver"
 	"github.com/kdeps/kdeps/pkg/utils"
@@ -29,33 +30,27 @@ func main() {
 	v.Version = version
 	v.Commit = commit
 
-	// Initialize filesystem and context
+	logger := logging.GetLogger()
 	fs := afero.NewOsFs()
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel() // Ensure context is canceled when main exits
 
 	graphID := uuid.New().String()
-	logger := logging.GetLogger()
+	actionDir := filepath.Join(os.TempDir(), "action")
+	agentDir := filepath.Join("/", "agent")
 
 	// Setup environment
 	env, err := setupEnvironment(fs)
 	if err != nil {
 		logger.Fatalf("failed to set up environment: %v", err)
 	}
 
-	// Get the system's temporary directory (cross-platform)
-	baseTempDir := os.TempDir()
-
-	// Define the desired subdirectory for actions
-	actionDir := filepath.Join(baseTempDir, "action")
-
-	// Ensure the action directory exists (creating it if necessary)
-	if err := fs.MkdirAll(actionDir, 0o777); err != nil {
-		logger.Fatalf("failed to create action directory: %s", err)
-	}
+	ctx = ktx.CreateContext(ctx, ktx.CtxKeyGraphID, graphID)
+	ctx = ktx.CreateContext(ctx, ktx.CtxKeyActionDir, actionDir)
+	ctx = ktx.CreateContext(ctx, ktx.CtxKeyAgentDir, agentDir)
 
 	if env.DockerMode == "1" {
-		dr, err := resolver.NewGraphResolver(fs, ctx, env, "/agent", actionDir, graphID, logger.With("requestID", graphID))
+		dr, err := resolver.NewGraphResolver(fs, ctx, env, logger.With("requestID", graphID))
 		if err != nil {
 			logger.Fatalf("failed to create graph resolver: %v", err)
 		}

diff --git a/pkg/docker/api_server.go b/pkg/docker/api_server.go
@@ -17,6 +17,7 @@ import (
 	"github.com/gin-gonic/gin"
 	"github.com/google/uuid"
 	"github.com/kdeps/kdeps/pkg/evaluator"
+	"github.com/kdeps/kdeps/pkg/ktx"
 	"github.com/kdeps/kdeps/pkg/logging"
 	"github.com/kdeps/kdeps/pkg/resolver"
 	"github.com/kdeps/kdeps/pkg/utils"
@@ -192,7 +193,9 @@ func APIServerHandler(ctx context.Context, route *apiserver.APIServerRoutes, bas
 		baseLogger := logging.GetLogger()
 		logger := baseLogger.With("requestID", graphID) // Now returns *logging.Logger
 
-		dr, err := resolver.NewGraphResolver(baseDr.Fs, ctx, baseDr.Environment, baseDr.AgentDir, baseDr.ActionDir, graphID, logger)
+		newCtx := ktx.UpdateContext(ctx, ktx.CtxKeyGraphID, graphID)
+
+		dr, err := resolver.NewGraphResolver(baseDr.Fs, newCtx, baseDr.Environment, logger)
 		if err != nil {
 			resp := APIResponse{
 				Success: false,

diff --git a/pkg/docker/docker_test.go b/pkg/docker/docker_test.go
@@ -465,7 +465,7 @@ func itWillInstallTheModels(arg1 string) error {
 }
 
 func kdepsWillCheckThePresenceOfTheFile(arg1 string) error {
-	dr, err := resolver.NewGraphResolver(testFs, ctx, environ, "/agent", "/tmp/action", "123", logger)
+	dr, err := resolver.NewGraphResolver(testFs, ctx, environ, logger)
 	if err != nil {
 		return err
 	}

diff --git a/pkg/ktx/constants.go b/pkg/ktx/constants.go
@@ -0,0 +1,7 @@
+package ktx
+
+const (
+	CtxKeyGraphID   ContextKey = "graphID"
+	CtxKeyActionDir ContextKey = "actionDir"
+	CtxKeyAgentDir  ContextKey = "agentDir"
+)
diff --git a/pkg/ktx/context.go b/pkg/ktx/context.go
@@ -0,0 +1,34 @@
+package ktx
+
+import (
+	"context"
+	"fmt"
+)
+
+// ContextKey is a custom type for context keys to avoid key collisions.
+type ContextKey string
+
+// Create adds a key-value pair to the context.
+func CreateContext(ctx context.Context, key ContextKey, value any) context.Context {
+	return context.WithValue(ctx, key, value)
+}
+
+// Read retrieves a value from the context.
+func ReadContext(ctx context.Context, key ContextKey) (any, bool) {
+	value := ctx.Value(key) // No type assertion needed
+	return value, value != nil
+}
+
+// Update modifies an existing value in the context.
+func UpdateContext(ctx context.Context, key ContextKey, newValue any) context.Context {
+	if ctx.Value(key) == nil {
+		fmt.Println("Key not found in context")
+		return ctx
+	}
+	return context.WithValue(ctx, key, newValue)
+}
+
+// Delete removes a key-value pair by returning a new context (contexts are immutable).
+func DeleteContext(ctx context.Context) context.Context {
+	return context.Background() // Returns a new empty context
+}
diff --git a/pkg/ktx/context_test.go b/pkg/ktx/context_test.go
@@ -0,0 +1,70 @@
+package ktx
+
+import (
+	"context"
+	"testing"
+)
+
+// Define test keys.
+const (
+	TestKey1 ContextKey = "testKey1"
+	TestKey2 ContextKey = "testKey2"
+)
+
+// Test CreateContext and ReadContext.
+func TestCreateAndReadContext(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+
+	// Create context with values.
+	ctx = CreateContext(ctx, TestKey1, "Hello")
+	ctx = CreateContext(ctx, TestKey2, 42)
+
+	// Read and check values.
+	if value, found := ReadContext(ctx, TestKey1); !found || value != "Hello" {
+		t.Errorf("Expected 'Hello', got %v", value)
+	}
+
+	if value, found := ReadContext(ctx, TestKey2); !found || value != 42 {
+		t.Errorf("Expected 42, got %v", value)
+	}
+}
+
+// Test UpdateContext.
+func TestUpdateContext(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	ctx = CreateContext(ctx, TestKey1, "InitialValue")
+
+	// Update value.
+	ctx = UpdateContext(ctx, TestKey1, "UpdatedValue")
+
+	// Verify updated value.
+	if value, found := ReadContext(ctx, TestKey1); !found || value != "UpdatedValue" {
+		t.Errorf("Expected 'UpdatedValue', got %v", value)
+	}
+
+	// Try updating a non-existing key (should not modify).
+	ctx = UpdateContext(ctx, "nonExistingKey", "NewValue")
+	if _, found := ReadContext(ctx, "nonExistingKey"); found {
+		t.Errorf("Expected non-existing key to remain absent")
+	}
+}
+
+// Test DeleteContext.
+func TestDeleteContext(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	ctx = CreateContext(ctx, TestKey1, "ToBeDeleted")
+
+	// Delete context.
+	ctx = DeleteContext(ctx)
+
+	// Verify it's empty.
+	if value, found := ReadContext(ctx, TestKey1); found {
+		t.Errorf("Expected key to be deleted, but got %v", value)
+	}
+}
diff --git a/pkg/resolver/resolver.go b/pkg/resolver/resolver.go
@@ -11,6 +11,7 @@ import (
 	"github.com/apple/pkl-go/pkl"
 	"github.com/kdeps/kartographer/graph"
 	"github.com/kdeps/kdeps/pkg/environment"
+	"github.com/kdeps/kdeps/pkg/ktx"
 	"github.com/kdeps/kdeps/pkg/logging"
 	"github.com/kdeps/kdeps/pkg/utils"
 	pklRes "github.com/kdeps/schema/gen/resource"
@@ -48,7 +49,23 @@ type ResourceNodeEntry struct {
 	File     string `pkl:"file"`
 }
 
-func NewGraphResolver(fs afero.Fs, ctx context.Context, env *environment.Environment, agentDir, actionDir, graphID string, logger *logging.Logger) (*DependencyResolver, error) {
+func NewGraphResolver(fs afero.Fs, ctx context.Context, env *environment.Environment, logger *logging.Logger) (*DependencyResolver, error) {
+	var agentDir, graphID, actionDir string
+
+	contextKeys := map[*string]ktx.ContextKey{
+		&agentDir:  ktx.CtxKeyAgentDir,
+		&graphID:   ktx.CtxKeyGraphID,
+		&actionDir: ktx.CtxKeyActionDir,
+	}
+
+	for ptr, key := range contextKeys {
+		if value, found := ktx.ReadContext(ctx, key); found {
+			if strValue, ok := value.(string); ok {
+				*ptr = strValue
+			}
+		}
+	}
+
 	workflowDir := filepath.Join(agentDir, "/workflow/")
 	projectDir := filepath.Join(agentDir, "/project/")
 	pklWfFile := filepath.Join(workflowDir, "workflow.pkl")

diff --git a/pkg/resolver/resolver_test.go b/pkg/resolver/resolver_test.go
@@ -379,7 +379,7 @@ func iLoadTheWorkflowResources() error {
 	logger := logging.GetLogger()
 	ctx = context.Background()
 
-	dr, err := resolver.NewGraphResolver(testFs, ctx, environ, agentDir, "/tmp/action", "123", logger)
+	dr, err := resolver.NewGraphResolver(testFs, ctx, environ, logger)
 	if err != nil {
 		log.Fatal(err)
 	}