740 lines
28 KiB
Go
740 lines
28 KiB
Go
package sqlite
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
|
|
"github.com/kandev/kandev/internal/agentctl/tracing"
|
|
"github.com/kandev/kandev/internal/db/dialect"
|
|
"github.com/kandev/kandev/internal/task/models"
|
|
)
|
|
|
|
// Message operations
|
|
|
|
// CreateMessage creates a new message
|
|
func (r *Repository) CreateMessage(ctx context.Context, message *models.Message) error {
|
|
if message.ID == "" {
|
|
message.ID = uuid.New().String()
|
|
}
|
|
if message.CreatedAt.IsZero() {
|
|
message.CreatedAt = time.Now().UTC()
|
|
}
|
|
if message.UpdatedAt.IsZero() {
|
|
message.UpdatedAt = message.CreatedAt
|
|
}
|
|
if message.AuthorType == "" {
|
|
message.AuthorType = models.MessageAuthorUser
|
|
}
|
|
|
|
requestsInput := 0
|
|
if message.RequestsInput {
|
|
requestsInput = 1
|
|
}
|
|
|
|
messageType := string(message.Type)
|
|
if messageType == "" {
|
|
messageType = string(models.MessageTypeMessage)
|
|
}
|
|
|
|
metadataJSON := "{}"
|
|
if message.Metadata != nil {
|
|
metadataBytes, err := json.Marshal(message.Metadata)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to serialize message metadata: %w", err)
|
|
}
|
|
metadataJSON = string(metadataBytes)
|
|
}
|
|
|
|
return r.insertMessageWithSessionLock(ctx, message, requestsInput, messageType, metadataJSON)
|
|
}
|
|
|
|
func (r *Repository) insertMessageRow(
|
|
ctx context.Context,
|
|
execer taskSessionExecutor,
|
|
message *models.Message,
|
|
requestsInput int,
|
|
messageType, metadataJSON string,
|
|
) error {
|
|
_, err := execer.ExecContext(ctx, r.db.Rebind(`
|
|
INSERT INTO task_session_messages (id, task_session_id, task_id, turn_id, author_type, author_id, content, requests_input, type, metadata, created_at, updated_at)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
`), message.ID, message.TaskSessionID, message.TaskID, message.TurnID, message.AuthorType, message.AuthorID, message.Content, requestsInput, messageType, metadataJSON, message.CreatedAt, message.UpdatedAt)
|
|
return err
|
|
}
|
|
|
|
// insertMessageWithSessionLock serializes message insertion with rollback of
|
|
// the message's turn on PostgreSQL. SQLite's writer pool is configured with a
|
|
// single connection in db.OpenSQLite, which provides equivalent serialization.
|
|
// If SQLite ever allows concurrent writers, add the equivalent session lock.
|
|
func (r *Repository) insertMessageWithSessionLock(
|
|
ctx context.Context,
|
|
message *models.Message,
|
|
requestsInput int,
|
|
messageType, metadataJSON string,
|
|
) error {
|
|
if !dialect.IsPostgres(r.db.DriverName()) {
|
|
return r.insertMessageRow(ctx, r.db, message, requestsInput, messageType, metadataJSON)
|
|
}
|
|
tx, err := r.db.BeginTxx(ctx, nil)
|
|
if err != nil {
|
|
return fmt.Errorf("begin message creation: %w", err)
|
|
}
|
|
defer func() { _ = tx.Rollback() }()
|
|
if err := lockSessionTurnWrites(ctx, tx, r.db.DriverName(), message.TaskSessionID); err != nil {
|
|
return err
|
|
}
|
|
if err := r.insertMessageRow(ctx, tx, message, requestsInput, messageType, metadataJSON); err != nil {
|
|
return err
|
|
}
|
|
if err := tx.Commit(); err != nil {
|
|
return fmt.Errorf("commit message creation: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// GetMessage retrieves a message by ID
|
|
func (r *Repository) GetMessage(ctx context.Context, id string) (*models.Message, error) {
|
|
message := &models.Message{}
|
|
var requestsInput int
|
|
var messageType string
|
|
var metadataJSON string
|
|
err := r.ro.QueryRowContext(ctx, r.ro.Rebind(`
|
|
SELECT id, task_session_id, task_id, turn_id, author_type, author_id, content, requests_input, type, metadata, created_at, updated_at
|
|
FROM task_session_messages WHERE id = ?
|
|
`), id).Scan(&message.ID, &message.TaskSessionID, &message.TaskID, &message.TurnID, &message.AuthorType, &message.AuthorID, &message.Content, &requestsInput, &messageType, &metadataJSON, &message.CreatedAt, &message.UpdatedAt)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
message.RequestsInput = requestsInput == 1
|
|
message.Type = models.MessageType(messageType)
|
|
|
|
if metadataJSON != "" && metadataJSON != "{}" {
|
|
if err := json.Unmarshal([]byte(metadataJSON), &message.Metadata); err != nil {
|
|
return nil, fmt.Errorf("failed to deserialize message metadata: %w", err)
|
|
}
|
|
}
|
|
|
|
return message, nil
|
|
}
|
|
|
|
// ListMessages returns all messages for a session ordered by creation time.
|
|
func (r *Repository) ListMessages(ctx context.Context, sessionID string) ([]*models.Message, error) {
|
|
rows, err := r.ro.QueryContext(ctx, r.ro.Rebind(`
|
|
SELECT id, task_session_id, task_id, turn_id, author_type, author_id, content, requests_input, type, metadata, created_at, updated_at
|
|
FROM task_session_messages WHERE task_session_id = ? ORDER BY created_at ASC
|
|
`), sessionID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer func() { _ = rows.Close() }()
|
|
msgs, _, err := scanMessageRows(rows, 0)
|
|
return msgs, err
|
|
}
|
|
|
|
// ListMessagesByTurnID returns all messages for a single turn ordered by
|
|
// creation time. Backed by idx_messages_turn_id, so it reads only the turn's
|
|
// own rows rather than the whole session's history.
|
|
func (r *Repository) ListMessagesByTurnID(ctx context.Context, turnID string) ([]*models.Message, error) {
|
|
rows, err := r.ro.QueryContext(ctx, r.ro.Rebind(`
|
|
SELECT id, task_session_id, task_id, turn_id, author_type, author_id, content, requests_input, type, metadata, created_at, updated_at
|
|
FROM task_session_messages WHERE turn_id = ? ORDER BY created_at ASC
|
|
`), turnID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer func() { _ = rows.Close() }()
|
|
msgs, _, err := scanMessageRows(rows, 0)
|
|
return msgs, err
|
|
}
|
|
|
|
// ListMessagesPaginated returns messages for a session ordered by creation time with pagination.
|
|
func (r *Repository) ListMessagesPaginated(ctx context.Context, sessionID string, opts models.ListMessagesOptions) ([]*models.Message, bool, error) {
|
|
ctx, span := tracing.Tracer("kandev-db").Start(ctx, "db.ListMessagesPaginated")
|
|
defer span.End()
|
|
|
|
limit := opts.Limit
|
|
if limit < 0 {
|
|
limit = 0
|
|
}
|
|
sortDir := "ASC"
|
|
if strings.EqualFold(opts.Sort, "desc") {
|
|
sortDir = "DESC"
|
|
}
|
|
cursor, err := r.resolveMessageCursor(ctx, sessionID, opts)
|
|
if err != nil {
|
|
return nil, false, err
|
|
}
|
|
query, args := buildListMessagesQuery(sessionID, opts, cursor, sortDir, limit)
|
|
rows, err := r.ro.QueryContext(ctx, r.ro.Rebind(query), args...)
|
|
if err != nil {
|
|
return nil, false, err
|
|
}
|
|
defer func() { _ = rows.Close() }()
|
|
return scanMessageRows(rows, limit)
|
|
}
|
|
|
|
// ListMessagesForPlugin returns messages matching the plugin Host data API
|
|
// filter (ADR 0047): session/task ids, a created_at time range, and message
|
|
// types, ordered oldest-first by (created_at, id) with SQL Limit/Offset. It
|
|
// backs internal/plugins' capability-gated Messages().List reader — always via
|
|
// the service layer, never called directly by a plugin.
|
|
func (r *Repository) ListMessagesForPlugin(ctx context.Context, filter models.PluginMessageFilter) ([]*models.Message, error) {
|
|
ctx, span := tracing.Tracer("kandev-db").Start(ctx, "db.ListMessagesForPlugin")
|
|
defer span.End()
|
|
|
|
query, args := buildPluginMessageQuery(filter)
|
|
rows, err := r.ro.QueryContext(ctx, r.ro.Rebind(query), args...)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer func() { _ = rows.Close() }()
|
|
msgs, _, err := scanMessageRows(rows, 0)
|
|
return msgs, err
|
|
}
|
|
|
|
// buildPluginMessageQuery assembles the WHERE/ORDER/LIMIT/OFFSET clauses for
|
|
// ListMessagesForPlugin. session_ids and task_ids each become an IN (...)
|
|
// clause (ANDed together when both are set); Since is an inclusive and Until
|
|
// an exclusive created_at bound; types filters by message type. A zero/negative
|
|
// Limit yields no LIMIT clause (returns all matching rows from the offset).
|
|
func buildPluginMessageQuery(filter models.PluginMessageFilter) (string, []interface{}) {
|
|
query := `
|
|
SELECT id, task_session_id, task_id, turn_id, author_type, author_id, content, requests_input, type, metadata, created_at, updated_at
|
|
FROM task_session_messages`
|
|
var conditions []string
|
|
var args []interface{}
|
|
|
|
if ph, inArgs := buildInPlaceholders(filter.SessionIDs); ph != "" {
|
|
conditions = append(conditions, "task_session_id IN ("+ph+")")
|
|
args = append(args, inArgs...)
|
|
}
|
|
if ph, inArgs := buildInPlaceholders(filter.TaskIDs); ph != "" {
|
|
conditions = append(conditions, "task_id IN ("+ph+")")
|
|
args = append(args, inArgs...)
|
|
}
|
|
if ph, inArgs := buildInPlaceholders(filter.Types); ph != "" {
|
|
conditions = append(conditions, "type IN ("+ph+")")
|
|
args = append(args, inArgs...)
|
|
}
|
|
if filter.Since != nil {
|
|
conditions = append(conditions, "created_at >= ?")
|
|
args = append(args, *filter.Since)
|
|
}
|
|
if filter.Until != nil {
|
|
conditions = append(conditions, "created_at < ?")
|
|
args = append(args, *filter.Until)
|
|
}
|
|
if len(conditions) > 0 {
|
|
query += " WHERE " + strings.Join(conditions, " AND ")
|
|
}
|
|
query += " ORDER BY created_at ASC, id ASC"
|
|
// OFFSET is only emitted alongside a LIMIT: standalone OFFSET is a syntax
|
|
// error in SQLite, and the only caller (the plugin message reader) always
|
|
// requests a positive Limit (page-limit+1), so Offset never travels alone.
|
|
if filter.Limit > 0 {
|
|
query += sqlLimitClause
|
|
args = append(args, filter.Limit)
|
|
if filter.Offset > 0 {
|
|
query += " OFFSET ?"
|
|
args = append(args, filter.Offset)
|
|
}
|
|
}
|
|
return query, args
|
|
}
|
|
|
|
func (r *Repository) resolveMessageCursor(ctx context.Context, sessionID string, opts models.ListMessagesOptions) (*models.Message, error) {
|
|
if opts.Before != "" {
|
|
cursor, err := r.GetMessage(ctx, opts.Before)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if cursor.TaskSessionID != sessionID {
|
|
return nil, fmt.Errorf("message cursor not found: %s", opts.Before)
|
|
}
|
|
return cursor, nil
|
|
}
|
|
if opts.After != "" {
|
|
cursor, err := r.GetMessage(ctx, opts.After)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if cursor.TaskSessionID != sessionID {
|
|
return nil, fmt.Errorf("message cursor not found: %s", opts.After)
|
|
}
|
|
return cursor, nil
|
|
}
|
|
return nil, nil
|
|
}
|
|
|
|
func buildListMessagesQuery(sessionID string, opts models.ListMessagesOptions, cursor *models.Message, sortDir string, limit int) (string, []interface{}) {
|
|
query := `
|
|
SELECT id, task_session_id, task_id, turn_id, author_type, author_id, content, requests_input, type, metadata, created_at, updated_at
|
|
FROM task_session_messages WHERE task_session_id = ?`
|
|
args := []interface{}{sessionID}
|
|
if cursor != nil {
|
|
if opts.Before != "" {
|
|
query += " AND (created_at < ? OR (created_at = ? AND id < ?))"
|
|
} else if opts.After != "" {
|
|
query += " AND (created_at > ? OR (created_at = ? AND id > ?))"
|
|
}
|
|
args = append(args, cursor.CreatedAt, cursor.CreatedAt, cursor.ID)
|
|
}
|
|
query += fmt.Sprintf(" ORDER BY created_at %s, id %s", sortDir, sortDir)
|
|
if limit > 0 {
|
|
query += sqlLimitClause
|
|
args = append(args, limit+1)
|
|
}
|
|
return query, args
|
|
}
|
|
|
|
func scanMessageRows(rows interface {
|
|
Next() bool
|
|
Scan(dest ...any) error
|
|
Err() error
|
|
}, limit int) ([]*models.Message, bool, error) {
|
|
var result []*models.Message
|
|
for rows.Next() {
|
|
message := &models.Message{}
|
|
var requestsInput int
|
|
var messageType string
|
|
var metadataJSON string
|
|
if err := rows.Scan(&message.ID, &message.TaskSessionID, &message.TaskID, &message.TurnID, &message.AuthorType, &message.AuthorID, &message.Content, &requestsInput, &messageType, &metadataJSON, &message.CreatedAt, &message.UpdatedAt); err != nil {
|
|
return nil, false, err
|
|
}
|
|
message.RequestsInput = requestsInput == 1
|
|
message.Type = models.MessageType(messageType)
|
|
if metadataJSON != "" && metadataJSON != "{}" {
|
|
if err := json.Unmarshal([]byte(metadataJSON), &message.Metadata); err != nil {
|
|
return nil, false, fmt.Errorf("failed to deserialize message metadata: %w", err)
|
|
}
|
|
}
|
|
result = append(result, message)
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return nil, false, err
|
|
}
|
|
if limit > 0 && len(result) > limit {
|
|
return result[:limit], true, nil
|
|
}
|
|
return result, false, nil
|
|
}
|
|
|
|
// SearchMessages returns messages in a session whose content matches the query
|
|
// (case-insensitive substring). Newest-first ordering. Limit is capped.
|
|
func (r *Repository) SearchMessages(ctx context.Context, sessionID string, opts models.SearchMessagesOptions) ([]*models.Message, error) {
|
|
query := strings.TrimSpace(opts.Query)
|
|
if query == "" {
|
|
return nil, nil
|
|
}
|
|
limit := opts.Limit
|
|
if limit <= 0 {
|
|
limit = 100
|
|
}
|
|
if limit > 500 {
|
|
limit = 500
|
|
}
|
|
like := dialect.Like(r.ro.DriverName())
|
|
// Escape LIKE metacharacters so a user query of "%" or "_" matches
|
|
// literally rather than as SQL wildcards.
|
|
escaper := strings.NewReplacer(`\`, `\\`, `%`, `\%`, `_`, `\_`)
|
|
pattern := "%" + escaper.Replace(query) + "%"
|
|
sql := fmt.Sprintf(`
|
|
SELECT id, task_session_id, task_id, turn_id, author_type, author_id, content, requests_input, type, metadata, created_at, updated_at
|
|
FROM task_session_messages
|
|
WHERE task_session_id = ? AND content %s ? ESCAPE '\'
|
|
ORDER BY created_at DESC, id DESC
|
|
LIMIT ?
|
|
`, like)
|
|
rows, err := r.ro.QueryContext(ctx, r.ro.Rebind(sql), sessionID, pattern, limit)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer func() { _ = rows.Close() }()
|
|
result, _, err := scanMessageRows(rows, 0)
|
|
return result, err
|
|
}
|
|
|
|
// DeleteMessage deletes a message by ID
|
|
func (r *Repository) DeleteMessage(ctx context.Context, id string) error {
|
|
result, err := r.db.ExecContext(ctx, r.db.Rebind(`DELETE FROM task_session_messages WHERE id = ?`), id)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
rows, _ := result.RowsAffected()
|
|
if rows == 0 {
|
|
return fmt.Errorf("message not found: %s", id)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (r *Repository) getMessageByMetadataField(ctx context.Context, sessionID, fieldName, fieldValue, orderSuffix string) (*models.Message, error) {
|
|
message := &models.Message{}
|
|
var requestsInput int
|
|
var messageType string
|
|
var metadataJSON string
|
|
drv := r.ro.DriverName()
|
|
query := fmt.Sprintf(`
|
|
SELECT id, task_session_id, task_id, turn_id, author_type, author_id, content, requests_input, type, metadata, created_at, updated_at
|
|
FROM task_session_messages WHERE task_session_id = ? AND %s = ?
|
|
%s
|
|
`, dialect.JSONExtract(drv, "metadata", fieldName), orderSuffix)
|
|
err := r.ro.QueryRowContext(ctx, r.ro.Rebind(query), sessionID, fieldValue).Scan(&message.ID, &message.TaskSessionID, &message.TaskID, &message.TurnID, &message.AuthorType, &message.AuthorID,
|
|
&message.Content, &requestsInput, &messageType, &metadataJSON, &message.CreatedAt, &message.UpdatedAt)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
message.RequestsInput = requestsInput == 1
|
|
message.Type = models.MessageType(messageType)
|
|
if metadataJSON != "" {
|
|
_ = json.Unmarshal([]byte(metadataJSON), &message.Metadata)
|
|
}
|
|
return message, nil
|
|
}
|
|
|
|
// GetMessageByToolCallID retrieves a tool message by session ID and tool_call_id in metadata.
|
|
// Searches all message types that have a tool_call_id in metadata (tool_call, tool_read, tool_edit,
|
|
// tool_execute, tool_search, todo, etc.).
|
|
//
|
|
// Permission_request messages also store tool_call_id but represent the user-approval card,
|
|
// not the tool call itself. Excluded so a tool_update never lands on a permission_request
|
|
// row — that would overwrite metadata.status (the user's approve/reject) and retype it to
|
|
// tool_execute, making the prompt buttons reappear after the turn ends.
|
|
func (r *Repository) GetMessageByToolCallID(ctx context.Context, sessionID, toolCallID string) (*models.Message, error) {
|
|
return r.getMessageByMetadataField(ctx, sessionID, "tool_call_id", toolCallID,
|
|
"AND type != 'permission_request' ORDER BY created_at ASC LIMIT 1")
|
|
}
|
|
|
|
// GetMessageByPendingID retrieves a message by session ID and pending_id in metadata
|
|
func (r *Repository) GetMessageByPendingID(ctx context.Context, sessionID, pendingID string) (*models.Message, error) {
|
|
return r.getMessageByMetadataField(ctx, sessionID, "pending_id", pendingID, "")
|
|
}
|
|
|
|
// FindMessageByPendingID finds the most-recent message by pending_id alone.
|
|
// This is useful when we only have the pending ID (e.g., from expired clarification responses).
|
|
// For multi-question clarification requests, prefer FindMessagesByPendingID to retrieve
|
|
// every related message in one shot.
|
|
func (r *Repository) FindMessageByPendingID(ctx context.Context, pendingID string) (*models.Message, error) {
|
|
message := &models.Message{}
|
|
var requestsInput int
|
|
var messageType string
|
|
var metadataJSON string
|
|
drv := r.ro.DriverName()
|
|
query := fmt.Sprintf(`
|
|
SELECT id, task_session_id, task_id, turn_id, author_type, author_id, content, requests_input, type, metadata, created_at, updated_at
|
|
FROM task_session_messages WHERE %s = ?
|
|
ORDER BY created_at DESC LIMIT 1
|
|
`, dialect.JSONExtract(drv, "metadata", "pending_id"))
|
|
err := r.ro.QueryRowContext(ctx, r.ro.Rebind(query), pendingID).Scan(&message.ID, &message.TaskSessionID, &message.TaskID, &message.TurnID, &message.AuthorType, &message.AuthorID,
|
|
&message.Content, &requestsInput, &messageType, &metadataJSON, &message.CreatedAt, &message.UpdatedAt)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
message.RequestsInput = requestsInput == 1
|
|
message.Type = models.MessageType(messageType)
|
|
if metadataJSON != "" {
|
|
_ = json.Unmarshal([]byte(metadataJSON), &message.Metadata)
|
|
}
|
|
return message, nil
|
|
}
|
|
|
|
// FindMessagesByPendingID returns every message that carries the given pending_id
|
|
// in its metadata, ordered by creation time (oldest first). Multi-question
|
|
// clarification requests emit one message per question, all sharing the same
|
|
// pending_id; this lookup lets the canceller / status-update path touch all of
|
|
// them without N round-trips.
|
|
func (r *Repository) FindMessagesByPendingID(ctx context.Context, pendingID string) ([]*models.Message, error) {
|
|
drv := r.ro.DriverName()
|
|
query := fmt.Sprintf(`
|
|
SELECT id, task_session_id, task_id, turn_id, author_type, author_id, content, requests_input, type, metadata, created_at, updated_at
|
|
FROM task_session_messages WHERE %s = ?
|
|
ORDER BY created_at ASC
|
|
`, dialect.JSONExtract(drv, "metadata", "pending_id"))
|
|
rows, err := r.ro.QueryContext(ctx, r.ro.Rebind(query), pendingID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer func() { _ = rows.Close() }()
|
|
result, _, err := scanMessageRows(rows, 0)
|
|
return result, err
|
|
}
|
|
|
|
// FindActiveClarificationMessagesBySessionID returns pending clarification rows
|
|
// owned by the session's newest durable turn. Older pending rows remain history.
|
|
// The canceller still drains live legacy waiters before this lookup. Persisted
|
|
// rows without their schema-required turn are malformed and need explicit data
|
|
// cleanup; they never become current input authority.
|
|
func (r *Repository) FindActiveClarificationMessagesBySessionID(ctx context.Context, sessionID string) ([]*models.Message, error) {
|
|
drv := r.ro.DriverName()
|
|
query := fmt.Sprintf(`
|
|
WITH current_turn AS (
|
|
SELECT turn_row.id
|
|
FROM task_session_turns turn_row
|
|
WHERE turn_row.task_session_id = ?
|
|
AND %s
|
|
ORDER BY turn_row.started_at DESC, turn_row.created_at DESC, turn_row.id DESC
|
|
LIMIT 1
|
|
)
|
|
SELECT m.id, m.task_session_id, m.task_id, m.turn_id, m.author_type, m.author_id,
|
|
m.content, m.requests_input, m.type, m.metadata, m.created_at, m.updated_at
|
|
FROM task_session_messages m
|
|
JOIN current_turn current ON current.id = m.turn_id
|
|
WHERE m.task_session_id = ?
|
|
AND m.type = 'clarification_request'
|
|
AND COALESCE(%s, '') IN ('', 'pending')
|
|
ORDER BY m.created_at ASC, m.id ASC
|
|
`, turnAuthorityPredicate(drv, "turn_row"), dialect.JSONExtract(drv, "m.metadata", "status"))
|
|
// First sessionID selects current-turn authority; second scopes messages so
|
|
// a malformed cross-session turn reference cannot leak another session's row.
|
|
rows, err := r.ro.QueryContext(ctx, r.ro.Rebind(query), sessionID, sessionID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer func() { _ = rows.Close() }()
|
|
result, _, err := scanMessageRows(rows, 0)
|
|
return result, err
|
|
}
|
|
|
|
// GetPendingActionsBySessionIDs returns the compact pending-action projection
|
|
// for sessions whose task-list state needs to render before messages are loaded.
|
|
func (r *Repository) GetPendingActionsBySessionIDs(ctx context.Context, sessionIDs []string) (map[string]models.TaskPendingAction, error) {
|
|
result := make(map[string]models.TaskPendingAction)
|
|
if len(sessionIDs) == 0 {
|
|
return result, nil
|
|
}
|
|
placeholders := make([]string, len(sessionIDs))
|
|
args := make([]interface{}, 0, len(sessionIDs))
|
|
for i, id := range sessionIDs {
|
|
placeholders[i] = "?"
|
|
args = append(args, id)
|
|
}
|
|
query := pendingActionsBySessionQuery(r.ro.DriverName(), placeholders)
|
|
rows, err := r.ro.QueryxContext(ctx, r.ro.Rebind(query), args...)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer func() { _ = rows.Close() }()
|
|
for rows.Next() {
|
|
var sessionID, action string
|
|
if err := rows.Scan(&sessionID, &action); err != nil {
|
|
return nil, err
|
|
}
|
|
// Clarification is the explicit per-session priority when malformed or
|
|
// transitional history contains both pending action kinds. Enforce it
|
|
// here so correctness never depends on UNION ALL row order.
|
|
if action == string(models.TaskPendingActionClarification) {
|
|
result[sessionID] = models.TaskPendingActionClarification
|
|
continue
|
|
}
|
|
if _, hasClarification := result[sessionID]; hasClarification {
|
|
continue
|
|
}
|
|
if action == string(models.TaskPendingActionPermission) {
|
|
result[sessionID] = models.TaskPendingActionPermission
|
|
}
|
|
}
|
|
return result, rows.Err()
|
|
}
|
|
|
|
func pendingActionsBySessionQuery(driverName string, placeholders []string) string {
|
|
placeholderList := strings.Join(placeholders, ",")
|
|
statusExpr := dialect.JSONExtract(driverName, "m.metadata", "status")
|
|
permissionOrderExpr := pendingActionMessageOrder(driverName, "m")
|
|
// Durable turns are authoritative. A message without a matching turn is
|
|
// malformed under the schema foreign key and must not reactivate old input.
|
|
// Terminal sessions quarantine pending history even when best-effort expiry
|
|
// persistence failed. Both message CTEs inherit the requested-session
|
|
// boundary through their task_session_id and turn_id joins to current_turn.
|
|
return fmt.Sprintf(`
|
|
WITH current_turn AS (
|
|
SELECT task_session_id, id AS turn_id
|
|
FROM (
|
|
SELECT task_session_id,
|
|
id,
|
|
ROW_NUMBER() OVER (
|
|
PARTITION BY task_session_id
|
|
ORDER BY started_at DESC, created_at DESC, id DESC
|
|
) AS rn
|
|
FROM task_session_turns turn_row
|
|
WHERE turn_row.task_session_id IN (%s)
|
|
AND %s
|
|
AND %s
|
|
) ranked
|
|
WHERE rn = 1
|
|
),
|
|
pending_clarifications AS (
|
|
SELECT DISTINCT m.task_session_id, 'clarification' AS action
|
|
FROM task_session_messages m
|
|
JOIN current_turn current
|
|
ON current.task_session_id = m.task_session_id
|
|
AND current.turn_id = m.turn_id
|
|
WHERE m.type = 'clarification_request'
|
|
AND COALESCE(%s, '') IN ('', 'pending')
|
|
),
|
|
latest_permissions AS (
|
|
SELECT m.task_session_id,
|
|
COALESCE(%s, '') AS status,
|
|
ROW_NUMBER() OVER (
|
|
PARTITION BY m.task_session_id
|
|
ORDER BY m.created_at DESC, %s DESC
|
|
) AS rn
|
|
FROM task_session_messages m
|
|
JOIN current_turn current
|
|
ON current.task_session_id = m.task_session_id
|
|
AND current.turn_id = m.turn_id
|
|
WHERE m.type = 'permission_request'
|
|
)
|
|
SELECT task_session_id, action
|
|
FROM pending_clarifications
|
|
UNION ALL
|
|
SELECT task_session_id, 'permission' AS action
|
|
FROM latest_permissions
|
|
WHERE rn = 1 AND status IN ('', 'pending')
|
|
`, placeholderList, turnAuthorityPredicate(driverName, "turn_row"),
|
|
nonTerminalSessionPredicate("turn_row"), statusExpr, statusExpr, permissionOrderExpr)
|
|
}
|
|
|
|
func pendingActionMessageOrder(driverName string, qualifier string) string {
|
|
column := "id"
|
|
if driverName == dialect.SQLite3 {
|
|
column = "rowid"
|
|
}
|
|
if qualifier == "" {
|
|
return column
|
|
}
|
|
return qualifier + "." + column
|
|
}
|
|
|
|
// FindMessageByPendingIDAndQuestion finds the message for a specific (pending_id,
|
|
// question_id) pair within a session. Used to flip per-question status (answered /
|
|
// rejected) on multi-question clarification bundles. The persistence layer stores
|
|
// the question id at metadata.question_id (flat, alongside metadata.question) so
|
|
// the JSON path lookup works on both SQLite and Postgres.
|
|
func (r *Repository) FindMessageByPendingIDAndQuestion(ctx context.Context, sessionID, pendingID, questionID string) (*models.Message, error) {
|
|
drv := r.ro.DriverName()
|
|
query := fmt.Sprintf(`
|
|
SELECT id, task_session_id, task_id, turn_id, author_type, author_id, content, requests_input, type, metadata, created_at, updated_at
|
|
FROM task_session_messages
|
|
WHERE task_session_id = ? AND %s = ? AND %s = ?
|
|
ORDER BY created_at ASC LIMIT 1
|
|
`,
|
|
dialect.JSONExtract(drv, "metadata", "pending_id"),
|
|
dialect.JSONExtract(drv, "metadata", "question_id"),
|
|
)
|
|
message := &models.Message{}
|
|
var requestsInput int
|
|
var messageType string
|
|
var metadataJSON string
|
|
err := r.ro.QueryRowContext(ctx, r.ro.Rebind(query), sessionID, pendingID, questionID).Scan(
|
|
&message.ID, &message.TaskSessionID, &message.TaskID, &message.TurnID, &message.AuthorType, &message.AuthorID,
|
|
&message.Content, &requestsInput, &messageType, &metadataJSON, &message.CreatedAt, &message.UpdatedAt,
|
|
)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
message.RequestsInput = requestsInput == 1
|
|
message.Type = models.MessageType(messageType)
|
|
if metadataJSON != "" {
|
|
_ = json.Unmarshal([]byte(metadataJSON), &message.Metadata)
|
|
}
|
|
return message, nil
|
|
}
|
|
|
|
// CompletePendingToolCallsForTurn marks all non-terminal tool call messages for a turn as "complete".
|
|
// This is a safety net to ensure no tool calls remain stuck in a non-terminal state (pending,
|
|
// running, in_progress, etc.) after a turn completes.
|
|
//
|
|
// Excludes permission_request messages: they share `tool_call_id` in metadata but their
|
|
// `status` is the user's approve/reject decision, not the tool call state. Forcing them to
|
|
// "complete" wipes "approved"/"rejected" and re-shows the prompt buttons in the UI.
|
|
func (r *Repository) CompletePendingToolCallsForTurn(ctx context.Context, turnID string) (int64, error) {
|
|
drv := r.db.DriverName()
|
|
query := fmt.Sprintf(`
|
|
UPDATE task_session_messages
|
|
SET metadata = %s, updated_at = CURRENT_TIMESTAMP
|
|
WHERE turn_id = ?
|
|
AND type != 'permission_request'
|
|
AND %s NOT IN ('complete', 'error')
|
|
AND %s
|
|
`, dialect.JSONSet(drv, "metadata", "status", "complete"),
|
|
dialect.JSONExtract(drv, "metadata", "status"),
|
|
dialect.JSONExtractIsNotNull(drv, "metadata", "tool_call_id"))
|
|
result, err := r.db.ExecContext(ctx, r.db.Rebind(query), turnID)
|
|
if err != nil {
|
|
return 0, fmt.Errorf("failed to complete pending 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)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to marshal metadata: %w", err)
|
|
}
|
|
|
|
requestsInput := 0
|
|
if message.RequestsInput {
|
|
requestsInput = 1
|
|
}
|
|
|
|
message.UpdatedAt = time.Now().UTC()
|
|
result, err := r.db.ExecContext(ctx, r.db.Rebind(`
|
|
UPDATE task_session_messages SET content = ?, requests_input = ?, type = ?, metadata = ?, updated_at = ?
|
|
WHERE id = ?
|
|
`), message.Content, requestsInput, string(message.Type), string(metadataJSON), message.UpdatedAt, message.ID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
rows, _ := result.RowsAffected()
|
|
if rows == 0 {
|
|
return fmt.Errorf("message not found: %s", message.ID)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// CountToolCallMessagesBySession returns the number of tool_call messages
|
|
// per session for the given session IDs. Sessions with zero tool calls are
|
|
// omitted from the result map. Powers the "ran N commands" segment of the
|
|
// inline session timeline entry without requiring the frontend to fetch the
|
|
// full message list for each row.
|
|
func (r *Repository) CountToolCallMessagesBySession(ctx context.Context, sessionIDs []string) (map[string]int, error) {
|
|
if len(sessionIDs) == 0 {
|
|
return map[string]int{}, nil
|
|
}
|
|
placeholders := make([]string, len(sessionIDs))
|
|
args := make([]interface{}, 0, len(sessionIDs))
|
|
for i, id := range sessionIDs {
|
|
placeholders[i] = "?"
|
|
args = append(args, id)
|
|
}
|
|
query := `SELECT task_session_id, COUNT(*) AS cnt
|
|
FROM task_session_messages
|
|
WHERE type = 'tool_call'
|
|
AND task_session_id IN (` + strings.Join(placeholders, ",") + `)
|
|
GROUP BY task_session_id`
|
|
rows, err := r.ro.QueryxContext(ctx, r.ro.Rebind(query), args...)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer func() { _ = rows.Close() }()
|
|
out := make(map[string]int)
|
|
for rows.Next() {
|
|
var sessionID string
|
|
var count int
|
|
if err := rows.Scan(&sessionID, &count); err != nil {
|
|
return nil, err
|
|
}
|
|
out[sessionID] = count
|
|
}
|
|
return out, rows.Err()
|
|
}
|