Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions go/api/database/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,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")
Expand Down
26 changes: 24 additions & 2 deletions go/core/internal/database/client_postgres.go
Original file line number Diff line number Diff line change
Expand Up @@ -98,7 +98,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,
Expand All @@ -111,6 +111,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) {
Expand Down Expand Up @@ -161,7 +166,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 {
Comment thread
mesutoezdil marked this conversation as resolved.
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})
})
Comment thread
mesutoezdil marked this conversation as resolved.
}

// ── Session Shares ─────────────────────────────────────────────────────────────
Expand Down
65 changes: 41 additions & 24 deletions go/core/internal/database/client_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -430,9 +448,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")

Expand Down Expand Up @@ -478,39 +500,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)
}

func TestLegacyTaskProtocolVersionRejected(t *testing.T) {
Expand Down
9 changes: 9 additions & 0 deletions go/core/internal/database/gen/events.sql.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

6 changes: 6 additions & 0 deletions go/core/internal/database/gen/querier.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

9 changes: 9 additions & 0 deletions go/core/internal/database/gen/session_shares.sql.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

12 changes: 12 additions & 0 deletions go/core/internal/database/gen/tasks.sql.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

3 changes: 3 additions & 0 deletions go/core/internal/database/queries/events.sql
Original file line number Diff line number Diff line change
Expand Up @@ -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;
3 changes: 3 additions & 0 deletions go/core/internal/database/queries/session_shares.sql
Original file line number Diff line number Diff line change
Expand Up @@ -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;
6 changes: 6 additions & 0 deletions go/core/internal/database/queries/tasks.sql
Original file line number Diff line number Diff line change
Expand Up @@ -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;
3 changes: 3 additions & 0 deletions go/core/internal/service/session/service.go
Original file line number Diff line number Diff line change
Expand Up @@ -178,6 +178,9 @@ func (s *Service) Create(ctx context.Context, request CreateRequest) (*database.
Source: request.Source,
}
if err := s.store.StoreSession(ctx, value); err != nil {
if errors.Is(err, database.ErrSessionIDInUse) {
return nil, serviceerrors.NewAlreadyExists("Session ID is already in use", err)
}
return nil, serviceerrors.NewInternal("Failed to create session", err)
}
stored, err := s.store.GetSession(ctx, id, userID)
Expand Down
12 changes: 12 additions & 0 deletions go/core/internal/service/session/service_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -196,6 +196,18 @@ func TestCreateEnforcesLegacySandboxSingleSession(t *testing.T) {
}
}

func TestCreateReportsAlreadyExistsOnDuplicateSessionID(t *testing.T) {
store := newSessionTestStore()
store.agents["default__NS__agent"] = &database.Agent{ID: "default__NS__agent"}
store.storeSessionError = database.ErrSessionIDInUse
dupID := "dup"

_, err := NewService(store).Create(sessionContext("user-a", ""), CreateRequest{ID: &dupID, AgentRef: "default/agent"})
if !serviceerrors.IsCode(err, serviceerrors.CodeAlreadyExists) {
t.Fatalf("Create() error = %v, want already exists", err)
}
}

func TestGetUsesShareOwnerAndReportsReadOnly(t *testing.T) {
store := newSessionTestStore()
store.sessions["shared"] = &database.Session{ID: "shared", UserID: "owner"}
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
DROP INDEX IF EXISTS session_id_active_unique;
5 changes: 5 additions & 0 deletions go/core/pkg/migrations/core/000017_session_id_unique.up.sql
Original file line number Diff line number Diff line change
@@ -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;