feat: add context.Context to database client interface (#1489)

## Description
Fixes: #1488

- Add context.Context as the first parameter to all Client interface
methods in the database layer
- Update `gormClient`, `fakeClient`, and all HTTP handler/controller
call sites to pass context through
- Enables request cancellation and tracing to propagate from HTTP
handlers and the controller reconciler down to database operations

# Testing

- Existing unit tests updated to pass `context.Background()` — all pass
locally
- Fake client updated to match new interface signatures
- No behavioral changes; purely a mechanical refactor to follow Go
context idioms

Signed-off-by: Jeremy Alvis <jeremy.alvis@solo.io>

Author: iplay88keys

go/api/database/client.go

@@ -1,6 +1,7 @@
 package database
 
 import (
+	"context"
 	"time"
 
 	"github.com/kagent-dev/kagent/go/api/v1alpha2"
@@ -20,63 +21,63 @@ type LangGraphCheckpointTuple struct {
 
 type Client interface {
 	// Store methods
-	StoreFeedback(feedback *Feedback) error
-	StoreSession(session *Session) error
-	StoreAgent(agent *Agent) error
-	StoreTask(task *protocol.Task) error
-	StorePushNotification(config *protocol.TaskPushNotificationConfig) error
-	StoreToolServer(toolServer *ToolServer) (*ToolServer, error)
-	StoreEvents(messages ...*Event) error
+	StoreFeedback(ctx context.Context, feedback *Feedback) error
+	StoreSession(ctx context.Context, session *Session) error
+	StoreAgent(ctx context.Context, agent *Agent) error
+	StoreTask(ctx context.Context, task *protocol.Task) error
+	StorePushNotification(ctx context.Context, config *protocol.TaskPushNotificationConfig) error
+	StoreToolServer(ctx context.Context, toolServer *ToolServer) (*ToolServer, error)
+	StoreEvents(ctx context.Context, messages ...*Event) error
 
 	// Delete methods
-	DeleteSession(sessionID string, userID string) error
-	DeleteAgent(agentID string) error
-	DeleteToolServer(serverName string, groupKind string) error
-	DeleteTask(taskID string) error
-	DeletePushNotification(taskID string) error
-	DeleteToolsForServer(serverName string, groupKind string) error
+	DeleteSession(ctx context.Context, sessionID string, userID string) error
+	DeleteAgent(ctx context.Context, agentID string) error
+	DeleteToolServer(ctx context.Context, serverName string, groupKind string) error
+	DeleteTask(ctx context.Context, taskID string) error
+	DeletePushNotification(ctx context.Context, taskID string) error
+	DeleteToolsForServer(ctx context.Context, serverName string, groupKind string) error
 
 	// Get methods
-	GetSession(sessionID string, userID string) (*Session, error)
-	GetAgent(name string) (*Agent, error)
-	GetTask(id string) (*protocol.Task, error)
-	GetTool(name string) (*Tool, error)
-	GetToolServer(name string) (*ToolServer, error)
-	GetPushNotification(taskID string, configID string) (*protocol.TaskPushNotificationConfig, error)
+	GetSession(ctx context.Context, sessionID string, userID string) (*Session, error)
+	GetAgent(ctx context.Context, name string) (*Agent, error)
+	GetTask(ctx context.Context, id string) (*protocol.Task, error)
+	GetTool(ctx context.Context, name string) (*Tool, error)
+	GetToolServer(ctx context.Context, name string) (*ToolServer, error)
+	GetPushNotification(ctx context.Context, taskID string, configID string) (*protocol.TaskPushNotificationConfig, error)
 
 	// List methods
-	ListTools() ([]Tool, error)
-	ListFeedback(userID string) ([]Feedback, error)
-	ListTasksForSession(sessionID string) ([]*protocol.Task, error)
-	ListSessions(userID string) ([]Session, error)
-	ListSessionsForAgent(agentID string, userID string) ([]Session, error)
-	ListAgents() ([]Agent, error)
-	ListToolServers() ([]ToolServer, error)
-	ListToolsForServer(serverName string, groupKind string) ([]Tool, error)
-	ListEventsForSession(sessionID, userID string, options QueryOptions) ([]*Event, error)
-	ListPushNotifications(taskID string) ([]*protocol.TaskPushNotificationConfig, error)
+	ListTools(ctx context.Context) ([]Tool, error)
+	ListFeedback(ctx context.Context, userID string) ([]Feedback, error)
+	ListTasksForSession(ctx context.Context, sessionID string) ([]*protocol.Task, error)
+	ListSessions(ctx context.Context, userID string) ([]Session, error)
+	ListSessionsForAgent(ctx context.Context, agentID string, userID string) ([]Session, error)
+	ListAgents(ctx context.Context) ([]Agent, error)
+	ListToolServers(ctx context.Context) ([]ToolServer, error)
+	ListToolsForServer(ctx context.Context, serverName string, groupKind string) ([]Tool, error)
+	ListEventsForSession(ctx context.Context, sessionID, userID string, options QueryOptions) ([]*Event, error)
+	ListPushNotifications(ctx context.Context, taskID string) ([]*protocol.TaskPushNotificationConfig, error)
 
 	// Helper methods
-	RefreshToolsForServer(serverName string, groupKind string, tools ...*v1alpha2.MCPTool) error
+	RefreshToolsForServer(ctx context.Context, serverName string, groupKind string, tools ...*v1alpha2.MCPTool) error
 
 	// LangGraph Checkpoint methods
-	StoreCheckpoint(checkpoint *LangGraphCheckpoint) error
-	StoreCheckpointWrites(writes []*LangGraphCheckpointWrite) error
-	ListCheckpoints(userID, threadID, checkpointNS string, checkpointID *string, limit int) ([]*LangGraphCheckpointTuple, error)
-	DeleteCheckpoint(userID, threadID string) error
+	StoreCheckpoint(ctx context.Context, checkpoint *LangGraphCheckpoint) error
+	StoreCheckpointWrites(ctx context.Context, writes []*LangGraphCheckpointWrite) error
+	ListCheckpoints(ctx context.Context, userID, threadID, checkpointNS string, checkpointID *string, limit int) ([]*LangGraphCheckpointTuple, error)
+	DeleteCheckpoint(ctx context.Context, userID, threadID string) error
 
 	// CrewAI methods
-	StoreCrewAIMemory(memory *CrewAIAgentMemory) error
-	SearchCrewAIMemoryByTask(userID, threadID, taskDescription string, limit int) ([]*CrewAIAgentMemory, error)
-	ResetCrewAIMemory(userID, threadID string) error
-	StoreCrewAIFlowState(state *CrewAIFlowState) error
-	GetCrewAIFlowState(userID, threadID string) (*CrewAIFlowState, error)
+	StoreCrewAIMemory(ctx context.Context, memory *CrewAIAgentMemory) error
+	SearchCrewAIMemoryByTask(ctx context.Context, userID, threadID, taskDescription string, limit int) ([]*CrewAIAgentMemory, error)
+	ResetCrewAIMemory(ctx context.Context, userID, threadID string) error
+	StoreCrewAIFlowState(ctx context.Context, state *CrewAIFlowState) error
+	GetCrewAIFlowState(ctx context.Context, userID, threadID string) (*CrewAIFlowState, error)
 
 	// Agent memory (vector search) methods
-	StoreAgentMemory(memory *Memory) error
-	StoreAgentMemories(memories []*Memory) error
-	SearchAgentMemory(agentName, userID string, embedding pgvector.Vector, limit int) ([]AgentMemorySearchResult, error)
-	ListAgentMemories(agentName, userID string) ([]Memory, error)
-	DeleteAgentMemory(agentName, userID string) error
-	PruneExpiredMemories() error
+	StoreAgentMemory(ctx context.Context, memory *Memory) error
+	StoreAgentMemories(ctx context.Context, memories []*Memory) error
+	SearchAgentMemory(ctx context.Context, agentName, userID string, embedding pgvector.Vector, limit int) ([]AgentMemorySearchResult, error)
+	ListAgentMemories(ctx context.Context, agentName, userID string) ([]Memory, error)
+	DeleteAgentMemory(ctx context.Context, agentName, userID string) error
+	PruneExpiredMemories(ctx context.Context) error
 }

go/core/internal/controller/reconciler/reconciler.go

@@ -85,7 +85,7 @@ func (a *kagentReconciler) ReconcileKagentAgent(ctx context.Context, req ctrl.Re
 	agent := &v1alpha2.Agent{}
 	if err := a.kube.Get(ctx, req.NamespacedName, agent); err != nil {
 		if apierrors.IsNotFound(err) {
-			return a.handleAgentDeletion(req)
+			return a.handleAgentDeletion(ctx, req)
 		}
 
 		return fmt.Errorf("failed to get agent %s: %w", req.NamespacedName, err)
@@ -99,9 +99,9 @@ func (a *kagentReconciler) ReconcileKagentAgent(ctx context.Context, req ctrl.Re
 	return a.reconcileAgentStatus(ctx, agent, err)
 }
 
-func (a *kagentReconciler) handleAgentDeletion(req ctrl.Request) error {
+func (a *kagentReconciler) handleAgentDeletion(ctx context.Context, req ctrl.Request) error {
 	id := utils.ConvertToPythonIdentifier(req.String())
-	if err := a.dbClient.DeleteAgent(id); err != nil {
+	if err := a.dbClient.DeleteAgent(ctx, id); err != nil {
 		return fmt.Errorf("failed to delete agent %s: %w",
 			req.String(), err)
 	}
@@ -205,11 +205,11 @@ func (a *kagentReconciler) ReconcileKagentMCPService(ctx context.Context, req ct
 				Name:      req.String(),
 				GroupKind: schema.GroupKind{Group: "", Kind: "Service"}.String(),
 			}
-			if err := a.dbClient.DeleteToolServer(dbService.Name, dbService.GroupKind); err != nil {
+			if err := a.dbClient.DeleteToolServer(ctx, dbService.Name, dbService.GroupKind); err != nil {
 				reconcileLog.Error(err, "failed to delete tool server for mcp service", "service", req.String())
 			}
 			reconcileLog.Info("mcp service was deleted", "service", req.String())
-			if err := a.dbClient.DeleteToolsForServer(dbService.Name, dbService.GroupKind); err != nil {
+			if err := a.dbClient.DeleteToolsForServer(ctx, dbService.Name, dbService.GroupKind); err != nil {
 				reconcileLog.Error(err, "failed to delete tools for mcp service", "service", req.String())
 			}
 			return nil
@@ -378,11 +378,11 @@ func (a *kagentReconciler) ReconcileKagentMCPServer(ctx context.Context, req ctr
 				Name:      req.String(),
 				GroupKind: schema.GroupKind{Group: "kagent.dev", Kind: "MCPServer"}.String(),
 			}
-			if err := a.dbClient.DeleteToolServer(dbServer.Name, dbServer.GroupKind); err != nil {
+			if err := a.dbClient.DeleteToolServer(ctx, dbServer.Name, dbServer.GroupKind); err != nil {
 				reconcileLog.Error(err, "failed to delete tool server for mcp server", "mcpServer", req.String())
 			}
 			reconcileLog.Info("mcp server was deleted", "mcpServer", req.String())
-			if err := a.dbClient.DeleteToolsForServer(dbServer.Name, dbServer.GroupKind); err != nil {
+			if err := a.dbClient.DeleteToolsForServer(ctx, dbServer.Name, dbServer.GroupKind); err != nil {
 				reconcileLog.Error(err, "failed to delete tools for mcp server", "mcpServer", req.String())
 			}
 			return nil
@@ -428,11 +428,11 @@ func (a *kagentReconciler) ReconcileKagentRemoteMCPServer(ctx context.Context, r
 				GroupKind: schema.GroupKind{Group: "kagent.dev", Kind: "RemoteMCPServer"}.String(),
 			}
 
-			if err := a.dbClient.DeleteToolServer(dbServer.Name, dbServer.GroupKind); err != nil {
+			if err := a.dbClient.DeleteToolServer(ctx, dbServer.Name, dbServer.GroupKind); err != nil {
 				l.Error(err, "failed to delete tool server for remote mcp server")
 			}
 
-			if err := a.dbClient.DeleteToolsForServer(dbServer.Name, dbServer.GroupKind); err != nil {
+			if err := a.dbClient.DeleteToolsForServer(ctx, dbServer.Name, dbServer.GroupKind); err != nil {
 				l.Error(err, "failed to delete tools for remote mcp server")
 			}
 
@@ -835,15 +835,15 @@ func (a *kagentReconciler) upsertAgent(ctx context.Context, agent *v1alpha2.Agen
 		Config: agentOutputs.Config,
 	}
 
-	if err := a.dbClient.StoreAgent(dbAgent); err != nil {
+	if err := a.dbClient.StoreAgent(ctx, dbAgent); err != nil {
 		return fmt.Errorf("failed to store agent %s: %w", id, err)
 	}
 
 	return nil
 }
 
 func (a *kagentReconciler) upsertToolServerForRemoteMCPServer(ctx context.Context, toolServer *database.ToolServer, remoteMcpServer *v1alpha2.RemoteMCPServer) ([]*v1alpha2.MCPTool, error) {
-	if _, err := a.dbClient.StoreToolServer(toolServer); err != nil {
+	if _, err := a.dbClient.StoreToolServer(ctx, toolServer); err != nil {
 		return nil, fmt.Errorf("failed to store toolServer %s: %w", toolServer.Name, err)
 	}
 
@@ -858,7 +858,7 @@ func (a *kagentReconciler) upsertToolServerForRemoteMCPServer(ctx context.Contex
 	}
 
 	// Refresh tools in database - uses transaction for atomicity
-	if err := a.dbClient.RefreshToolsForServer(toolServer.Name, toolServer.GroupKind, tools...); err != nil {
+	if err := a.dbClient.RefreshToolsForServer(ctx, toolServer.Name, toolServer.GroupKind, tools...); err != nil {
 		return nil, fmt.Errorf("failed to refresh tools for toolServer %s: %w", toolServer.Name, err)
 	}
 
@@ -956,7 +956,7 @@ func (a *kagentReconciler) listTools(ctx context.Context, tsp mcp.Transport, too
 
 func (a *kagentReconciler) getDiscoveredMCPTools(ctx context.Context, serverRef string) ([]*v1alpha2.MCPTool, error) {
 	// This function is currently only used for RemoteMCPServer
-	allTools, err := a.dbClient.ListToolsForServer(serverRef, schema.GroupKind{Group: "kagent.dev", Kind: "RemoteMCPServer"}.String())
+	allTools, err := a.dbClient.ListToolsForServer(ctx, serverRef, schema.GroupKind{Group: "kagent.dev", Kind: "RemoteMCPServer"}.String())
 	if err != nil {
 		return nil, err
 	}