fix: complete tool calls on turn end (#216)
This commit is contained in:
@@ -263,7 +263,10 @@ func (m *mockRepository) GetActiveTurnBySessionID(ctx context.Context, sessionID
|
||||
return nil, nil
|
||||
}
|
||||
func (m *mockRepository) UpdateTurn(ctx context.Context, turn *models.Turn) error { return nil }
|
||||
func (m *mockRepository) CompleteTurn(ctx context.Context, id string) error { return nil }
|
||||
func (m *mockRepository) CompleteTurn(ctx context.Context, id string) error { return nil }
|
||||
func (m *mockRepository) CompleteRunningToolCallsForTurn(ctx context.Context, turnID string) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
func (m *mockRepository) ListTurnsBySession(ctx context.Context, sessionID string) ([]*models.Turn, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
@@ -150,6 +150,9 @@ func (m *mockRepository) UpdateTurn(ctx context.Context, turn *models.Turn) erro
|
||||
func (m *mockRepository) CompleteTurn(ctx context.Context, id string) error {
|
||||
return nil
|
||||
}
|
||||
func (m *mockRepository) CompleteRunningToolCallsForTurn(ctx context.Context, turnID string) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
func (m *mockRepository) ListTurnsBySession(ctx context.Context, sessionID string) ([]*models.Turn, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
@@ -60,6 +60,7 @@ type Repository interface {
|
||||
GetActiveTurnBySessionID(ctx context.Context, sessionID string) (*models.Turn, error)
|
||||
UpdateTurn(ctx context.Context, turn *models.Turn) error
|
||||
CompleteTurn(ctx context.Context, id string) error
|
||||
CompleteRunningToolCallsForTurn(ctx context.Context, turnID string) (int64, error)
|
||||
ListTurnsBySession(ctx context.Context, sessionID string) ([]*models.Turn, error)
|
||||
|
||||
// Task Session operations
|
||||
|
||||
@@ -267,6 +267,24 @@ func (r *Repository) GetMessageByPendingID(ctx context.Context, sessionID, pendi
|
||||
return message, nil
|
||||
}
|
||||
|
||||
// CompleteRunningToolCallsForTurn marks all "running" tool call messages for a turn as "complete".
|
||||
// This is a safety net to ensure no tool calls remain stuck in "running" state after a turn completes.
|
||||
func (r *Repository) CompleteRunningToolCallsForTurn(ctx context.Context, turnID string) (int64, error) {
|
||||
result, err := r.db.ExecContext(ctx, `
|
||||
UPDATE task_session_messages
|
||||
SET metadata = json_set(metadata, '$.status', 'complete')
|
||||
WHERE turn_id = ?
|
||||
AND json_extract(metadata, '$.status') = 'running'
|
||||
AND json_extract(metadata, '$.tool_call_id') IS NOT NULL
|
||||
`, turnID)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("failed to complete running tool calls for turn %s: %w", turnID, err)
|
||||
}
|
||||
|
||||
rows, _ := result.RowsAffected()
|
||||
return rows, nil
|
||||
}
|
||||
|
||||
// UpdateMessage updates an existing message
|
||||
func (r *Repository) UpdateMessage(ctx context.Context, message *models.Message) error {
|
||||
metadataJSON, err := json.Marshal(message.Metadata)
|
||||
|
||||
@@ -1059,3 +1059,94 @@ func TestSQLiteRepository_UpdateTaskSessionState(t *testing.T) {
|
||||
t.Error("expected error for updating nonexistent session status")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSQLiteRepository_CompleteRunningToolCallsForTurn(t *testing.T) {
|
||||
repo, cleanup := createTestSQLiteRepo(t)
|
||||
defer cleanup()
|
||||
ctx := context.Background()
|
||||
|
||||
// Setup
|
||||
board := &models.Board{ID: "board-1", Name: "Test Board"}
|
||||
_ = repo.CreateBoard(ctx, board)
|
||||
task := &models.Task{ID: "task-1", BoardID: "board-1", WorkflowStepID: "step-1", Title: "Test Task"}
|
||||
_ = repo.CreateTask(ctx, task)
|
||||
sessionID := setupSQLiteTestSession(ctx, repo, task.ID, "session-1")
|
||||
turnID := setupSQLiteTestTurn(ctx, repo, sessionID, task.ID, "turn-1")
|
||||
|
||||
// Create a tool call message with status "running"
|
||||
runningTool := &models.Message{
|
||||
ID: "msg-running-1", TaskSessionID: sessionID, TaskID: task.ID, TurnID: turnID,
|
||||
AuthorType: models.MessageAuthorAgent, Content: "Running tool",
|
||||
Type: models.MessageTypeToolCall,
|
||||
Metadata: map[string]interface{}{"tool_call_id": "tc-1", "status": "running"},
|
||||
}
|
||||
// Create a tool call message already "complete"
|
||||
completeTool := &models.Message{
|
||||
ID: "msg-complete-1", TaskSessionID: sessionID, TaskID: task.ID, TurnID: turnID,
|
||||
AuthorType: models.MessageAuthorAgent, Content: "Complete tool",
|
||||
Type: models.MessageTypeToolCall,
|
||||
Metadata: map[string]interface{}{"tool_call_id": "tc-2", "status": "complete"},
|
||||
}
|
||||
// Create a regular message (no tool_call_id) with status "running" — should NOT be affected
|
||||
regularMsg := &models.Message{
|
||||
ID: "msg-regular-1", TaskSessionID: sessionID, TaskID: task.ID, TurnID: turnID,
|
||||
AuthorType: models.MessageAuthorAgent, Content: "Regular message",
|
||||
Type: models.MessageTypeMessage,
|
||||
Metadata: map[string]interface{}{"status": "running"},
|
||||
}
|
||||
// Create a second running tool call
|
||||
runningTool2 := &models.Message{
|
||||
ID: "msg-running-2", TaskSessionID: sessionID, TaskID: task.ID, TurnID: turnID,
|
||||
AuthorType: models.MessageAuthorAgent, Content: "Running tool 2",
|
||||
Type: models.MessageTypeToolCall,
|
||||
Metadata: map[string]interface{}{"tool_call_id": "tc-3", "status": "running"},
|
||||
}
|
||||
|
||||
for _, msg := range []*models.Message{runningTool, completeTool, regularMsg, runningTool2} {
|
||||
if err := repo.CreateMessage(ctx, msg); err != nil {
|
||||
t.Fatalf("failed to create message %s: %v", msg.ID, err)
|
||||
}
|
||||
}
|
||||
|
||||
// Execute
|
||||
affected, err := repo.CompleteRunningToolCallsForTurn(ctx, turnID)
|
||||
if err != nil {
|
||||
t.Fatalf("CompleteRunningToolCallsForTurn failed: %v", err)
|
||||
}
|
||||
|
||||
// Should have updated exactly 2 running tool call messages
|
||||
if affected != 2 {
|
||||
t.Errorf("expected 2 affected rows, got %d", affected)
|
||||
}
|
||||
|
||||
// Verify running tool calls are now "complete"
|
||||
msg1, _ := repo.GetMessage(ctx, "msg-running-1")
|
||||
if msg1.Metadata["status"] != "complete" {
|
||||
t.Errorf("expected msg-running-1 status 'complete', got %v", msg1.Metadata["status"])
|
||||
}
|
||||
msg2, _ := repo.GetMessage(ctx, "msg-running-2")
|
||||
if msg2.Metadata["status"] != "complete" {
|
||||
t.Errorf("expected msg-running-2 status 'complete', got %v", msg2.Metadata["status"])
|
||||
}
|
||||
|
||||
// Verify already-complete tool call is unchanged
|
||||
msg3, _ := repo.GetMessage(ctx, "msg-complete-1")
|
||||
if msg3.Metadata["status"] != "complete" {
|
||||
t.Errorf("expected msg-complete-1 status 'complete', got %v", msg3.Metadata["status"])
|
||||
}
|
||||
|
||||
// Verify regular message (no tool_call_id) was NOT affected
|
||||
msg4, _ := repo.GetMessage(ctx, "msg-regular-1")
|
||||
if msg4.Metadata["status"] != "running" {
|
||||
t.Errorf("expected msg-regular-1 status 'running' (unchanged), got %v", msg4.Metadata["status"])
|
||||
}
|
||||
|
||||
// Running again should affect 0 rows
|
||||
affected2, err := repo.CompleteRunningToolCallsForTurn(ctx, turnID)
|
||||
if err != nil {
|
||||
t.Fatalf("second CompleteRunningToolCallsForTurn failed: %v", err)
|
||||
}
|
||||
if affected2 != 0 {
|
||||
t.Errorf("expected 0 affected rows on second call, got %d", affected2)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2416,6 +2416,15 @@ func (s *Service) CompleteTurn(ctx context.Context, turnID string) error {
|
||||
return err
|
||||
}
|
||||
|
||||
// Safety net: mark any tool calls still in "running" state as "complete"
|
||||
if affected, err := s.repo.CompleteRunningToolCallsForTurn(ctx, turnID); err != nil {
|
||||
s.logger.Warn("failed to complete running tool calls for turn", zap.String("turn_id", turnID), zap.Error(err))
|
||||
} else if affected > 0 {
|
||||
s.logger.Info("completed stale running tool calls on turn end",
|
||||
zap.String("turn_id", turnID),
|
||||
zap.Int64("affected", affected))
|
||||
}
|
||||
|
||||
// Fetch the completed turn to get the completed_at timestamp
|
||||
turn, err := s.repo.GetTurn(ctx, turnID)
|
||||
if err != nil {
|
||||
|
||||
Reference in New Issue
Block a user