diff --git a/go/api/database/client.go b/go/api/database/client.go index b58942479..629d9fff5 100644 --- a/go/api/database/client.go +++ b/go/api/database/client.go @@ -10,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") + // ErrTaskOwnedByAnotherUser means a task with this id already belongs to a // different user. var ErrTaskOwnedByAnotherUser = errors.New("task id owned by another user") diff --git a/go/core/internal/database/client_postgres.go b/go/core/internal/database/client_postgres.go index 42fef4327..d52985d1e 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" @@ -92,7 +93,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, @@ -105,6 +106,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) { @@ -155,7 +161,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 75f4eb3a0..37097b454 100644 --- a/go/core/internal/database/client_test.go +++ b/go/core/internal/database/client_test.go @@ -240,6 +240,24 @@ func TestStoreSessionIdempotence(t *testing.T) { require.Error(t, err, "another user's session must not be readable") } +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) @@ -429,9 +447,13 @@ func TestNullOwnedTaskAccess(t *testing.T) { err = client.StoreTask(ctx, &a2a.Task{ID: "t-legacy", ContextID: "s-mine"}, "bob") require.ErrorIs(t, err, dbpkg.ErrTaskOwnedByAnotherUser, "the claim must stick") - // A session id used by two users is ambiguous: the NULL-owned task stays - // hidden from both, and neither can claim it. + // A session id used by two users across its history is ambiguous: the + // NULL-owned task stays hidden from both, and neither can claim it. Two + // live sessions can no longer share an id (session_id_active_unique), so + // the ambiguity is built from alice's session having existed and been + // deleted before bob's session took over the same id. require.NoError(t, client.StoreSession(ctx, &dbpkg.Session{ID: "s-shared", UserID: "alice"})) + require.NoError(t, client.DeleteSession(ctx, "s-shared", "alice")) require.NoError(t, client.StoreSession(ctx, &dbpkg.Session{ID: "s-shared", UserID: "bob"})) seedNullTask("t-ambiguous", "s-shared") @@ -477,39 +499,34 @@ func TestNullOwnedTaskAgainstLaterSessionIsInaccessible(t *testing.T) { require.ErrorIs(t, err, dbpkg.ErrTaskOwnedByAnotherUser, "bob must not be able to delete the orphaned task") } -// TestListTasksForSessionIsScopedToOwner: session ids are not globally unique -// (session's key is (id, user_id)), so listing tasks by session id alone -// would leak one user's tasks to another user holding the same session id. +// TestListTasksForSessionIsScopedToOwner: a session id is only unique among +// live sessions (session_id_active_unique), so it can still be reused by a +// different user once the original owner's session is deleted. Listing tasks +// by session id alone must not resurface the previous owner's (now +// cascade-deleted) tasks to whoever reuses the id, and writing the new +// owner's task must not touch a stale row from the old owner. func TestListTasksForSessionIsScopedToOwner(t *testing.T) { db := setupTestDB(t) client := NewClient(db) ctx := context.Background() - require.NoError(t, client.StoreSession(ctx, &dbpkg.Session{ID: "s-shared", UserID: "alice"})) - require.NoError(t, client.StoreSession(ctx, &dbpkg.Session{ID: "s-shared", UserID: "bob"})) - require.NoError(t, client.StoreTask(ctx, &a2a.Task{ID: "t-alice", ContextID: "s-shared"}, "alice")) - require.NoError(t, client.StoreTask(ctx, &a2a.Task{ID: "t-bob", ContextID: "s-shared"}, "bob")) - - bobBefore, err := client.GetSession(ctx, "s-shared", "bob") - require.NoError(t, err) - time.Sleep(10 * time.Millisecond) + require.NoError(t, client.StoreSession(ctx, &dbpkg.Session{ID: "s-reused", UserID: "alice"})) + require.NoError(t, client.StoreTask(ctx, &a2a.Task{ID: "t-alice", ContextID: "s-reused"}, "alice")) + require.NoError(t, client.DeleteSession(ctx, "s-reused", "alice"), + "deleting alice's session cascades to t-alice") - tasks, err := client.ListTasksForSession(ctx, "s-shared", "alice") - require.NoError(t, err) - require.Len(t, tasks, 1) - assert.Equal(t, a2a.TaskID("t-alice"), tasks[0].ID) + require.NoError(t, client.StoreSession(ctx, &dbpkg.Session{ID: "s-reused", UserID: "bob"})) + require.NoError(t, client.StoreTask(ctx, &a2a.Task{ID: "t-bob", ContextID: "s-reused"}, "bob")) - tasks, err = client.ListTasksForSession(ctx, "s-shared", "bob") + tasks, err := client.ListTasksForSession(ctx, "s-reused", "bob") require.NoError(t, err) - require.Len(t, tasks, 1) + require.Len(t, tasks, 1, "alice's cascade-deleted task must not resurface for bob") assert.Equal(t, a2a.TaskID("t-bob"), tasks[0].ID) - // Storing alice's task must not touch bob's same-id session. - require.NoError(t, client.StoreTask(ctx, &a2a.Task{ID: "t-alice", ContextID: "s-shared"}, "alice")) - bobAfter, err := client.GetSession(ctx, "s-shared", "bob") + // alice's session is gone; she gets nothing back for the id she used to own. + tasks, err = client.ListTasksForSession(ctx, "s-reused", "alice") require.NoError(t, err) - assert.Equal(t, bobBefore.UpdatedAt, bobAfter.UpdatedAt, - "another user's task write must not advance this session's updated_at") + assert.Empty(t, tasks) } // TestStoreAgentIdempotence verifies that calling StoreAgent multiple times 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 ba8ae5c8c..102671ded 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) @@ -74,9 +75,14 @@ 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, arg SoftDeleteTaskParams) error + // SoftDeleteTasksBySession cascades from an already owner-verified session + // delete (the caller checked GetSession(id, userID) first), so it trusts + // session_id alone and does not re-check ownership per task. + 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 bde040957..87645d3e7 100644 --- a/go/core/internal/database/gen/tasks.sql.go +++ b/go/core/internal/database/gen/tasks.sql.go @@ -128,6 +128,18 @@ func (q *Queries) SoftDeleteTask(ctx context.Context, arg SoftDeleteTaskParams) return err } +const softDeleteTasksBySession = `-- name: SoftDeleteTasksBySession :exec +UPDATE task SET deleted_at = NOW() WHERE session_id = $1 AND deleted_at IS NULL +` + +// SoftDeleteTasksBySession cascades from an already owner-verified session +// delete (the caller checked GetSession(id, userID) first), so it trusts +// session_id alone and does not re-check ownership per task. +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 14e20793f..0b156a472 100644 --- a/go/core/internal/database/queries/tasks.sql +++ b/go/core/internal/database/queries/tasks.sql @@ -78,3 +78,9 @@ WHERE task.id = $1 AND task.deleted_at IS NULL SELECT MIN(s.user_id) FROM session s WHERE s.id = task.session_id AND s.created_at <= task.created_at HAVING COUNT(DISTINCT s.user_id) = 1))); + +-- SoftDeleteTasksBySession cascades from an already owner-verified session +-- delete (the caller checked GetSession(id, userID) first), so it trusts +-- session_id alone and does not re-check ownership per task. +-- 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 15bc04a76..b664bfa7b 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/000008_session_id_unique.down.sql b/go/core/pkg/migrations/core/000008_session_id_unique.down.sql new file mode 100644 index 000000000..567e78c81 --- /dev/null +++ b/go/core/pkg/migrations/core/000008_session_id_unique.down.sql @@ -0,0 +1 @@ +DROP INDEX IF EXISTS session_id_active_unique; diff --git a/go/core/pkg/migrations/core/000008_session_id_unique.up.sql b/go/core/pkg/migrations/core/000008_session_id_unique.up.sql new file mode 100644 index 000000000..9fe2b4c4e --- /dev/null +++ b/go/core/pkg/migrations/core/000008_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;