Files
kandev/apps/backend/internal/task/repository/sqlite/message.go
T
2026-09-22 10:31:53 +01:00

1249 lines
48 KiB
Go

package sqlite
import (
"context"
"database/sql"
"encoding/json"
"errors"
"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"
"github.com/kandev/kandev/internal/task/repository/admission"
)
// Message operations
// GetLastMessageTimeBySessionIDs returns the newest task_session_messages
// updated_at for each requested session, one chunked query over
// idx_messages_session_updated. Sessions with no messages are absent from
// the result; callers fall back to the session row's own timestamps. The
// session reconciliation sweep uses this to measure per-session event
// silence without loading any transcript content.
func (r *Repository) GetLastMessageTimeBySessionIDs(ctx context.Context, sessionIDs []string) (map[string]time.Time, error) {
result := make(map[string]time.Time, len(sessionIDs))
if len(sessionIDs) == 0 {
return result, nil
}
for _, chunk := range chunkIDs(sessionIDs, sqliteMaxHostParams) {
placeholders, ids := buildInPlaceholders(chunk)
query := `
SELECT task_session_id, MAX(updated_at)
FROM task_session_messages
WHERE task_session_id IN (` + placeholders + `)
GROUP BY task_session_id`
rows, err := r.ro.QueryContext(ctx, r.ro.Rebind(query), ids...)
if err != nil {
return nil, fmt.Errorf("load last message time by session: %w", err)
}
for rows.Next() {
var (
sessionID string
rawLastMessage interface{}
)
if err := rows.Scan(&sessionID, &rawLastMessage); err != nil {
_ = rows.Close()
return nil, fmt.Errorf("scan last message time by session: %w", err)
}
lastMessage, err := parseTaskActivityTime(rawLastMessage)
if err != nil {
_ = rows.Close()
return nil, fmt.Errorf("parse last message time for session %q: %w", sessionID, err)
}
result[sessionID] = lastMessage
}
if err := rows.Err(); err != nil {
_ = rows.Close()
return nil, fmt.Errorf("read last message time by session: %w", err)
}
_ = rows.Close()
}
return result, nil
}
// 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.AuthorType == "" {
message.AuthorType = models.MessageAuthorUser
}
requestsInput := 0
if message.RequestsInput {
requestsInput = 1
}
messageType := string(message.Type)
if messageType == "" {
messageType = string(models.MessageTypeMessage)
}
metadataJSON, err := r.externalizeMessagePayload(ctx, message)
if err != nil {
return err
}
if message.AuthorType == models.MessageAuthorUser {
// User messages get their prompt-ordering timestamp and ordinal from
// the atomic per-session create boundary (see
// createUserMessageWithBoundary). Agent/tool messages stay on the hot
// path below.
return r.createUserMessageWithBoundary(ctx, message, requestsInput, messageType, metadataJSON)
}
if message.CreatedAt.IsZero() {
message.CreatedAt = time.Now().UTC()
}
if message.UpdatedAt.IsZero() {
message.UpdatedAt = message.CreatedAt
}
return r.insertMessageWithSessionLock(ctx, message, requestsInput, messageType, metadataJSON)
}
// CreateMessageWithInitialTaskBrief persists a user message while atomically
// deciding whether the supplied prepared-session candidate owns the session's
// first prompt slot.
func (r *Repository) CreateMessageWithInitialTaskBrief(
ctx context.Context,
message *models.Message,
candidate *admission.InitialTaskBriefCandidate,
) error {
if candidate == nil {
return r.CreateMessage(ctx, message)
}
if message.ID == "" {
message.ID = uuid.New().String()
}
if message.AuthorType == "" {
message.AuthorType = models.MessageAuthorUser
}
if message.AuthorType != models.MessageAuthorUser {
return fmt.Errorf("initial task brief requires a user message")
}
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.createUserMessageWithBoundaryAndInitialTaskBrief(
ctx, message, requestsInput, messageType, metadataJSON, candidate,
)
}
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, prompt_seq, payload_digest, payload_size)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
`), message.ID, message.TaskSessionID, message.TaskID, message.TurnID, message.AuthorType, message.AuthorID, message.Content, requestsInput, messageType, metadataJSON, message.CreatedAt, message.UpdatedAt, message.PromptIndex, message.PayloadDigest, message.PayloadSize)
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.insertMessageWithPayloadGuard(ctx, 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, payload_digest, payload_size
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, &message.PayloadDigest, &message.PayloadSize)
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"
}
var (
cursor *models.Message
err error
)
if opts.Around != "" {
cursor, err = r.resolveAroundTarget(ctx, sessionID, opts.Around)
} else {
cursor, err = r.resolveMessageCursor(ctx, sessionID, opts)
}
if err != nil {
return nil, false, err
}
if cursor != nil && opts.AuthorType == string(models.MessageAuthorUser) &&
cursor.AuthorType != models.MessageAuthorUser {
return nil, false, fmt.Errorf("message cursor not found: %s", cursor.ID)
}
var cursorKey string
if cursor != nil {
cursorKey, err = r.messageCursorKey(ctx, cursor.ID)
if err != nil {
if opts.Around != "" && errors.Is(err, sql.ErrNoRows) {
return nil, false, fmt.Errorf("%w: %s", ErrMessageNotFound, cursor.ID)
}
return nil, false, err
}
}
if opts.Around != "" {
return r.listMessagesAround(ctx, sessionID, cursor, cursorKey, limit)
}
query, args := buildListMessagesQuery(r.ro.DriverName(), sessionID, opts, cursor, cursorKey, 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 scanPromptIndexedMessageRows(rows, limit)
}
// messageCursorKey returns the normalized-microsecond ordering key of the
// cursor row, derived through the same SQL expression the per-row key uses.
// On SQLite the expression yields the key text directly; on PostgreSQL the
// stored timestamp is scanned and formatted to the exact key bytes (the
// stored value is UTC by convention, and a naive timestamp carries no offset
// whose parsing could diverge from the per-row expression).
func (r *Repository) messageCursorKey(ctx context.Context, id string) (string, error) {
drv := r.ro.DriverName()
query := fmt.Sprintf(
"SELECT %s FROM task_session_messages WHERE id = ?",
dialect.NormalizedMicrosecond(drv, "created_at"),
)
if dialect.IsPostgres(drv) {
var ts time.Time
if err := r.ro.QueryRowContext(ctx, r.ro.Rebind(query), id).Scan(&ts); err != nil {
return "", err
}
return formatPromptKey(ts), nil
}
var key string
if err := r.ro.QueryRowContext(ctx, r.ro.Rebind(query), id).Scan(&key); err != nil {
return "", err
}
return key, nil
}
// 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
}
// resolveAroundTarget validates and loads the target message for an around-window read.
func (r *Repository) resolveAroundTarget(ctx context.Context, sessionID, id string) (*models.Message, error) {
target, err := r.GetMessage(ctx, id)
if err != nil {
if errors.Is(err, sql.ErrNoRows) {
return nil, fmt.Errorf("%w: %s", ErrMessageNotFound, id)
}
return nil, err
}
if target.TaskSessionID != sessionID {
return nil, fmt.Errorf("%w: %s", ErrMessageNotFound, id)
}
return target, nil
}
// listMessagesAround returns the target and newer messages in newest-first order,
// along with whether the bounded window was truncated.
func (r *Repository) listMessagesAround(
ctx context.Context,
sessionID string,
target *models.Message,
targetKey string,
limit int,
) ([]*models.Message, bool, error) {
nm := dialect.NormalizedMicrosecond(r.ro.DriverName(), "created_at")
bound := "?"
if dialect.IsPostgres(r.ro.DriverName()) {
bound = "CAST(? AS timestamp)"
}
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,
CASE WHEN author_type = 'user' THEN prompt_seq ELSE 0 END AS prompt_index
FROM task_session_messages
WHERE task_session_id = ? AND (%s > %s OR (%s = %s AND id >= ?))
ORDER BY %s ASC, id ASC`, nm, bound, nm, bound, nm)
args := []interface{}{sessionID, targetKey, targetKey, target.ID}
if limit > 0 {
query += sqlLimitClause
args = append(args, limit+1)
}
rows, err := r.ro.QueryContext(ctx, r.ro.Rebind(query), args...)
if err != nil {
return nil, false, err
}
defer func() { _ = rows.Close() }()
messages, hasMore, err := scanPromptIndexedMessageRows(rows, limit)
if err != nil {
return nil, false, err
}
for left, right := 0, len(messages)-1; left < right; left, right = left+1, right-1 {
messages[left], messages[right] = messages[right], messages[left]
}
return messages, hasMore, nil
}
// buildListMessagesQuery assembles the paginated message list query and bound arguments, ordering rows by the normalized-microsecond key with cursor and limit bounds.
func buildListMessagesQuery(driverName, sessionID string, opts models.ListMessagesOptions, cursor *models.Message, cursorKey string, sortDir string, limit int) (string, []interface{}) {
nm := dialect.NormalizedMicrosecond(driverName, "created_at")
// The cursor bound is the normalized key literal `YYYY-MM-DD HH:MM:SS.ffffff`
// (space separator, zero-padded 6-digit fraction, no `Z`, UTC). PostgreSQL
// casts it to `timestamp` so it compares against its native per-row key;
// SQLite compares the text key directly.
bound := "?"
if dialect.IsPostgres(driverName) {
bound = "CAST(? AS timestamp)"
}
// prompt_index is the durable per-session ordinal persisted at creation
// (prompt_seq), so it is a plain column: absolute rather than
// page-relative by construction, no window pass needed. The cursor
// predicate, ORDER BY, and LIMIT (limit+1) use the same normalized key as
// the rest of the package, so every page is contiguous in ordinal order
// and the sentinel `promptNumber === 1` stop can never skip a
// higher-ordinal prompt sharing a microsecond.
query := `
SELECT id, task_session_id, task_id, turn_id, author_type, author_id, content, requests_input, type, metadata, created_at, updated_at,
CASE WHEN author_type = 'user' THEN prompt_seq ELSE 0 END AS prompt_index
FROM task_session_messages
WHERE task_session_id = ?`
args := []interface{}{sessionID}
if len(opts.AuthorTypes) > 0 {
placeholders := make([]string, len(opts.AuthorTypes))
for index, author := range opts.AuthorTypes {
placeholders[index] = "?"
args = append(args, author)
}
query += " AND author_type IN (" + strings.Join(placeholders, ",") + ")"
} else if opts.AuthorType != "" {
query += " AND author_type = ?"
args = append(args, opts.AuthorType)
}
if opts.TaskID != "" {
query += " AND task_id = ?"
args = append(args, opts.TaskID)
}
if cursor != nil {
if opts.Before != "" {
query += fmt.Sprintf(" AND (%s < %s OR (%s = %s AND id < ?))", nm, bound, nm, bound)
} else if opts.After != "" {
query += fmt.Sprintf(" AND (%s > %s OR (%s = %s AND id > ?))", nm, bound, nm, bound)
}
args = append(args, cursorKey, cursorKey, cursor.ID)
}
query += fmt.Sprintf(" ORDER BY %s %s, id %s", nm, sortDir, sortDir)
if limit > 0 {
query += sqlLimitClause
args = append(args, limit+1)
}
return query, args
}
// scanPromptIndexedMessageRows scans the 13-column indexed message projection
// (the 12 legacy columns plus the computed prompt_index), returning the
// page-truncated slice and whether more rows remain, like scanMessageRows.
func scanPromptIndexedMessageRows(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, &message.PromptIndex); 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
}
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, "")
}
// GetPermissionMessageByIdentity retrieves a permission_request message by its
// full (task, session, request, pending) identity. A provider may reuse
// pending_id for a later, unrelated request once the original is resolved, so
// matching request_id too prevents a delayed event about the old request from
// being applied to the new one.
func (r *Repository) GetPermissionMessageByIdentity(ctx context.Context, taskID, sessionID, requestID, pendingID string) (*models.Message, error) {
return r.getPermissionMessageByIdentity(ctx, taskID, sessionID, requestID, 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 and message ID (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) {
query := findMessagesByPendingIDQuery(r.ro.DriverName())
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
}
// FindMessagesByPendingIDs returns all messages for the requested pending IDs
// in one or more bounded queries, grouped by pending ID. The chunking keeps
// the query below SQLite's host-parameter limit while replacing the inbox's
// per-row message hydration loop.
func (r *Repository) FindMessagesByPendingIDs(
ctx context.Context, pendingIDs []string,
) (map[string][]*models.Message, error) {
result := make(map[string][]*models.Message, len(pendingIDs))
if len(pendingIDs) == 0 {
return result, nil
}
driverName := r.ro.DriverName()
pendingExpr := dialect.JSONExtract(driverName, "metadata", "pending_id")
for _, chunk := range chunkIDs(pendingIDs, sqliteMaxHostParams) {
placeholders, args := buildInPlaceholders(chunk)
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 IN (%s)
ORDER BY %s ASC, created_at ASC, id ASC
`, pendingExpr, placeholders, pendingExpr)
rows, err := r.ro.QueryContext(ctx, r.ro.Rebind(query), args...)
if err != nil {
return nil, err
}
messages, _, scanErr := scanMessageRows(rows, 0)
_ = rows.Close()
if scanErr != nil {
return nil, scanErr
}
for _, message := range messages {
pendingID, _ := message.Metadata["pending_id"].(string)
if pendingID != "" {
result[pendingID] = append(result[pendingID], message)
}
}
}
return result, nil
}
// findMessagesByPendingIDQuery is shared with query-plan tests so the
// production lookup and its planner witness cannot drift apart.
func findMessagesByPendingIDQuery(driverName string) string {
return 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, id ASC
`, dialect.JSONExtract(driverName, "metadata", "pending_id"))
}
// 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()
predicate, orderBy := currentTurnAuthority(drv, "turn_row")
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 %s
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
`, predicate, orderBy, 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.
predicate, orderBy := currentTurnAuthority(driverName, "turn_row")
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 %s
) 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')
`, orderBy, placeholderList, predicate,
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) {
tx, err := r.db.BeginTxx(ctx, nil)
if err != nil {
return 0, err
}
defer func() { _ = tx.Rollback() }()
rows, err := r.completePendingToolCallsForTurnTx(ctx, tx, turnID)
if err != nil {
return 0, err
}
if err := tx.Commit(); err != nil {
return 0, err
}
return rows, nil
}
// UpdateMessage updates an existing message
func (r *Repository) UpdateMessage(ctx context.Context, message *models.Message) error {
metadataJSON, err := r.externalizeMessagePayload(ctx, message)
if err != nil {
return err
}
requestsInput := 0
if message.RequestsInput {
requestsInput = 1
}
message.UpdatedAt = time.Now().UTC()
if !dialect.IsPostgres(r.db.DriverName()) {
return r.updateMessageWithPayloadGuard(ctx, message, []byte(metadataJSON), requestsInput)
}
result, err := r.db.ExecContext(ctx, r.db.Rebind(`
UPDATE task_session_messages SET content = ?, requests_input = ?, type = ?, metadata = ?, payload_digest = ?, payload_size = ?, updated_at = ?
WHERE id = ?
`), message.Content, requestsInput, string(message.Type), metadataJSON, message.PayloadDigest, message.PayloadSize, 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
}
// ClaimPermissionResolution installs the first durable dispatch claim for one
// exact pending permission. The conditional UPDATE is the serialization point
// on both SQLite and PostgreSQL.
func (r *Repository) ClaimPermissionResolution(ctx context.Context, request models.PermissionResolutionClaimRequest) (*models.PermissionResolutionClaimResult, error) {
request.Audit.Result = models.PermissionResolutionDispatching
auditJSON, err := json.Marshal(request.Audit)
if err != nil {
return nil, fmt.Errorf("marshal permission resolution claim: %w", err)
}
driver := r.db.DriverName()
now := time.Now().UTC()
query := fmt.Sprintf(`
UPDATE task_session_messages
SET metadata = %s, updated_at = ?
WHERE task_id = ? AND task_session_id = ? AND type = 'permission_request'
AND %s = ? AND %s = ?
AND COALESCE(%s, '') = ''
AND %s IS NULL
`, permissionClaimJSONExpression(driver),
permissionJSONExtract(driver, "metadata", "request_id"),
permissionJSONExtract(driver, "metadata", "pending_id"),
permissionJSONExtract(driver, "metadata", "status"),
permissionJSONExtract(driver, "metadata", "permission_resolution"))
result, err := r.db.ExecContext(ctx, r.db.Rebind(query), string(auditJSON), now,
request.TaskID, request.SessionID, request.Audit.RequestID, request.Audit.PendingID)
if err != nil {
return nil, fmt.Errorf("claim permission resolution: %w", err)
}
rows, err := result.RowsAffected()
if err != nil {
return nil, fmt.Errorf("claim permission resolution rows: %w", err)
}
message, getErr := r.getPermissionMessageByIdentity(ctx, request.TaskID, request.SessionID, request.Audit.RequestID, request.Audit.PendingID)
if rows == 1 {
if getErr != nil {
return nil, fmt.Errorf("read claimed permission resolution: %w", getErr)
}
return &models.PermissionResolutionClaimResult{Outcome: models.PermissionClaimed, Message: message}, nil
}
if errors.Is(getErr, sql.ErrNoRows) {
return &models.PermissionResolutionClaimResult{Outcome: models.PermissionClaimNotFound}, nil
}
if getErr != nil {
return nil, fmt.Errorf("read unclaimed permission resolution: %w", getErr)
}
audit, ok := permissionAuditFromMetadata(message.Metadata)
if ok && audit.Result == models.PermissionResolutionDispatching {
return &models.PermissionResolutionClaimResult{Outcome: models.PermissionClaimInProgress, Message: message}, nil
}
return &models.PermissionResolutionClaimResult{Outcome: models.PermissionClaimAlreadyFinal, Message: message}, nil
}
// FinalizePermissionResolution closes only the caller's dispatch claim. A
// failed finalization never reopens the request for another provider response.
func (r *Repository) FinalizePermissionResolution(ctx context.Context, request models.PermissionResolutionFinalizeRequest) (*models.PermissionResolutionFinalizeResult, error) {
driver := r.db.DriverName()
finalizedAt := request.FinalizedAt.UTC()
query := fmt.Sprintf(`
UPDATE task_session_messages
SET metadata = %s, updated_at = ?
WHERE task_id = ? AND task_session_id = ? AND type = 'permission_request'
AND %s = ? AND %s = ?
AND %s = ? AND %s = ?
`, permissionFinalizeJSONExpression(driver),
permissionJSONExtract(driver, "metadata", "request_id"),
permissionJSONExtract(driver, "metadata", "pending_id"),
permissionJSONExtract(driver, "metadata", "permission_resolution", "claim_id"),
permissionJSONExtract(driver, "metadata", "permission_resolution", "result"))
result, err := r.db.ExecContext(ctx, r.db.Rebind(query), string(request.Result), finalizedAt.Format(time.RFC3339Nano), string(request.Status), finalizedAt,
request.TaskID, request.SessionID, request.RequestID, request.PendingID, request.ClaimID, string(models.PermissionResolutionDispatching))
if err != nil {
return nil, fmt.Errorf("finalize permission resolution: %w", err)
}
rows, err := result.RowsAffected()
if err != nil {
return nil, fmt.Errorf("finalize permission resolution rows: %w", err)
}
message, getErr := r.getPermissionMessageByIdentity(ctx, request.TaskID, request.SessionID, request.RequestID, request.PendingID)
if rows == 1 {
if getErr != nil {
return nil, fmt.Errorf("read finalized permission resolution: %w", getErr)
}
return &models.PermissionResolutionFinalizeResult{Outcome: models.PermissionFinalized, Message: message}, nil
}
if errors.Is(getErr, sql.ErrNoRows) {
return &models.PermissionResolutionFinalizeResult{Outcome: models.PermissionFinalizeNotFound}, nil
}
if getErr != nil {
return nil, fmt.Errorf("read unfinalized permission resolution: %w", getErr)
}
audit, ok := permissionAuditFromMetadata(message.Metadata)
if !ok || audit.ClaimID != request.ClaimID {
return &models.PermissionResolutionFinalizeResult{Outcome: models.PermissionFinalizeClaimMismatch, Message: message}, nil
}
return &models.PermissionResolutionFinalizeResult{Outcome: models.PermissionFinalizeAlreadyFinal, Message: message}, nil
}
func permissionClaimJSONExpression(driver string) string {
if dialect.IsPostgres(driver) {
return `jsonb_set(COALESCE(NULLIF(metadata, ''), '{}')::jsonb, '{permission_resolution}', ?::jsonb, true)::text`
}
return `json_set(COALESCE(NULLIF(metadata, ''), '{}'), '$.permission_resolution', json(?))`
}
func permissionFinalizeJSONExpression(driver string) string {
if dialect.IsPostgres(driver) {
return `jsonb_set(jsonb_set(jsonb_set(COALESCE(NULLIF(metadata, ''), '{}')::jsonb, '{permission_resolution,result}', to_jsonb(?::text), true), '{permission_resolution,finalized_at}', to_jsonb(?::text), true), '{status}', to_jsonb(?::text), true)::text`
}
return `json_set(COALESCE(NULLIF(metadata, ''), '{}'), '$.permission_resolution.result', ?, '$.permission_resolution.finalized_at', ?, '$.status', ?)`
}
func permissionJSONExtract(driver, column string, path ...string) string {
if dialect.IsPostgres(driver) {
return fmt.Sprintf("COALESCE(NULLIF(%s, ''), '{}')::jsonb#>>'{%s}'", column, strings.Join(path, ","))
}
return fmt.Sprintf("json_extract(CASE WHEN json_valid(%s) THEN %s ELSE '{}' END, '$.%s')", column, column, strings.Join(path, "."))
}
func (r *Repository) getPermissionMessageByIdentity(ctx context.Context, taskID, sessionID, requestID, pendingID string) (*models.Message, error) {
driver := r.db.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_id = ? AND task_session_id = ? AND type = 'permission_request'
AND %s = ? AND %s = ?
LIMIT 1
`, permissionJSONExtract(driver, "metadata", "request_id"), permissionJSONExtract(driver, "metadata", "pending_id"))
message := &models.Message{}
var requestsInput int
var messageType string
var metadataJSON string
err := r.db.QueryRowContext(ctx, r.db.Rebind(query), taskID, sessionID, requestID, 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 != "" {
if err := json.Unmarshal([]byte(metadataJSON), &message.Metadata); err != nil {
return nil, fmt.Errorf("deserialize permission message metadata: %w", err)
}
}
return message, nil
}
func permissionAuditFromMetadata(metadata map[string]any) (models.PermissionResolutionAudit, bool) {
raw, ok := metadata["permission_resolution"]
if !ok {
return models.PermissionResolutionAudit{}, false
}
encoded, err := json.Marshal(raw)
if err != nil {
return models.PermissionResolutionAudit{}, false
}
var audit models.PermissionResolutionAudit
if err := json.Unmarshal(encoded, &audit); err != nil {
return models.PermissionResolutionAudit{}, false
}
return audit, audit.ClaimID != ""
}
func (r *Repository) GetPermissionResolutionAudit(ctx context.Context, taskID, sessionID, requestID, pendingID string) (*models.PermissionResolutionAudit, error) {
message, err := r.getPermissionMessageByIdentity(ctx, taskID, sessionID, requestID, pendingID)
if err != nil {
if errors.Is(err, sql.ErrNoRows) {
return nil, nil
}
return nil, err
}
audit, ok := permissionAuditFromMetadata(message.Metadata)
if !ok {
return nil, nil
}
return &audit, 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()
}