From 74d0b6474c8dcd8ccb8422f82f44f9f4fa264cbd Mon Sep 17 00:00:00 2001
From: Yash Datta <yd2590@columbia.edu>
Date: Thu, 23 Jul 2026 04:33:18 +0200
Subject: [PATCH 1/4] fix: rebuild WIMSE URI when identity_type changes in
 UpdateIdentity
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The WIMSE URI embeds identity_type as a path segment; when UpdateIdentity
changed identity_type the stored URI was not rebuilt, leaving GetIdentityByWIMSEURI
broken and old JWT sub claims dangling. Rebuilds the URI in lockstep (matching
the pattern in reconcileDiscovered) and cascade-revokes active credentials with
reason identity_retyped so old tokens stop resolving.

Adds a domain unit test pinning that BuildWIMSEURI embeds identity_type as a
distinct path segment, and two integration tests covering the retype round-trip
(old URI → 404, new URI → 200) and URI-stable cases.

Fixes #239

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 domain/identity_test.go            | 36 ++++++++++++++++++++
 internal/service/identity.go       | 17 ++++++++++
 tests/integration/identity_test.go | 53 ++++++++++++++++++++++++++++++
 3 files changed, 106 insertions(+)

diff --git a/domain/identity_test.go b/domain/identity_test.go
index 770b358..52364b6 100644
--- a/domain/identity_test.go
+++ b/domain/identity_test.go
@@ -136,6 +136,42 @@ func TestBuildWIMSEURI_LengthCap(t *testing.T) {
 	}
 }
 
+// TestBuildWIMSEURI_TypeSegmentVaries pins that identity_type is embedded as a
+// path segment and differs for each type — the invariant that UpdateIdentity
+// must maintain when retyping: the same (domain, account, project, external_id)
+// produces different URIs for different identity types.
+func TestBuildWIMSEURI_TypeSegmentVaries(t *testing.T) {
+	const (
+		wimseDomain = "highflame.dev"
+		account     = "acct"
+		project     = "proj"
+		extID       = "my-identity"
+	)
+	cases := []struct {
+		typ     IdentityType
+		segment string
+	}{
+		{IdentityTypeAgent, "/agent/"},
+		{IdentityTypeApplication, "/application/"},
+		{IdentityTypeMCPServer, "/mcp_server/"},
+		{IdentityTypeService, "/service/"},
+	}
+	uris := make(map[string]bool)
+	for _, tc := range cases {
+		uri, err := BuildWIMSEURI(wimseDomain, account, project, tc.typ, extID)
+		if err != nil {
+			t.Fatalf("BuildWIMSEURI(%q) returned error: %v", tc.typ, err)
+		}
+		if !strings.Contains(uri, tc.segment) {
+			t.Errorf("URI for type %q = %q, want segment %q", tc.typ, uri, tc.segment)
+		}
+		if uris[uri] {
+			t.Errorf("duplicate URI across identity types: %q", uri)
+		}
+		uris[uri] = true
+	}
+}
+
 // TestSubTypeValidForIdentityType pins the sub_type ↔ identity_type matrix.
 // Regression: code_agent must be valid for identity_type=agent (a delegated
 // code agent, e.g. forge's per-sandbox mint) — it was previously accepted only
diff --git a/internal/service/identity.go b/internal/service/identity.go
index 36c243b..e09fb46 100644
--- a/internal/service/identity.go
+++ b/internal/service/identity.go
@@ -673,6 +673,7 @@ func (s *IdentityService) UpdateIdentity(ctx context.Context, id, accountID, pro
 	if req.TrustLevel != "" {
 		identity.TrustLevel = req.TrustLevel
 	}
+	priorType := identity.IdentityType
 	if req.IdentityType != "" {
 		if !req.IdentityType.Valid() {
 			return nil, fmt.Errorf("%w: invalid identity_type: %s", ErrInvalidIdentityField, req.IdentityType)
@@ -691,6 +692,22 @@ func (s *IdentityService) UpdateIdentity(ctx context.Context, id, accountID, pro
 		!identity.SubType.ValidForIdentityType(identity.IdentityType) {
 		return nil, fmt.Errorf("%w: invalid sub_type %q for identity_type %q", ErrInvalidIdentityField, identity.SubType, identity.IdentityType)
 	}
+	// When identity_type changes the WIMSE URI must be rebuilt in lockstep:
+	// the URI embeds identity_type as a path segment, so a stale URI breaks
+	// GetIdentityByWIMSEURI and leaves dangling JWT sub claims. Cascade-revoke
+	// active credentials so old tokens (which carry the prior URI as sub) stop
+	// resolving — same sweep as deactivation, with a distinct audit reason.
+	if identity.IdentityType != priorType {
+		newURI, err := domain.BuildWIMSEURI(
+			s.wimseDomain, identity.AccountID, identity.ProjectID,
+			identity.IdentityType, identity.ExternalID,
+		)
+		if err != nil {
+			return nil, err
+		}
+		identity.WIMSEURI = newURI
+		s.runDeactivationCleanup(ctx, identity, "identity_retyped")
+	}
 	if req.OwnerUserID != "" {
 		identity.OwnerUserID = req.OwnerUserID
 	}
diff --git a/tests/integration/identity_test.go b/tests/integration/identity_test.go
index ec399b4..df6bfb1 100644
--- a/tests/integration/identity_test.go
+++ b/tests/integration/identity_test.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"net/http"
+	"net/url"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -634,3 +635,55 @@ func TestUpdateIdentityRiskMetadata(t *testing.T) {
 	assert.Equal(t, "high", reclassified["risk_tier"], "risk_tier should be unchanged by capability_tier update")
 	assert.Equal(t, "ial3", reclassified["ial"], "ial should be unchanged by capability_tier update")
 }
+
+// TestUpdateIdentity_TypeChange_WIMSEURIRebuild verifies that PATCH
+// /identities/{id} with a new identity_type rebuilds the WIMSE URI in
+// lockstep. The old URI must no longer resolve (404) and the new URI must
+// resolve (200) — otherwise GetIdentityByWIMSEURI and JWT sub-claim lookups
+// break silently after a retype.
+func TestUpdateIdentity_TypeChange_WIMSEURIRebuild(t *testing.T) {
+	externalID := uid("retype-agent")
+	identity := registerIdentity(t, externalID, nil)
+
+	oldURI := identity.WIMSEURI
+	require.Contains(t, oldURI, "/agent/", "initial URI must embed identity_type=agent")
+
+	// Retype agent → application (supply a compatible sub_type).
+	resp := doRequest(t, http.MethodPatch, adminPath("/identities/"+identity.ID), map[string]any{
+		"identity_type": "application",
+		"sub_type":      "api_service",
+	}, adminHeaders())
+	require.Equal(t, http.StatusOK, resp.StatusCode)
+	body := decode(t, resp)
+
+	newURI, _ := body["wimse_uri"].(string)
+	assert.Contains(t, newURI, "/application/", "WIMSE URI must be rebuilt to reflect the new identity_type")
+	assert.NotContains(t, newURI, "/agent/", "old identity_type segment must not appear in the rebuilt URI")
+	assert.NotEqual(t, oldURI, newURI, "URI must change when identity_type changes")
+
+	// New URI resolves; stale URI does not.
+	newResp := get(t, adminPath("/identities/by-wimse")+"?uri="+url.QueryEscape(newURI), adminHeaders())
+	assert.Equal(t, http.StatusOK, newResp.StatusCode, "new wimse_uri must resolve after retype")
+	_ = newResp.Body.Close()
+
+	oldResp := get(t, adminPath("/identities/by-wimse")+"?uri="+url.QueryEscape(oldURI), adminHeaders())
+	assert.Equal(t, http.StatusNotFound, oldResp.StatusCode, "stale wimse_uri must no longer resolve after retype")
+	_ = oldResp.Body.Close()
+}
+
+// TestUpdateIdentity_TypeUnchanged_URIStable verifies that an update that
+// does not touch identity_type leaves the WIMSE URI unchanged — the rebuild
+// must be a no-op when the type is the same.
+func TestUpdateIdentity_TypeUnchanged_URIStable(t *testing.T) {
+	externalID := uid("stable-uri-agent")
+	identity := registerIdentity(t, externalID, nil)
+	originalURI := identity.WIMSEURI
+
+	resp := doRequest(t, http.MethodPatch, adminPath("/identities/"+identity.ID), map[string]any{
+		"name": "renamed-agent",
+	}, adminHeaders())
+	require.Equal(t, http.StatusOK, resp.StatusCode)
+	body := decode(t, resp)
+
+	assert.Equal(t, originalURI, body["wimse_uri"], "URI must not change when identity_type is not updated")
+}

From 3bdca5e1a91a261243fd16d455c50679beec63e0 Mon Sep 17 00:00:00 2001
From: Yash Datta <yd2590@columbia.edu>
Date: Thu, 23 Jul 2026 04:35:19 +0200
Subject: [PATCH 2/4] fix: move identity_retyped cleanup to after repo.Update
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

runDeactivationCleanup must run after the DB write succeeds — if repo.Update
fails, calling cleanup first irreversibly revokes credentials while the
identity row remains unchanged, leaving it permanently broken.

Follows the same post-persist ordering used for the deactivated and expired
cleanup paths at the bottom of UpdateIdentity.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 internal/service/identity.go | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/internal/service/identity.go b/internal/service/identity.go
index e09fb46..6805328 100644
--- a/internal/service/identity.go
+++ b/internal/service/identity.go
@@ -694,9 +694,7 @@ func (s *IdentityService) UpdateIdentity(ctx context.Context, id, accountID, pro
 	}
 	// When identity_type changes the WIMSE URI must be rebuilt in lockstep:
 	// the URI embeds identity_type as a path segment, so a stale URI breaks
-	// GetIdentityByWIMSEURI and leaves dangling JWT sub claims. Cascade-revoke
-	// active credentials so old tokens (which carry the prior URI as sub) stop
-	// resolving — same sweep as deactivation, with a distinct audit reason.
+	// GetIdentityByWIMSEURI and leaves dangling JWT sub claims.
 	if identity.IdentityType != priorType {
 		newURI, err := domain.BuildWIMSEURI(
 			s.wimseDomain, identity.AccountID, identity.ProjectID,
@@ -706,7 +704,6 @@ func (s *IdentityService) UpdateIdentity(ctx context.Context, id, accountID, pro
 			return nil, err
 		}
 		identity.WIMSEURI = newURI
-		s.runDeactivationCleanup(ctx, identity, "identity_retyped")
 	}
 	if req.OwnerUserID != "" {
 		identity.OwnerUserID = req.OwnerUserID
@@ -807,6 +804,16 @@ func (s *IdentityService) UpdateIdentity(ctx context.Context, id, accountID, pro
 		return nil, err
 	}
 
+	// Post-persist cleanups — all run after repo.Update succeeds so a DB
+	// failure cannot leave credentials revoked while the identity row is
+	// unchanged (the authoritative outcome must land before the sweep).
+
+	// Retype: cascade-revoke active credentials so old tokens (which carry
+	// the prior URI as sub) stop resolving — same sweep as deactivation,
+	// with a distinct audit reason.
+	if identity.IdentityType != priorType {
+		s.runDeactivationCleanup(ctx, identity, "identity_retyped")
+	}
 	// Fresh transition into deactivated or expired: sweep linked API keys,
 	// cascade-revoke active credentials, and emit a retirement/expiry signal.
 	// Centralized here so every update path (PUT /identities/{id},

From 75da38460960fa79fca6e75b68328e2675c23f37 Mon Sep 17 00:00:00 2001
From: Yash Datta <yd2590@columbia.edu>
Date: Thu, 23 Jul 2026 04:48:23 +0200
Subject: [PATCH 3/4] fix: bump google.golang.org/grpc to v1.82.1
 (GHSA-hrxh-6v49-42gf HIGH)

---
 go.mod |  6 +++---
 go.sum | 12 ++++++------
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/go.mod b/go.mod
index f0f7d76..6e11d39 100644
--- a/go.mod
+++ b/go.mod
@@ -98,9 +98,9 @@ require (
 	golang.org/x/net v0.55.0 // indirect
 	golang.org/x/sys v0.45.0 // indirect
 	golang.org/x/text v0.37.0 // indirect
-	google.golang.org/genproto/googleapis/api v0.0.0-20260401024825-9d38bb4040a9 // indirect
-	google.golang.org/genproto/googleapis/rpc v0.0.0-20260401024825-9d38bb4040a9 // indirect
-	google.golang.org/grpc v1.80.0 // indirect
+	google.golang.org/genproto/googleapis/api v0.0.0-20260414002931-afd174a4e478 // indirect
+	google.golang.org/genproto/googleapis/rpc v0.0.0-20260414002931-afd174a4e478 // indirect
+	google.golang.org/grpc v1.82.1 // indirect
 	google.golang.org/protobuf v1.36.11 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 	mellium.im/sasl v0.3.2 // indirect
diff --git a/go.sum b/go.sum
index d7b4754..8440ab4 100644
--- a/go.sum
+++ b/go.sum
@@ -231,12 +231,12 @@ golang.org/x/text v0.37.0/go.mod h1:a5sjxXGs9hsn/AJVwuElvCAo9v8QYLzvavO5z2PiM38=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 gonum.org/v1/gonum v0.17.0 h1:VbpOemQlsSMrYmn7T2OUvQ4dqxQXU+ouZFQsZOx50z4=
 gonum.org/v1/gonum v0.17.0/go.mod h1:El3tOrEuMpv2UdMrbNlKEh9vd86bmQ6vqIcDwxEOc1E=
-google.golang.org/genproto/googleapis/api v0.0.0-20260401024825-9d38bb4040a9 h1:VPWxll4HlMw1Vs/qXtN7BvhZqsS9cdAittCNvVENElA=
-google.golang.org/genproto/googleapis/api v0.0.0-20260401024825-9d38bb4040a9/go.mod h1:7QBABkRtR8z+TEnmXTqIqwJLlzrZKVfAUm7tY3yGv0M=
-google.golang.org/genproto/googleapis/rpc v0.0.0-20260401024825-9d38bb4040a9 h1:m8qni9SQFH0tJc1X0vmnpw/0t+AImlSvp30sEupozUg=
-google.golang.org/genproto/googleapis/rpc v0.0.0-20260401024825-9d38bb4040a9/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8=
-google.golang.org/grpc v1.80.0 h1:Xr6m2WmWZLETvUNvIUmeD5OAagMw3FiKmMlTdViWsHM=
-google.golang.org/grpc v1.80.0/go.mod h1:ho/dLnxwi3EDJA4Zghp7k2Ec1+c2jqup0bFkw07bwF4=
+google.golang.org/genproto/googleapis/api v0.0.0-20260414002931-afd174a4e478 h1:yQugLulqltosq0B/f8l4w9VryjV+N/5gcW0jQ3N8Qec=
+google.golang.org/genproto/googleapis/api v0.0.0-20260414002931-afd174a4e478/go.mod h1:C6ADNqOxbgdUUeRTU+LCHDPB9ttAMCTff6auwCVa4uc=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20260414002931-afd174a4e478 h1:RmoJA1ujG+/lRGNfUnOMfhCy5EipVMyvUE+KNbPbTlw=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20260414002931-afd174a4e478/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8=
+google.golang.org/grpc v1.82.1 h1:NnAxzGRA0677vCa4BUkOAnO5+FfQqVl9iUXeD0IqcGE=
+google.golang.org/grpc v1.82.1/go.mod h1:yzTZ1TB1Z3SG+LIYaI+WiE8D5+PZ3ArnrSp8zF3+/ZA=
 google.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBNVE=
 google.golang.org/protobuf v1.36.11/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=

From ce2786c04df66dad05a88d5b78b1a680901d32c3 Mon Sep 17 00:00:00 2001
From: Yash Datta <yd2590@columbia.edu>
Date: Thu, 23 Jul 2026 05:05:26 +0200
Subject: [PATCH 4/4] fix: align make test with CI unit-test step (exclude
 integration, 300s timeout)

make test was running ./... including tests/integration/ with a 120s timeout,
but testcontainers startup alone takes ~60s, causing the integration binary to
consistently hit the timeout. CI runs unit and integration tests as separate
jobs with 300s timeouts.

Update make test to match CI: exclude ./tests/... (identical to the
go list grep in pr-check.yml) and raise the timeout to 300s. Add
make test-all for the combined suite and make test-integration now also
uses 300s to accommodate container startup.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 Makefile | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/Makefile b/Makefile
index 384a656..dfdf240 100644
--- a/Makefile
+++ b/Makefile
@@ -18,11 +18,14 @@ build: ## Build the zeroid binary
 run: build ## Build and run zeroid locally
 	./$(BINARY) -config zeroid.yaml
 
-test: ## Run all tests (unit + integration)
-	go test ./... -v -race -count=1 -timeout=120s
+test: ## Run unit tests (matches CI unit-test step; integration requires Docker — use make test-integration)
+	go test $$(go list ./... | grep -v '/tests') -v -race -count=1 -timeout=300s
 
-test-integration: ## Run integration tests only (requires Docker)
-	go test ./tests/integration/ -v -count=1 -timeout=120s
+test-integration: ## Run integration tests only (requires Docker, testcontainers)
+	go test ./tests/integration/ -v -count=1 -timeout=300s
+
+test-all: ## Run unit + integration tests (requires Docker)
+	go test ./... -v -race -count=1 -timeout=300s
 
 lint: ## Run go vet
 	go vet ./...