diff --git a/packages/api/internal/handlers/sandbox.go b/packages/api/internal/handlers/sandbox.go index 4c2ca21bc6..539acf8128 100644 --- a/packages/api/internal/handlers/sandbox.go +++ b/packages/api/internal/handlers/sandbox.go @@ -154,7 +154,7 @@ func (a *APIStore) startSandboxInternal( zap.String("end_time", endTime.Format("2006-01-02 15:04:05 -07:00")), zap.String("auto_resume_policy", autoResumePolicy), zap.Bool("auto_pause", sbx.AutoPause), - zap.String("parent_template_id", sbx.BaseTemplateID), + zap.String("template_id", sbx.BaseTemplateID), ) return sbx, nil diff --git a/packages/api/internal/handlers/sandbox_get.go b/packages/api/internal/handlers/sandbox_get.go index e052f423d4..e32f67be3e 100644 --- a/packages/api/internal/handlers/sandbox_get.go +++ b/packages/api/internal/handlers/sandbox_get.go @@ -109,7 +109,7 @@ func (a *APIStore) GetSandboxesSandboxID(c *gin.Context, id string) { // Sandbox exists and belongs to the team - return running sandbox sbx sandbox := api.SandboxDetail{ ClientID: sbx.ClientID, - TemplateID: sbx.TemplateID, + TemplateID: sbx.BaseTemplateID, Alias: sbx.Alias, SandboxID: sbx.SandboxID, StartedAt: sbx.StartTime, @@ -202,7 +202,7 @@ func (a *APIStore) GetSandboxesSandboxID(c *gin.Context, id string) { sandbox := api.SandboxDetail{ ClientID: consts.ClientID, // for backwards compatibility we need to return a client id - TemplateID: lastSnapshot.Snapshot.EnvID, + TemplateID: lastSnapshot.Snapshot.BaseEnvID, SandboxID: lastSnapshot.Snapshot.SandboxID, StartedAt: lastSnapshot.Snapshot.SandboxStartedAt.Time, CpuCount: cpuCount, diff --git a/packages/api/internal/orchestrator/create_instance_test.go b/packages/api/internal/orchestrator/create_instance_test.go index 641ade070f..7ab2112586 100644 --- a/packages/api/internal/orchestrator/create_instance_test.go +++ b/packages/api/internal/orchestrator/create_instance_test.go @@ -147,6 +147,7 @@ func TestCreateSandbox_StaleDataAfterConcurrentPause(t *testing.T) { ) require.Nil(t, apiErr) assert.Equal(t, "tpl-v1", sbx1.TemplateID) + assert.Equal(t, "base-tpl", sbx1.BaseTemplateID) // Clean up reservation. o.sandboxStore.Remove(t.Context(), team.Team.ID, sandboxID) @@ -172,6 +173,8 @@ func TestCreateSandbox_StaleDataAfterConcurrentPause(t *testing.T) { // The sandbox SHOULD have been created with V2 (fresh) data. assert.Equal(t, "tpl-v2", sbx2.TemplateID, "CreateSandbox must use the latest snapshot data, not stale pre-lock values") + assert.Equal(t, "base-tpl", sbx2.BaseTemplateID, + "CreateSandbox must preserve the base template ID") assert.Equal(t, "v2", sbx2.Metadata["snapshot"], "CreateSandbox must use the latest metadata, not stale pre-lock values") }) diff --git a/packages/api/internal/sandbox/sandbox.go b/packages/api/internal/sandbox/sandbox.go index 48dcd53ff1..f3a4b3456f 100644 --- a/packages/api/internal/sandbox/sandbox.go +++ b/packages/api/internal/sandbox/sandbox.go @@ -111,7 +111,7 @@ type Sandbox struct { func (s Sandbox) ToAPISandbox() *api.Sandbox { return &api.Sandbox{ SandboxID: s.SandboxID, - TemplateID: s.TemplateID, + TemplateID: s.BaseTemplateID, ClientID: s.ClientID, Alias: s.Alias, EnvdVersion: s.EnvdVersion, diff --git a/packages/db/queries/get_sandbox_record.sql.go b/packages/db/queries/get_sandbox_record.sql.go index 891331d970..ae419e73fa 100644 --- a/packages/db/queries/get_sandbox_record.sql.go +++ b/packages/db/queries/get_sandbox_record.sql.go @@ -15,7 +15,7 @@ import ( const getSandboxRecordByTeamAndSandboxID = `-- name: GetSandboxRecordByTeamAndSandboxID :one SELECT sl.sandbox_id, - sl.env_id AS template_id, + COALESCE(s.base_env_id, sl.env_id) AS template_id, sl.vcpu, sl.ram_mb, sl.total_disk_size_mb, @@ -26,10 +26,13 @@ SELECT FROM billing.sandbox_logs sl LEFT JOIN public.teams t ON t.id = sl.team_id LEFT JOIN public.clusters c ON c.id = t.cluster_id +LEFT JOIN public.snapshots s + ON s.sandbox_id = sl.sandbox_id + AND s.team_id = sl.team_id LEFT JOIN LATERAL ( SELECT ea.alias FROM public.env_aliases ea - WHERE ea.env_id = sl.env_id + WHERE ea.env_id = COALESCE(s.base_env_id, sl.env_id) ORDER BY ea.alias LIMIT 1 ) template_alias ON TRUE diff --git a/packages/db/queries/sandboxes/get_sandbox_record.sql b/packages/db/queries/sandboxes/get_sandbox_record.sql index 15849be542..5785af32e7 100644 --- a/packages/db/queries/sandboxes/get_sandbox_record.sql +++ b/packages/db/queries/sandboxes/get_sandbox_record.sql @@ -1,7 +1,7 @@ -- name: GetSandboxRecordByTeamAndSandboxID :one SELECT sl.sandbox_id, - sl.env_id AS template_id, + COALESCE(s.base_env_id, sl.env_id) AS template_id, sl.vcpu, sl.ram_mb, sl.total_disk_size_mb, @@ -12,10 +12,13 @@ SELECT FROM billing.sandbox_logs sl LEFT JOIN public.teams t ON t.id = sl.team_id LEFT JOIN public.clusters c ON c.id = t.cluster_id +LEFT JOIN public.snapshots s + ON s.sandbox_id = sl.sandbox_id + AND s.team_id = sl.team_id LEFT JOIN LATERAL ( SELECT ea.alias FROM public.env_aliases ea - WHERE ea.env_id = sl.env_id + WHERE ea.env_id = COALESCE(s.base_env_id, sl.env_id) ORDER BY ea.alias LIMIT 1 ) template_alias ON TRUE diff --git a/tests/integration/internal/tests/api/sandboxes/sandbox_connect_test.go b/tests/integration/internal/tests/api/sandboxes/sandbox_connect_test.go index 0a9a720e58..116ad388e8 100644 --- a/tests/integration/internal/tests/api/sandboxes/sandbox_connect_test.go +++ b/tests/integration/internal/tests/api/sandboxes/sandbox_connect_test.go @@ -33,6 +33,7 @@ func TestSandboxConnect(t *testing.T) { require.Equal(t, http.StatusCreated, sbxConnect.StatusCode()) require.NotNil(t, sbxConnect.JSON201) assert.Equal(t, sbxConnect.JSON201.SandboxID, sbxId) + assert.Equal(t, sbx.TemplateID, sbxConnect.JSON201.TemplateID) // Check if the sandbox is running res, err := c.GetSandboxesSandboxIDWithResponse(t.Context(), sbxId, setup.WithAPIKey()) @@ -40,6 +41,7 @@ func TestSandboxConnect(t *testing.T) { require.Equal(t, http.StatusOK, res.StatusCode()) require.NotNil(t, res.JSON200) assert.Equal(t, api.Running, res.JSON200.State) + assert.Equal(t, sbx.TemplateID, res.JSON200.TemplateID) }) t.Run("connect to running sandbox", func(t *testing.T) { diff --git a/tests/integration/internal/tests/api/sandboxes/sandbox_detail_test.go b/tests/integration/internal/tests/api/sandboxes/sandbox_detail_test.go index 71ea9640b7..2ce712eaa7 100644 --- a/tests/integration/internal/tests/api/sandboxes/sandbox_detail_test.go +++ b/tests/integration/internal/tests/api/sandboxes/sandbox_detail_test.go @@ -62,6 +62,7 @@ func TestSandboxDetailReturnsLifecycleAndNetworkConfig(t *testing.T) { returnedSbx := response.JSON200 assert.Equal(t, sbx.SandboxID, returnedSbx.SandboxID) + assert.Equal(t, sbx.TemplateID, returnedSbx.TemplateID) assert.Equal(t, expectedState, returnedSbx.State) require.NotNil(t, returnedSbx.AllowInternetAccess) @@ -98,6 +99,7 @@ func TestSandboxDetailPaused(t *testing.T) { require.Equal(t, http.StatusOK, response.StatusCode()) returnedSbx := response.JSON200 assert.Equal(t, sbx.SandboxID, returnedSbx.SandboxID) + assert.Equal(t, sbx.TemplateID, returnedSbx.TemplateID) } func TestSandboxDetailPausingSandbox(t *testing.T) { diff --git a/tests/integration/internal/tests/api/sandboxes/sandbox_resume_test.go b/tests/integration/internal/tests/api/sandboxes/sandbox_resume_test.go index dde4689af6..f1550cd4f4 100644 --- a/tests/integration/internal/tests/api/sandboxes/sandbox_resume_test.go +++ b/tests/integration/internal/tests/api/sandboxes/sandbox_resume_test.go @@ -42,6 +42,7 @@ func TestSandboxResume(t *testing.T) { require.Equal(t, http.StatusCreated, sbxResume.StatusCode()) require.NotNil(t, sbxResume.JSON201) assert.Equal(t, sbxResume.JSON201.SandboxID, sbxId) + assert.Equal(t, sbx.TemplateID, sbxResume.JSON201.TemplateID) }) t.Run("concurrent resumes", func(t *testing.T) { @@ -76,6 +77,7 @@ func TestSandboxResume(t *testing.T) { require.Equal(t, http.StatusOK, res.StatusCode()) require.NotNil(t, res.JSON200) assert.Equal(t, api.Running, res.JSON200.State) + assert.Equal(t, sbx.TemplateID, res.JSON200.TemplateID) assert.True(t, resumed.Load(), "at least one resume should succeed") }) diff --git a/tests/integration/internal/tests/api/sandboxes/sandbox_test.go b/tests/integration/internal/tests/api/sandboxes/sandbox_test.go index 69d06a1264..69d7f86351 100644 --- a/tests/integration/internal/tests/api/sandboxes/sandbox_test.go +++ b/tests/integration/internal/tests/api/sandboxes/sandbox_test.go @@ -38,6 +38,8 @@ func TestSandboxCreate(t *testing.T) { }) assert.Equal(t, http.StatusCreated, resp.StatusCode()) + require.NotNil(t, resp.JSON201) + assert.Equal(t, setup.SandboxTemplateID, resp.JSON201.TemplateID) } func TestSandboxResumeUnknownSandbox(t *testing.T) { @@ -96,6 +98,7 @@ func TestSandboxResumeWithSecuredEnvd(t *testing.T) { }) assert.Equal(t, sbxResume.JSON201.SandboxID, sbxCreate.JSON201.SandboxID) + assert.Equal(t, sbxCreate.JSON201.TemplateID, sbxResume.JSON201.TemplateID) assert.Equal(t, sbxResume.JSON201.EnvdAccessToken, sbxCreate.JSON201.EnvdAccessToken) }