diff --git a/go/api/database/client.go b/go/api/database/client.go index e80dc8647..f52f79d3e 100644 --- a/go/api/database/client.go +++ b/go/api/database/client.go @@ -2,6 +2,7 @@ package database import ( "context" + "errors" "time" a2a "github.com/a2aproject/a2a-go/v2/a2a" @@ -9,6 +10,10 @@ import ( "github.com/pgvector/pgvector-go" ) +// ErrSessionIDInUse means the requested session id is already active on a +// different session (possibly owned by another user). +var ErrSessionIDInUse = errors.New("session id already in use") + type QueryOptions struct { Limit int After time.Time diff --git a/go/core/internal/database/client_postgres.go b/go/core/internal/database/client_postgres.go index 4e0bdc66d..ebb6c7ad3 100644 --- a/go/core/internal/database/client_postgres.go +++ b/go/core/internal/database/client_postgres.go @@ -11,6 +11,7 @@ import ( a2a "github.com/a2aproject/a2a-go/v2/a2a" "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" "github.com/jackc/pgx/v5/pgtype" "github.com/jackc/pgx/v5/pgxpool" dbpkg "github.com/kagent-dev/kagent/go/api/database" @@ -83,7 +84,7 @@ func (c *postgresClient) DeleteAgent(ctx context.Context, agentID string) error // ── Sessions ────────────────────────────────────────────────────────────────── func (c *postgresClient) StoreSession(ctx context.Context, session *dbpkg.Session) error { - return c.withTx(ctx, func(q *dbgen.Queries) error { + err := c.withTx(ctx, func(q *dbgen.Queries) error { params := dbgen.UpsertSessionParams{ ID: session.ID, UserID: session.UserID, @@ -96,6 +97,11 @@ func (c *postgresClient) StoreSession(ctx context.Context, session *dbpkg.Sessio } return q.UpsertSession(ctx, params) }) + var pgErr *pgconn.PgError + if errors.As(err, &pgErr) && pgErr.ConstraintName == "session_id_active_unique" { + return dbpkg.ErrSessionIDInUse + } + return err } func (c *postgresClient) GetSession(ctx context.Context, sessionID, userID string) (*dbpkg.Session, error) { @@ -146,7 +152,24 @@ func (c *postgresClient) ListSessionsForAgentAllUsers(ctx context.Context, agent } func (c *postgresClient) DeleteSession(ctx context.Context, sessionID, userID string) error { - return c.q.SoftDeleteSession(ctx, dbgen.SoftDeleteSessionParams{ID: sessionID, UserID: userID}) + return c.withTx(ctx, func(q *dbgen.Queries) error { + if _, err := q.GetSession(ctx, dbgen.GetSessionParams{ID: sessionID, UserID: userID}); err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return nil + } + return err + } + if err := q.SoftDeleteTasksBySession(ctx, &sessionID); err != nil { + return err + } + if err := q.SoftDeleteEventsBySession(ctx, &sessionID); err != nil { + return err + } + if err := q.DeleteSessionSharesBySession(ctx, sessionID); err != nil { + return err + } + return q.SoftDeleteSession(ctx, dbgen.SoftDeleteSessionParams{ID: sessionID, UserID: userID}) + }) } // ── Session Shares ───────────────────────────────────────────────────────────── diff --git a/go/core/internal/database/client_test.go b/go/core/internal/database/client_test.go index a8c263d7d..13e7e5d6e 100644 --- a/go/core/internal/database/client_test.go +++ b/go/core/internal/database/client_test.go @@ -234,6 +234,24 @@ func TestStoreSessionIdempotence(t *testing.T) { assert.Equal(t, "Updated", *retrieved.Name, "Session should have updated name") } +func TestStoreSessionRejectsIDUsedByAnotherUser(t *testing.T) { + db := setupTestDB(t) + client := NewClient(db) + ctx := context.Background() + + agentID := "agent-1" + first := &dbpkg.Session{ID: "shared-id", UserID: "user-a", AgentID: &agentID} + require.NoError(t, client.StoreSession(ctx, first), "first user should get the id") + + second := &dbpkg.Session{ID: "shared-id", UserID: "user-b", AgentID: &agentID} + err := client.StoreSession(ctx, second) + require.ErrorIs(t, err, dbpkg.ErrSessionIDInUse, "a second user must not claim an id already active for another user") + + // Once the first user's session is gone, the id is free again. + require.NoError(t, client.DeleteSession(ctx, "shared-id", "user-a")) + require.NoError(t, client.StoreSession(ctx, second), "id should be reusable after the original session is deleted") +} + func TestListSessionsOrdersByRecentActivity(t *testing.T) { db := setupTestDB(t) client := NewClient(db) diff --git a/go/core/internal/database/gen/events.sql.go b/go/core/internal/database/gen/events.sql.go index d34d42298..7b4faf749 100644 --- a/go/core/internal/database/gen/events.sql.go +++ b/go/core/internal/database/gen/events.sql.go @@ -329,3 +329,12 @@ func (q *Queries) SoftDeleteEvent(ctx context.Context, id string) error { _, err := q.db.Exec(ctx, softDeleteEvent, id) return err } + +const softDeleteEventsBySession = `-- name: SoftDeleteEventsBySession :exec +UPDATE event SET deleted_at = NOW() WHERE session_id = $1 AND deleted_at IS NULL +` + +func (q *Queries) SoftDeleteEventsBySession(ctx context.Context, sessionID *string) error { + _, err := q.db.Exec(ctx, softDeleteEventsBySession, sessionID) + return err +} diff --git a/go/core/internal/database/gen/querier.go b/go/core/internal/database/gen/querier.go index 078a06502..357633893 100644 --- a/go/core/internal/database/gen/querier.go +++ b/go/core/internal/database/gen/querier.go @@ -13,6 +13,7 @@ type Querier interface { DeleteAgentMemory(ctx context.Context, arg DeleteAgentMemoryParams) error DeleteExpiredMemories(ctx context.Context) error DeleteSessionShare(ctx context.Context, arg DeleteSessionShareParams) error + DeleteSessionSharesBySession(ctx context.Context, sessionID string) error ExtendMemoryTTL(ctx context.Context) error GetAgent(ctx context.Context, id string) (Agent, error) GetCheckpoint(ctx context.Context, arg GetCheckpointParams) (LgCheckpoint, error) @@ -60,9 +61,11 @@ type Querier interface { SoftDeleteCheckpointWrites(ctx context.Context, arg SoftDeleteCheckpointWritesParams) error SoftDeleteCheckpoints(ctx context.Context, arg SoftDeleteCheckpointsParams) error SoftDeleteEvent(ctx context.Context, id string) error + SoftDeleteEventsBySession(ctx context.Context, sessionID *string) error SoftDeletePushNotification(ctx context.Context, taskID string) error SoftDeleteSession(ctx context.Context, arg SoftDeleteSessionParams) error SoftDeleteTask(ctx context.Context, id string) error + SoftDeleteTasksBySession(ctx context.Context, sessionID *string) error SoftDeleteToolServer(ctx context.Context, arg SoftDeleteToolServerParams) error SoftDeleteToolsForServer(ctx context.Context, arg SoftDeleteToolsForServerParams) error TaskExists(ctx context.Context, id string) (bool, error) diff --git a/go/core/internal/database/gen/session_shares.sql.go b/go/core/internal/database/gen/session_shares.sql.go index 9e84e5363..37a91d6ce 100644 --- a/go/core/internal/database/gen/session_shares.sql.go +++ b/go/core/internal/database/gen/session_shares.sql.go @@ -57,6 +57,15 @@ func (q *Queries) DeleteSessionShare(ctx context.Context, arg DeleteSessionShare return err } +const deleteSessionSharesBySession = `-- name: DeleteSessionSharesBySession :exec +DELETE FROM session_share WHERE session_id = $1 +` + +func (q *Queries) DeleteSessionSharesBySession(ctx context.Context, sessionID string) error { + _, err := q.db.Exec(ctx, deleteSessionSharesBySession, sessionID) + return err +} + const getSessionShareByToken = `-- name: GetSessionShareByToken :one SELECT id, token, session_id, user_id, read_only, created_at FROM session_share WHERE token = $1 diff --git a/go/core/internal/database/gen/tasks.sql.go b/go/core/internal/database/gen/tasks.sql.go index be4e93d75..fd3173495 100644 --- a/go/core/internal/database/gen/tasks.sql.go +++ b/go/core/internal/database/gen/tasks.sql.go @@ -73,6 +73,15 @@ func (q *Queries) SoftDeleteTask(ctx context.Context, id string) error { return err } +const softDeleteTasksBySession = `-- name: SoftDeleteTasksBySession :exec +UPDATE task SET deleted_at = NOW() WHERE session_id = $1 AND deleted_at IS NULL +` + +func (q *Queries) SoftDeleteTasksBySession(ctx context.Context, sessionID *string) error { + _, err := q.db.Exec(ctx, softDeleteTasksBySession, sessionID) + return err +} + const taskExists = `-- name: TaskExists :one SELECT EXISTS ( SELECT 1 FROM task WHERE id = $1 AND deleted_at IS NULL diff --git a/go/core/internal/database/queries/events.sql b/go/core/internal/database/queries/events.sql index 9f916a03d..1813256c3 100644 --- a/go/core/internal/database/queries/events.sql +++ b/go/core/internal/database/queries/events.sql @@ -57,3 +57,6 @@ LIMIT $2; -- name: SoftDeleteEvent :exec UPDATE event SET deleted_at = NOW() WHERE id = $1 AND deleted_at IS NULL; + +-- name: SoftDeleteEventsBySession :exec +UPDATE event SET deleted_at = NOW() WHERE session_id = $1 AND deleted_at IS NULL; diff --git a/go/core/internal/database/queries/session_shares.sql b/go/core/internal/database/queries/session_shares.sql index 48f215352..7e6bec957 100644 --- a/go/core/internal/database/queries/session_shares.sql +++ b/go/core/internal/database/queries/session_shares.sql @@ -21,3 +21,6 @@ WHERE token = $1 AND session_id = $2 AND user_id = $3; INSERT INTO session_share_access (user_id, share_id, accessed_at) VALUES ($1, $2, NOW()) ON CONFLICT (user_id, share_id) DO UPDATE SET accessed_at = NOW(); + +-- name: DeleteSessionSharesBySession :exec +DELETE FROM session_share WHERE session_id = $1; diff --git a/go/core/internal/database/queries/tasks.sql b/go/core/internal/database/queries/tasks.sql index d41fcdf92..bab66114e 100644 --- a/go/core/internal/database/queries/tasks.sql +++ b/go/core/internal/database/queries/tasks.sql @@ -33,3 +33,6 @@ WHERE upserted_task.session_id IS NOT NULL -- name: SoftDeleteTask :exec UPDATE task SET deleted_at = NOW() WHERE id = $1 AND deleted_at IS NULL; + +-- name: SoftDeleteTasksBySession :exec +UPDATE task SET deleted_at = NOW() WHERE session_id = $1 AND deleted_at IS NULL; diff --git a/go/core/internal/httpserver/handlers/sessions.go b/go/core/internal/httpserver/handlers/sessions.go index 63935605e..aed1d35e4 100644 --- a/go/core/internal/httpserver/handlers/sessions.go +++ b/go/core/internal/httpserver/handlers/sessions.go @@ -2,6 +2,7 @@ package handlers import ( "context" + stderrors "errors" "fmt" "net/http" "strconv" @@ -178,6 +179,10 @@ func (h *SessionsHandler) HandleCreateSession(w ErrorResponseWriter, r *http.Req "name", sessionRequest.Name) if err := h.DatabaseService.StoreSession(r.Context(), session); err != nil { + if stderrors.Is(err, database.ErrSessionIDInUse) { + w.RespondWithError(errors.NewConflictError("Session ID is already in use", err)) + return + } w.RespondWithError(errors.NewInternalServerError("Failed to create session", err)) return } diff --git a/go/core/pkg/migrations/core/000007_session_id_unique.down.sql b/go/core/pkg/migrations/core/000007_session_id_unique.down.sql new file mode 100644 index 000000000..567e78c81 --- /dev/null +++ b/go/core/pkg/migrations/core/000007_session_id_unique.down.sql @@ -0,0 +1 @@ +DROP INDEX IF EXISTS session_id_active_unique; diff --git a/go/core/pkg/migrations/core/000007_session_id_unique.up.sql b/go/core/pkg/migrations/core/000007_session_id_unique.up.sql new file mode 100644 index 000000000..9fe2b4c4e --- /dev/null +++ b/go/core/pkg/migrations/core/000007_session_id_unique.up.sql @@ -0,0 +1,5 @@ +-- A session's client-supplied id was only unique per (id, user_id), so two +-- different users could create sessions with the same id. Deleting one then +-- cascaded to the other user's tasks and events, which are keyed by session_id +-- alone. This index makes id unique among live sessions so that can't happen. +CREATE UNIQUE INDEX IF NOT EXISTS session_id_active_unique ON session (id) WHERE deleted_at IS NULL;