refactor: rename CorrelateApprovalByToolID to LinkConversationEventToApprovalUsingToolID

Rename the method for better clarity and move correlation logic to execute
immediately after approval creation for all approval types, not just auto-approved
This commit is contained in:
Sundeep Malladi
2025-08-14 15:30:34 -05:00
parent 9a55e0c2f4
commit 009eaf417f
4 changed files with 14 additions and 14 deletions

View File

@@ -232,7 +232,7 @@ func (m *manager) correlateApproval(ctx context.Context, approval *store.Approva
}
// Correlate by tool ID
if err := m.store.CorrelateApprovalByToolID(ctx, approval.SessionID, toolCall.ToolID, approval.ID); err != nil {
if err := m.store.LinkConversationEventToApprovalUsingToolID(ctx, approval.SessionID, toolCall.ToolID, approval.ID); err != nil {
return fmt.Errorf("failed to correlate approval: %w", err)
}
@@ -345,6 +345,14 @@ func (m *manager) CreateApprovalWithToolUseID(ctx context.Context, sessionID, to
// Publish event for real-time updates
m.publishNewApprovalEvent(approval)
if err := m.store.LinkConversationEventToApprovalUsingToolID(ctx, sessionID, toolUseID, approval.ID); err != nil {
// Log but don't fail
// TODO(1): Don't ship if above LinkConversationEventToApprovalUsingToolID does not retry
// it's possible, albeit unlikely, that the raw_event has not made it to
// conversation_events yet
return nil, fmt.Errorf("failed to correlate approval: %w", err)
}
// Handle status-specific post-creation tasks
switch status {
case store.ApprovalStatusLocalPending:
@@ -355,15 +363,7 @@ func (m *manager) CreateApprovalWithToolUseID(ctx context.Context, sessionID, to
"session_id", session.ID)
}
case store.ApprovalStatusLocalApproved:
// For auto-approved, update correlation status immediately if we can correlate
// Try to correlate with tool call by tool_use_id
if err := m.store.CorrelateApprovalByToolID(ctx, sessionID, toolUseID, approval.ID); err != nil {
// Log but don't fail
slog.Warn("failed to correlate auto-approved approval",
"error", err,
"approval_id", approval.ID,
"tool_use_id", toolUseID)
}
// For auto-approved, update correlation status immediately
// Update approval status
if err := m.store.UpdateApprovalStatus(ctx, approval.ID, store.ApprovalStatusApproved); err != nil {
slog.Warn("failed to update approval status in conversation events",

View File

@@ -248,7 +248,7 @@ func TestManager_CorrelateApproval(t *testing.T) {
mockStore.EXPECT().GetUncorrelatedPendingToolCall(ctx, sessionID, toolName).Return(pendingToolCall, nil)
// Mock correlating by tool ID
mockStore.EXPECT().CorrelateApprovalByToolID(ctx, sessionID, "tool-123", gomock.Any()).Return(nil)
mockStore.EXPECT().LinkConversationEventToApprovalUsingToolID(ctx, sessionID, "tool-123", gomock.Any()).Return(nil)
// Mock event publishing
mockEventBus.EXPECT().Publish(gomock.Any())

View File

@@ -1825,8 +1825,8 @@ func (s *SQLiteStore) CorrelateApproval(ctx context.Context, sessionID string, t
return nil
}
// CorrelateApprovalByToolID correlates an approval with a specific tool call by tool_id
func (s *SQLiteStore) CorrelateApprovalByToolID(ctx context.Context, sessionID string, toolID string, approvalID string) error {
// LinkConversationEventToApprovalUsingToolID correlates an approval with a specific tool call by tool_id
func (s *SQLiteStore) LinkConversationEventToApprovalUsingToolID(ctx context.Context, sessionID string, toolID string, approvalID string) error {
// Update the tool call directly by tool_id
updateQuery := `
UPDATE conversation_events

View File

@@ -31,7 +31,7 @@ type ConversationStore interface {
GetToolCallByID(ctx context.Context, toolID string) (*ConversationEvent, error)
MarkToolCallCompleted(ctx context.Context, toolID string, sessionID string) error
CorrelateApproval(ctx context.Context, sessionID string, toolName string, approvalID string) error
CorrelateApprovalByToolID(ctx context.Context, sessionID string, toolID string, approvalID string) error
LinkConversationEventToApprovalUsingToolID(ctx context.Context, sessionID string, toolID string, approvalID string) error
UpdateApprovalStatus(ctx context.Context, approvalID string, status string) error
// MCP server operations