package agent

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"sync/atomic"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/drivercontract"
	"github.com/tidbcloud/AgentTeam/desktop/internal/consts"
	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
)

func bindCodexPrimaryActiveThreadForIsolationTest(driver *CodexAppServerDriver) {
	driver.updateThreadSnapshot(appServerThread{
		ID:     "thread-primary",
		Status: appServerThreadStatus{Type: "active"},
		Turns: []appServerTurn{{
			ID:     "turn-primary",
			Status: appServerTurnStatus{Type: "inProgress"},
		}},
	})
}

func TestCodexAppServerForeignThreadNotificationsDoNotAffectPrimary(t *testing.T) {
	driver, _ := newTestCodexDriver(t, nil)
	driver.translatedLines = make(chan string, 32)
	bindCodexPrimaryActiveThreadForIsolationTest(driver)

	notifications := []appServerEnvelope{
		{
			Method: "thread/status/changed",
			Params: rawRPCResult(t, appServerThreadStatusNotification{
				ThreadID: "thread-child",
				Status:   appServerThreadStatus{Type: "idle"},
			}),
		},
		{
			Method: "turn/started",
			Params: rawRPCResult(t, appServerTurnNotification{
				ThreadID: "thread-child",
				Turn: appServerTurn{
					ID:     "turn-child",
					Status: appServerTurnStatus{Type: "inProgress"},
				},
			}),
		},
		{
			Method: "item/reasoning/summaryTextDelta",
			Params: rawRPCResult(t, map[string]any{
				"threadId": "thread-child",
				"turnId":   "turn-child",
				"itemId":   "reasoning-child",
				"delta":    "child reasoning",
			}),
		},
		{
			Method: "item/agentMessage/delta",
			Params: rawRPCResult(t, appServerItemDeltaNotification{
				ThreadID: "thread-child",
				TurnID:   "turn-child",
				ItemID:   "message-child",
				Delta:    "child text",
			}),
		},
		{
			Method: "item/plan/delta",
			Params: rawRPCResult(t, appServerItemDeltaNotification{
				ThreadID: "thread-child",
				TurnID:   "turn-child",
				ItemID:   "plan-child",
				Delta:    "child plan",
			}),
		},
		{
			Method: "item/commandExecution/outputDelta",
			Params: rawRPCResult(t, appServerItemDeltaNotification{
				ThreadID: "thread-child",
				TurnID:   "turn-child",
				ItemID:   "command-child",
				Delta:    "child command output",
			}),
		},
		{
			Method: "item/started",
			Params: rawRPCResult(t, appServerItemNotification{
				ThreadID: "thread-child",
				TurnID:   "turn-child",
				Item: appServerItem{
					Type:     "userMessage",
					ID:       "message-child-input",
					ClientID: "primary-runtime-input-id",
				},
			}),
		},
		{
			Method: "item/started",
			Params: rawRPCResult(t, appServerItemNotification{
				ThreadID: "thread-child",
				TurnID:   "turn-child",
				Item: appServerItem{
					Type:    "commandExecution",
					ID:      "command-child",
					Command: "go test ./...",
				},
			}),
		},
		{
			Method: "item/started",
			Params: rawRPCResult(t, appServerItemNotification{
				ThreadID: "thread-child",
				TurnID:   "turn-child",
				Item: appServerItem{
					Type:           "collabAgentToolCall",
					ID:             "collab-child",
					Tool:           "wait",
					Status:         "inProgress",
					SenderThreadID: "thread-child",
				},
			}),
		},
		{
			Method: "item/completed",
			Params: rawRPCResult(t, appServerItemNotification{
				ThreadID: "thread-child",
				TurnID:   "turn-child",
				Item: appServerItem{
					Type:          "subAgentActivity",
					ID:            "activity-child",
					Kind:          "started",
					AgentThreadID: "thread-grandchild",
					AgentPath:     "/root/child/grandchild",
				},
			}),
		},
		{
			Method: consts.ParsedEventKindError,
			Params: rawRPCResult(t, map[string]any{
				"threadId":  "thread-child",
				"turnId":    "turn-child",
				"willRetry": false,
				"error": map[string]any{
					"message":        "child failed",
					"codexErrorInfo": "contextWindowExceeded",
				},
			}),
		},
		{
			Method: "turn/completed",
			Params: rawRPCResult(t, appServerTurnNotification{
				ThreadID: "thread-child",
				Turn: appServerTurn{
					ID:     "turn-child",
					Status: appServerTurnStatus{Type: "completed"},
				},
			}),
		},
	}

	for _, notification := range notifications {
		driver.handleNotification(notification)
	}

	status, activeTurnID, ok := driver.currentThreadState("thread-primary")
	require.True(t, ok)
	require.Equal(t, appServerThreadStatus{Type: "active"}, status)
	require.Equal(t, "turn-primary", activeTurnID)
	require.Equal(t, 0, driver.currentCompletedTurnCount("thread-primary"))
	driver.mu.Lock()
	require.True(t, driver.canSteerLocked("thread-primary"))
	driver.mu.Unlock()
	require.Empty(t, driver.translatedLines)

	driver.mu.Lock()
	require.False(t, driver.compactQueued)
	require.False(t, driver.compactionInProgress)
	require.Empty(t, driver.compactionTurnID)
	driver.mu.Unlock()
}

func TestCodexAppServerStalePrimaryCompletionAndIdleDoNotEndCurrentTurn(t *testing.T) {
	driver, _ := newTestCodexDriver(t, nil)
	driver.translatedLines = make(chan string, 4)
	bindCodexPrimaryActiveThreadForIsolationTest(driver)

	driver.handleNotification(appServerEnvelope{
		Method: "turn/completed",
		Params: rawRPCResult(t, appServerTurnNotification{
			ThreadID: "thread-primary",
			Turn: appServerTurn{
				ID:     "turn-stale",
				Status: appServerTurnStatus{Type: "completed"},
			},
		}),
	})
	driver.handleNotification(appServerEnvelope{
		Method: "thread/status/changed",
		Params: rawRPCResult(t, appServerThreadStatusNotification{
			ThreadID: "thread-primary",
			Status:   appServerThreadStatus{Type: "idle"},
		}),
	})

	status, activeTurnID, ok := driver.currentThreadState("thread-primary")
	require.True(t, ok)
	require.Equal(t, appServerThreadStatus{Type: "active"}, status)
	require.Equal(t, "turn-primary", activeTurnID)
	driver.mu.Lock()
	require.True(t, driver.canSteerLocked("thread-primary"))
	driver.mu.Unlock()
	require.Empty(t, driver.translatedLines, "stale completion/status must not emit turn_end or runtime_idle")

	driver.handleNotification(appServerEnvelope{
		Method: "turn/completed",
		Params: rawRPCResult(t, appServerTurnNotification{
			ThreadID: "thread-primary",
			Turn: appServerTurn{
				ID:     "turn-primary",
				Status: appServerTurnStatus{Type: "completed"},
			},
		}),
	})

	status, activeTurnID, ok = driver.currentThreadState("thread-primary")
	require.True(t, ok)
	require.Equal(t, appServerThreadStatus{Type: "idle"}, status)
	require.Empty(t, activeTurnID)
	events := driver.ParseLine(<-driver.translatedLines)
	require.Equal(t, consts.ParsedEventKindTurnEnd, events[0].Kind)
}

func TestCodexAppServerUsageIsDeduplicatedPerThreadWithoutReplacingPrimaryState(t *testing.T) {
	driver, _ := newTestCodexDriver(t, nil)
	driver.translatedLines = make(chan string, 8)
	bindCodexPrimaryActiveThreadForIsolationTest(driver)
	driver.updateEffectiveModel("gpt-5.2")
	driver.handleNotification(appServerEnvelope{
		Method: "item/completed",
		Params: rawRPCResult(t, appServerItemNotification{
			ThreadID: "thread-primary",
			TurnID:   "turn-primary",
			Item: appServerItem{
				Type:              "collabAgentToolCall",
				ID:                "spawn-child",
				Tool:              "spawnAgent",
				Status:            "completed",
				Model:             "gpt-5-mini",
				ReceiverThreadIDs: []string{"thread-child"},
			},
		}),
	})
	spawnEvents := driver.ParseLine(<-driver.translatedLines)
	require.Len(t, spawnEvents, 1)
	require.Equal(t, consts.ParsedEventKindToolOutput, spawnEvents[0].Kind)

	primaryUsage := appServerThreadTokenUsage{
		Last:  appServerTokenUsageBreakdown{InputTokens: 10, OutputTokens: 2, TotalTokens: 12},
		Total: appServerTokenUsageBreakdown{InputTokens: 10, OutputTokens: 2, TotalTokens: 12},
	}
	childUsage := appServerThreadTokenUsage{
		Last:  appServerTokenUsageBreakdown{InputTokens: 20, OutputTokens: 4, TotalTokens: 24},
		Total: appServerTokenUsageBreakdown{InputTokens: 20, OutputTokens: 4, TotalTokens: 24},
	}
	primaryUsageAdvanced := appServerThreadTokenUsage{
		Last:  appServerTokenUsageBreakdown{InputTokens: 3, OutputTokens: 1, TotalTokens: 4},
		Total: appServerTokenUsageBreakdown{InputTokens: 13, OutputTokens: 3, TotalTokens: 16},
	}

	emitUsage := func(threadID string, usage appServerThreadTokenUsage) {
		driver.handleNotification(appServerEnvelope{
			Method: "thread/tokenUsage/updated",
			Params: rawRPCResult(t, appServerThreadTokenUsageNotification{
				ThreadID:   threadID,
				TurnID:     "turn-usage",
				TokenUsage: usage,
			}),
		})
	}

	emitUsage("thread-primary", primaryUsage)
	emitUsage("thread-child", childUsage)
	emitUsage("thread-primary", primaryUsage)
	emitUsage("thread-child", childUsage)
	emitUsage("thread-primary", primaryUsageAdvanced)

	require.Len(t, driver.translatedLines, 3)
	var gotInputTokens []uint64
	var gotSessionIDs []string
	var gotModels []string
	for range 3 {
		events := driver.ParseLine(<-driver.translatedLines)
		require.Len(t, events, 1)
		require.Equal(t, consts.ParsedEventKindUsage, events[0].Kind)
		gotInputTokens = append(gotInputTokens, events[0].InputTokens)
		gotSessionIDs = append(gotSessionIDs, events[0].SessionID)
		gotModels = append(gotModels, events[0].Model)
	}
	require.Equal(t, []uint64{10, 20, 3}, gotInputTokens)
	require.Equal(t, []string{"thread-primary", "thread-child", "thread-primary"}, gotSessionIDs)
	require.Equal(t, []string{"gpt-5.2", "gpt-5-mini", "gpt-5.2"}, gotModels)

	status, activeTurnID, ok := driver.currentThreadState("thread-primary")
	require.True(t, ok)
	require.Equal(t, appServerThreadStatus{Type: "active"}, status)
	require.Equal(t, "turn-primary", activeTurnID)

	driver.mu.Lock()
	require.Equal(t, primaryUsageAdvanced, driver.lastTokenUsage)
	require.Equal(t, primaryUsageAdvanced, driver.lastTokenUsageByThread["thread-primary"])
	require.Equal(t, childUsage, driver.lastTokenUsageByThread["thread-child"])
	require.Equal(t, "gpt-5.2", driver.modelByThread["thread-primary"])
	require.Equal(t, "gpt-5-mini", driver.modelByThread["thread-child"])
	driver.mu.Unlock()
}

func TestCodexAppServerThreadModelNotificationsAttributeChildUsage(t *testing.T) {
	driver, _ := newTestCodexDriver(t, nil)
	driver.translatedLines = make(chan string, 4)
	bindCodexPrimaryActiveThreadForIsolationTest(driver)
	driver.updateEffectiveModel("gpt-parent")

	driver.handleNotification(appServerEnvelope{
		Method: "thread/settings/updated",
		Params: rawRPCResult(t, map[string]any{
			"threadId": "thread-child",
			"threadSettings": map[string]any{
				"model": "gpt-child",
			},
		}),
	})
	driver.handleNotification(appServerEnvelope{
		Method: "model/rerouted",
		Params: rawRPCResult(t, map[string]any{
			"threadId":  "thread-child",
			"turnId":    "turn-child",
			"fromModel": "gpt-child",
			"toModel":   "gpt-child-rerouted",
			"reason":    "fallback",
		}),
	})
	driver.handleNotification(appServerEnvelope{
		Method: "thread/tokenUsage/updated",
		Params: rawRPCResult(t, appServerThreadTokenUsageNotification{
			ThreadID: "thread-child",
			TurnID:   "turn-child",
			TokenUsage: appServerThreadTokenUsage{
				Last:  appServerTokenUsageBreakdown{InputTokens: 7, OutputTokens: 2, TotalTokens: 9},
				Total: appServerTokenUsageBreakdown{InputTokens: 7, OutputTokens: 2, TotalTokens: 9},
			},
		}),
	})
	driver.handleNotification(appServerEnvelope{
		Method: "turn/completed",
		Params: rawRPCResult(t, appServerTurnNotification{
			ThreadID: "thread-child",
			Turn: appServerTurn{
				ID:     "turn-child",
				Status: appServerTurnStatus{Type: "completed"},
			},
		}),
	})
	driver.handleNotification(appServerEnvelope{
		Method: "thread/tokenUsage/updated",
		Params: rawRPCResult(t, appServerThreadTokenUsageNotification{
			ThreadID: "thread-child",
			TurnID:   "turn-child-2",
			TokenUsage: appServerThreadTokenUsage{
				Last:  appServerTokenUsageBreakdown{InputTokens: 3, OutputTokens: 1, TotalTokens: 4},
				Total: appServerTokenUsageBreakdown{InputTokens: 10, OutputTokens: 3, TotalTokens: 13},
			},
		}),
	})

	require.Len(t, driver.translatedLines, 2)
	first := driver.ParseLine(<-driver.translatedLines)
	second := driver.ParseLine(<-driver.translatedLines)
	require.Len(t, first, 1)
	require.Len(t, second, 1)
	require.Equal(t, "gpt-child-rerouted", first[0].Model)
	require.Equal(t, "gpt-child", second[0].Model)
	require.Equal(t, "thread-child", first[0].SessionID)
	require.Equal(t, "thread-child", second[0].SessionID)

	driver.mu.Lock()
	require.Equal(t, "gpt-parent", driver.effectiveModel)
	require.Equal(t, "gpt-child", driver.modelByThread["thread-child"])
	require.Empty(t, driver.modelByThreadTurn)
	driver.mu.Unlock()
}

func TestCodexAppServerV2ChildUsageHydratesAuthoritativeModelBeforeEmission(t *testing.T) {
	driver, writer := newTestCodexDriver(t, func(req capturedRPCRequest) appServerEnvelope {
		return appServerEnvelope{Result: rawRPCResult(t, appServerThreadResult{
			Thread: appServerThread{ID: "thread-child"},
			Model:  "gpt-child-override",
		})}
	})
	driver.translatedLines = make(chan string, 4)
	bindCodexPrimaryActiveThreadForIsolationTest(driver)
	driver.updateEffectiveModel("gpt-parent")

	// Multi-agent v2 only projects a Started item to the parent; that
	// item intentionally contains no model.
	driver.handleNotification(appServerEnvelope{
		Method: "item/completed",
		Params: rawRPCResult(t, appServerItemNotification{
			ThreadID: "thread-primary",
			TurnID:   "turn-primary",
			Item: appServerItem{
				Type:          "subAgentActivity",
				ID:            "spawn-v2-child",
				Kind:          "started",
				AgentThreadID: "thread-child",
				AgentPath:     "/root/child",
			},
		}),
	})
	activity := driver.ParseLine(<-driver.translatedLines)
	require.Len(t, activity, 1)
	require.Equal(t, consts.ParsedEventKindSubagentActivity, activity[0].Kind)

	driver.handleNotification(appServerEnvelope{
		Method: "thread/tokenUsage/updated",
		Params: rawRPCResult(t, appServerThreadTokenUsageNotification{
			ThreadID: "thread-child",
			TurnID:   "turn-child",
			TokenUsage: appServerThreadTokenUsage{
				Last:  appServerTokenUsageBreakdown{InputTokens: 11, OutputTokens: 4, TotalTokens: 15},
				Total: appServerTokenUsageBreakdown{InputTokens: 11, OutputTokens: 4, TotalTokens: 15},
			},
		}),
	})

	require.Eventually(t, func() bool { return len(driver.translatedLines) == 1 }, time.Second, time.Millisecond)
	usage := driver.ParseLine(<-driver.translatedLines)
	require.Len(t, usage, 1)
	require.Equal(t, consts.ParsedEventKindUsage, usage[0].Kind)
	require.Equal(t, "thread-child", usage[0].SessionID)
	require.Equal(t, "gpt-child-override", usage[0].Model)
	require.Equal(t, uint64(11), usage[0].InputTokens)

	require.Len(t, writer.requests, 2)
	require.Equal(t, "thread/read", writer.requests[0].Method)
	require.Equal(t, "thread-child", writer.requests[0].Params["threadId"])
	require.Equal(t, false, writer.requests[0].Params["includeTurns"])
	require.Equal(t, "thread/resume", writer.requests[1].Method)
	require.Equal(t, "thread-child", writer.requests[1].Params["threadId"])
	require.Equal(t, true, writer.requests[1].Params["excludeTurns"])
	require.Eventually(t, func() bool {
		driver.mu.Lock()
		defer driver.mu.Unlock()
		return driver.modelByThread["thread-child"] == "gpt-child-override" &&
			len(driver.pendingTokenUsageByThread["thread-child"]) == 0 &&
			len(driver.modelHydrationInFlight) == 0
	}, time.Second, time.Millisecond)
}

func TestCodexAppServerChildModelHydrationDoesNotResumeNotLoadedChild(t *testing.T) {
	driver, writer := newTestCodexDriver(t, func(req capturedRPCRequest) appServerEnvelope {
		if req.Method != "thread/read" {
			return appServerEnvelope{Error: &appServerResponseError{
				Code:    -32601,
				Message: "model hydration must not resume a notLoaded child",
			}}
		}
		return appServerEnvelope{Result: rawRPCResult(t, appServerThreadResult{
			Thread: appServerThread{
				ID:     "thread-child",
				Status: appServerThreadStatus{Type: "notLoaded"},
			},
		})}
	})
	driver.translatedLines = make(chan string, 2)
	bindCodexPrimaryActiveThreadForIsolationTest(driver)
	driver.updateEffectiveModel("gpt-parent")

	driver.handleNotification(childUsageNotification(t, 9, 3))

	require.Eventually(t, func() bool {
		driver.mu.Lock()
		defer driver.mu.Unlock()
		return len(driver.modelHydrationInFlight) == 0 &&
			len(driver.pendingTokenUsageByThread["thread-child"]) == 0
	}, time.Second, time.Millisecond)
	require.Equal(t, []string{"thread/read"}, capturedRPCMethods(writer.requests), "notLoaded child must never be cold-resumed for usage recording")
	driver.mu.Lock()
	require.Empty(t, driver.modelByThread["thread-child"])
	_, hasLastUsage := driver.lastTokenUsageByThread["thread-child"]
	driver.mu.Unlock()
	require.False(t, hasLastUsage)
	require.Empty(t, driver.translatedLines, "unknown usage is dropped instead of reviving the child")
}

func TestCodexAppServerChildModelHydrationDoesNotOverwriteNewerSettings(t *testing.T) {
	driver, _ := newTestCodexDriver(t, nil)
	writer := &rawCaptureWriteCloser{}
	driver.rawStdin = writer
	driver.translatedLines = make(chan string, 8)
	bindCodexPrimaryActiveThreadForIsolationTest(driver)
	driver.updateEffectiveModel("gpt-parent")

	// U1 arrives before the child's model is known. Hydration first confirms the
	// child is still loaded, then starts a thread/resume whose authoritative
	// snapshot is model A.
	driver.handleNotification(childUsageNotification(t, 11, 4))
	require.Eventually(t, func() bool { return len(writer.Writes()) == 1 }, time.Second, time.Millisecond)
	var readRequest struct {
		ID int `json:"id"`
	}
	require.NoError(t, json.Unmarshal(writer.Writes()[0], &readRequest))
	require.NotZero(t, readRequest.ID)
	driver.mu.Lock()
	readPending, ok := driver.pending[readRequest.ID]
	delete(driver.pending, readRequest.ID)
	driver.mu.Unlock()
	require.True(t, ok)
	readPending.responseCh <- appServerEnvelope{Result: rawRPCResult(t, appServerThreadResult{
		Thread: appServerThread{ID: "thread-child", Status: appServerThreadStatus{Type: "idle"}},
	})}
	require.Eventually(t, func() bool { return len(writer.Writes()) == 2 }, time.Second, time.Millisecond)
	var resumeRequest struct {
		ID int `json:"id"`
	}
	require.NoError(t, json.Unmarshal(writer.Writes()[1], &resumeRequest))
	require.NotZero(t, resumeRequest.ID)

	driver.handleNotification(appServerEnvelope{
		Method: "thread/settings/updated",
		Params: rawRPCResult(t, map[string]any{
			"threadId": "thread-child",
			"threadSettings": map[string]any{
				"model": "gpt-new-settings",
			},
		}),
	})
	// A later settings notification establishes model B for future usage, but it
	// must not retroactively reattribute U1 while the model-A snapshot is still
	// in flight. U2 is likewise held behind U1 to preserve notification order,
	// while retaining B as the model observed at U2's arrival.
	require.Never(t, func() bool { return len(driver.translatedLines) != 0 }, 20*time.Millisecond, time.Millisecond)
	driver.handleNotification(childUsageNotification(t, 12, 4))
	require.Never(t, func() bool { return len(driver.translatedLines) != 0 }, 20*time.Millisecond, time.Millisecond)

	driver.mu.Lock()
	pending, ok := driver.pending[resumeRequest.ID]
	delete(driver.pending, resumeRequest.ID)
	driver.mu.Unlock()
	require.True(t, ok)
	pending.responseCh <- appServerEnvelope{Result: rawRPCResult(t, appServerThreadResult{
		Thread: appServerThread{ID: "thread-child"},
		Model:  "gpt-stale-snapshot",
	})}
	require.Eventually(t, func() bool { return len(driver.translatedLines) == 2 }, time.Second, time.Millisecond)
	first := driver.ParseLine(<-driver.translatedLines)
	require.Len(t, first, 1)
	require.Equal(t, "gpt-stale-snapshot", first[0].Model)
	require.Equal(t, uint64(11), first[0].InputTokens)
	second := driver.ParseLine(<-driver.translatedLines)
	require.Len(t, second, 1)
	require.Equal(t, "gpt-new-settings", second[0].Model)
	require.Equal(t, uint64(12), second[0].InputTokens)
	require.Eventually(t, func() bool {
		driver.mu.Lock()
		defer driver.mu.Unlock()
		return len(driver.modelHydrationInFlight) == 0 &&
			len(driver.pendingTokenUsageByThread["thread-child"]) == 0
	}, time.Second, time.Millisecond)

	driver.handleNotification(childUsageNotification(t, 13, 4))
	third := driver.ParseLine(<-driver.translatedLines)
	require.Len(t, third, 1)
	require.Equal(t, "gpt-new-settings", third[0].Model)
	require.Equal(t, uint64(13), third[0].InputTokens)
	driver.mu.Lock()
	require.Equal(t, "gpt-new-settings", driver.modelByThread["thread-child"])
	driver.mu.Unlock()
}

func TestCodexAppServerChildModelHydrationDispatchFenceRejectsPostCheckCloseAndReset(t *testing.T) {
	prepareCheckedHydration := func(t *testing.T) (*CodexAppServerDriver, *rpcCaptureWriter, uint64) {
		t.Helper()
		driver, writer := newTestCodexDriver(t, nil)
		const generation uint64 = 7
		driver.launchGeneration = generation
		driver.modelHydrationInFlight = map[string]uint64{"thread-child": generation}
		driver.pendingTokenUsageByThread = map[string][]pendingAppServerThreadTokenUsage{
			"thread-child": {{
				threadID: "thread-child",
				turnID:   "turn-child",
				usage: appServerThreadTokenUsage{
					Last: appServerTokenUsageBreakdown{InputTokens: 1, TotalTokens: 1},
				},
			}},
		}
		require.Equal(t, codexThreadModelHydrationNeeded, driver.threadModelHydrationState(generation, "thread-child"))
		return driver, writer, generation
	}

	t.Run("close wins after state check", func(t *testing.T) {
		driver, writer, generation := prepareCheckedHydration(t)
		driver.setThreadModelHydrationBlocked([]string{"thread-child"}, true)

		_, err := driver.readRunningThreadModel(context.Background(), generation, "thread-child")
		require.ErrorIs(t, err, errCodexAppServerRequestDispatchRejected)
		require.Empty(t, writer.requests, "closed child must not receive a late thread/resume")
	})

	t.Run("new launch wins after state check", func(t *testing.T) {
		driver, oldWriter, generation := prepareCheckedHydration(t)
		newWriter := &rawCaptureWriteCloser{}
		driver.mu.Lock()
		driver.resetAppServerStateForLaunchLocked(drivercontract.LaunchConfig{}, newWriter)
		driver.mu.Unlock()

		_, err := driver.readRunningThreadModel(context.Background(), generation, "thread-child")
		require.ErrorIs(t, err, errCodexAppServerRequestDispatchRejected)
		require.Empty(t, oldWriter.requests, "old generation must not dispatch after reset")
		require.Empty(t, newWriter.Writes(), "old hydration must not write to the new transport")
	})
}

func TestCodexAppServerChildModelHydrationFlushesUsageBeforeNewUsage(t *testing.T) {
	driver, _ := newTestCodexDriver(t, func(req capturedRPCRequest) appServerEnvelope {
		return appServerEnvelope{Result: rawRPCResult(t, appServerThreadResult{
			Thread: appServerThread{ID: "thread-child"},
			Model:  "gpt-child",
		})}
	})
	// An unbuffered output makes the hydration goroutine pause after installing
	// the model but before publishing U1. U2 must wait behind tokenUsageMu.
	driver.translatedLines = make(chan string)
	bindCodexPrimaryActiveThreadForIsolationTest(driver)
	driver.updateEffectiveModel("gpt-parent")

	driver.handleNotification(childUsageNotification(t, 1, 1))
	require.Eventually(t, func() bool {
		driver.mu.Lock()
		defer driver.mu.Unlock()
		return driver.modelByThread["thread-child"] == "gpt-child" &&
			len(driver.pendingTokenUsageByThread["thread-child"]) == 0
	}, time.Second, time.Millisecond)

	u2Done := make(chan struct{})
	go func() {
		driver.handleNotification(childUsageNotification(t, 2, 1))
		close(u2Done)
	}()
	select {
	case <-u2Done:
		t.Fatal("new child usage overtook the pending hydration flush")
	case <-time.After(20 * time.Millisecond):
	}

	first := driver.ParseLine(<-driver.translatedLines)
	require.Len(t, first, 1)
	require.Equal(t, uint64(1), first[0].InputTokens)
	second := driver.ParseLine(<-driver.translatedLines)
	require.Len(t, second, 1)
	require.Equal(t, uint64(2), second[0].InputTokens)
	select {
	case <-u2Done:
	case <-time.After(time.Second):
		t.Fatal("second child usage did not finish")
	}
}

func TestCodexAppServerChildModelHydrationRetriesTransientFailure(t *testing.T) {
	var resumeAttempts atomic.Int32
	driver, writer := newTestCodexDriver(t, func(req capturedRPCRequest) appServerEnvelope {
		switch req.Method {
		case "thread/read":
			return appServerEnvelope{Result: rawRPCResult(t, appServerThreadResult{
				Thread: appServerThread{ID: "thread-child", Status: appServerThreadStatus{Type: "idle"}},
			})}
		case "thread/resume":
			if resumeAttempts.Add(1) == 1 {
				return appServerEnvelope{Error: &appServerResponseError{
					Code:    -32000,
					Message: "transient resume failure",
				}}
			}
			return appServerEnvelope{Result: rawRPCResult(t, appServerThreadResult{
				Thread: appServerThread{ID: "thread-child"},
				Model:  "gpt-child-after-retry",
			})}
		default:
			t.Fatalf("unexpected request %s", req.Method)
			return appServerEnvelope{}
		}
	})
	driver.translatedLines = make(chan string, 2)
	bindCodexPrimaryActiveThreadForIsolationTest(driver)
	driver.updateEffectiveModel("gpt-parent")

	driver.handleNotification(childUsageNotification(t, 7, 2))
	require.Eventually(t, func() bool { return len(driver.translatedLines) == 1 }, time.Second, time.Millisecond)
	usage := driver.ParseLine(<-driver.translatedLines)
	require.Len(t, usage, 1)
	require.Equal(t, "gpt-child-after-retry", usage[0].Model)
	require.Equal(t, int32(2), resumeAttempts.Load())
	require.Equal(t, []string{"thread/read", "thread/resume", "thread/read", "thread/resume"}, capturedRPCMethods(writer.requests))
	require.Eventually(t, func() bool {
		driver.mu.Lock()
		defer driver.mu.Unlock()
		return len(driver.modelHydrationInFlight) == 0 &&
			len(driver.pendingTokenUsageByThread["thread-child"]) == 0
	}, time.Second, time.Millisecond)
}

func TestCodexAppServerChildModelHydrationStopsAndCleansUpAfterClose(t *testing.T) {
	firstFailure := make(chan struct{})
	var attempts atomic.Int32
	driver, writer := newTestCodexDriver(t, func(req capturedRPCRequest) appServerEnvelope {
		if attempts.Add(1) == 1 {
			close(firstFailure)
		}
		return appServerEnvelope{Error: &appServerResponseError{
			Code:    -32000,
			Message: "connection closing",
		}}
	})
	driver.translatedLines = make(chan string, 2)
	bindCodexPrimaryActiveThreadForIsolationTest(driver)
	driver.updateEffectiveModel("gpt-parent")

	driver.handleNotification(childUsageNotification(t, 5, 1))
	select {
	case <-firstFailure:
	case <-time.After(time.Second):
		t.Fatal("timed out waiting for the first hydration failure")
	}
	require.NoError(t, driver.closeRawStdin())

	require.Eventually(t, func() bool {
		driver.mu.Lock()
		defer driver.mu.Unlock()
		return len(driver.modelHydrationInFlight) == 0 &&
			len(driver.pendingTokenUsageByThread["thread-child"]) == 0
	}, time.Second, time.Millisecond)
	require.Equal(t, int32(1), attempts.Load())
	require.Len(t, writer.requests, 1)
	require.Empty(t, driver.translatedLines)
}

func TestCodexAppServerChildModelHydrationDropsDelayedSuccessAfterClose(t *testing.T) {
	driver, _ := newTestCodexDriver(t, nil)
	writer := &rawCaptureWriteCloser{}
	driver.rawStdin = writer
	driver.translatedLines = make(chan string, 4)
	bindCodexPrimaryActiveThreadForIsolationTest(driver)
	driver.updateEffectiveModel("gpt-parent")

	driver.handleNotification(childUsageNotification(t, 6, 2))
	require.Eventually(t, func() bool { return len(writer.Writes()) == 1 }, time.Second, time.Millisecond)
	var request struct {
		ID int `json:"id"`
	}
	require.NoError(t, json.Unmarshal(writer.Writes()[0], &request))
	require.NotZero(t, request.ID)

	driver.handleNotification(appServerEnvelope{
		Method: "item/completed",
		Params: rawRPCResult(t, appServerItemNotification{
			ThreadID: "thread-primary",
			TurnID:   "turn-primary",
			Item: appServerItem{
				Type:          "subAgentActivity",
				ID:            "interrupt-during-hydration",
				Kind:          "interrupted",
				AgentThreadID: "thread-child",
				AgentPath:     "/root/child",
			},
		}),
	})
	activity := driver.ParseLine(<-driver.translatedLines)
	require.Len(t, activity, 1)
	require.Equal(t, consts.ParsedEventKindSubagentActivity, activity[0].Kind)

	driver.mu.Lock()
	pending, ok := driver.pending[request.ID]
	delete(driver.pending, request.ID)
	driver.mu.Unlock()
	require.True(t, ok)
	pending.responseCh <- appServerEnvelope{Result: rawRPCResult(t, appServerThreadResult{
		Thread: appServerThread{ID: "thread-child"},
		Model:  "gpt-closed-child",
	})}

	require.Eventually(t, func() bool {
		driver.mu.Lock()
		defer driver.mu.Unlock()
		return len(driver.modelHydrationInFlight) == 0
	}, time.Second, time.Millisecond)
	driver.mu.Lock()
	require.Empty(t, driver.modelByThread["thread-child"])
	require.Empty(t, driver.pendingTokenUsageByThread["thread-child"])
	_, hasLastUsage := driver.lastTokenUsageByThread["thread-child"]
	require.False(t, hasLastUsage)
	driver.mu.Unlock()
	require.Empty(t, driver.translatedLines)
}

func TestCodexAppServerChildClosePublishesOnlyPendingUsageWithCapturedModel(t *testing.T) {
	driver, _ := newTestCodexDriver(t, nil)
	writer := &rawCaptureWriteCloser{}
	driver.rawStdin = writer
	driver.translatedLines = make(chan string, 8)
	bindCodexPrimaryActiveThreadForIsolationTest(driver)
	driver.updateEffectiveModel("gpt-parent")

	// U1 has no model and starts a delayed hydration snapshot A.
	driver.handleNotification(childUsageNotification(t, 11, 4))
	require.Eventually(t, func() bool { return len(writer.Writes()) == 1 }, time.Second, time.Millisecond)
	var request struct {
		ID int `json:"id"`
	}
	require.NoError(t, json.Unmarshal(writer.Writes()[0], &request))
	require.NotZero(t, request.ID)

	// Settings B applies only to future usage. U2 captures B but remains queued
	// behind U1 until the in-flight hydration reaches a terminal state.
	driver.handleNotification(appServerEnvelope{
		Method: "thread/settings/updated",
		Params: rawRPCResult(t, map[string]any{
			"threadId": "thread-child",
			"threadSettings": map[string]any{
				"model": "gpt-new-settings",
			},
		}),
	})
	driver.handleNotification(childUsageNotification(t, 12, 4))
	require.Empty(t, driver.translatedLines)

	// Close wins before snapshot A returns. The close fence must reject A for U1
	// while still publishing U2 with the model it captured before the fence.
	driver.handleNotification(appServerEnvelope{
		Method: "item/completed",
		Params: rawRPCResult(t, appServerItemNotification{
			ThreadID: "thread-primary",
			TurnID:   "turn-primary",
			Item: appServerItem{
				Type:          "subAgentActivity",
				ID:            "interrupt-after-u2",
				Kind:          "interrupted",
				AgentThreadID: "thread-child",
				AgentPath:     "/root/child",
			},
		}),
	})
	activity := driver.ParseLine(<-driver.translatedLines)
	require.Len(t, activity, 1)
	require.Equal(t, consts.ParsedEventKindSubagentActivity, activity[0].Kind)

	driver.mu.Lock()
	pending, ok := driver.pending[request.ID]
	delete(driver.pending, request.ID)
	driver.mu.Unlock()
	require.True(t, ok)
	pending.responseCh <- appServerEnvelope{Result: rawRPCResult(t, appServerThreadResult{
		Thread: appServerThread{ID: "thread-child"},
		Model:  "gpt-stale-snapshot",
	})}

	require.Eventually(t, func() bool { return len(driver.translatedLines) == 1 }, time.Second, time.Millisecond)
	usage := driver.ParseLine(<-driver.translatedLines)
	require.Len(t, usage, 1)
	require.Equal(t, consts.ParsedEventKindUsage, usage[0].Kind)
	require.Equal(t, "gpt-new-settings", usage[0].Model)
	require.Equal(t, uint64(12), usage[0].InputTokens)
	driver.mu.Lock()
	require.Equal(t, "gpt-new-settings", driver.modelByThread["thread-child"], "late snapshot A must not install")
	require.Empty(t, driver.pendingTokenUsageByThread["thread-child"])
	require.Empty(t, driver.modelHydrationInFlight)
	_, blocked := driver.modelHydrationBlocked["thread-child"]
	require.True(t, blocked)
	driver.mu.Unlock()

	// A later explicit interaction reopens the child. Replaying U2 must remain
	// deduplicated because its captured B source was already published.
	driver.handleNotification(appServerEnvelope{
		Method: "item/completed",
		Params: rawRPCResult(t, appServerItemNotification{
			ThreadID: "thread-primary",
			TurnID:   "turn-primary",
			Item: appServerItem{
				Type:          "subAgentActivity",
				ID:            "interact-after-close",
				Kind:          "interacted",
				AgentThreadID: "thread-child",
				AgentPath:     "/root/child",
			},
		}),
	})
	interaction := driver.ParseLine(<-driver.translatedLines)
	require.Len(t, interaction, 1)
	require.Equal(t, consts.ParsedEventKindSubagentActivity, interaction[0].Kind)
	driver.handleNotification(childUsageNotification(t, 12, 4))
	time.Sleep(2 * codexThreadModelHydrationBackoff)
	require.Empty(t, driver.translatedLines)
}

func TestCodexAppServerClosedChildRejectsLateUnknownUsage(t *testing.T) {
	tests := []struct {
		name  string
		block func(*testing.T, *CodexAppServerDriver)
	}{
		{
			name: "sub-agent interrupted",
			block: func(t *testing.T, driver *CodexAppServerDriver) {
				driver.handleNotification(appServerEnvelope{
					Method: "item/completed",
					Params: rawRPCResult(t, appServerItemNotification{
						ThreadID: "thread-primary",
						TurnID:   "turn-primary",
						Item: appServerItem{
							Type:          "subAgentActivity",
							ID:            "interrupt-child",
							Kind:          "interrupted",
							AgentThreadID: "thread-child",
							AgentPath:     "/root/child",
						},
					}),
				})
			},
		},
		{
			name: "close-agent tool",
			block: func(t *testing.T, driver *CodexAppServerDriver) {
				driver.handleNotification(appServerEnvelope{
					Method: "item/started",
					Params: rawRPCResult(t, appServerItemNotification{
						ThreadID: "thread-primary",
						TurnID:   "turn-primary",
						Item: appServerItem{
							Type:              "collabAgentToolCall",
							ID:                "close-child",
							Tool:              "closeAgent",
							ReceiverThreadIDs: []string{"thread-child"},
						},
					}),
				})
			},
		},
	}

	for _, testCase := range tests {
		t.Run(testCase.name, func(t *testing.T) {
			var hydrationRequests atomic.Int32
			driver, _ := newTestCodexDriver(t, func(req capturedRPCRequest) appServerEnvelope {
				hydrationRequests.Add(1)
				return appServerEnvelope{Error: &appServerResponseError{Code: -32000, Message: "unexpected hydration"}}
			})
			driver.translatedLines = make(chan string, 4)
			bindCodexPrimaryActiveThreadForIsolationTest(driver)
			driver.updateEffectiveModel("gpt-parent")

			testCase.block(t, driver)
			require.NotEmpty(t, driver.translatedLines)
			<-driver.translatedLines
			driver.handleNotification(childUsageNotification(t, 9, 2))

			time.Sleep(2 * codexThreadModelHydrationBackoff)
			require.Equal(t, int32(0), hydrationRequests.Load())
			driver.mu.Lock()
			require.Empty(t, driver.pendingTokenUsageByThread["thread-child"])
			require.Empty(t, driver.modelHydrationInFlight)
			_, hasLastUsage := driver.lastTokenUsageByThread["thread-child"]
			require.False(t, hasLastUsage)
			driver.mu.Unlock()
			require.Empty(t, driver.translatedLines)
		})
	}
}

func TestCodexAppServerFailedCloseRemainsBlockedUntilExplicitResume(t *testing.T) {
	driver, writer := newTestCodexDriver(t, func(req capturedRPCRequest) appServerEnvelope {
		return appServerEnvelope{Result: rawRPCResult(t, appServerThreadResult{
			Thread: appServerThread{ID: "thread-child"},
			Model:  "gpt-live-child",
		})}
	})
	driver.translatedLines = make(chan string, 4)
	bindCodexPrimaryActiveThreadForIsolationTest(driver)
	driver.updateEffectiveModel("gpt-parent")

	for _, notification := range []struct {
		method string
		status string
	}{
		{method: "item/started", status: "inProgress"},
		{method: "item/completed", status: "failed"},
	} {
		driver.handleNotification(appServerEnvelope{
			Method: notification.method,
			Params: rawRPCResult(t, appServerItemNotification{
				ThreadID: "thread-primary",
				TurnID:   "turn-primary",
				Item: appServerItem{
					Type:              "collabAgentToolCall",
					ID:                "failed-close-child",
					Tool:              "closeAgent",
					Status:            notification.status,
					ReceiverThreadIDs: []string{"thread-child"},
				},
			}),
		})
	}
	require.Len(t, driver.translatedLines, 2)
	<-driver.translatedLines
	<-driver.translatedLines

	driver.handleNotification(childUsageNotification(t, 8, 3))
	time.Sleep(2 * codexThreadModelHydrationBackoff)
	require.Empty(t, driver.translatedLines)
	require.Empty(t, writer.requests)
	driver.mu.Lock()
	_, blocked := driver.modelHydrationBlocked["thread-child"]
	require.True(t, blocked)
	driver.mu.Unlock()

	// A later successful resume is an explicit liveness signal and is the only
	// thing here that should reopen child model hydration.
	driver.handleNotification(appServerEnvelope{
		Method: "item/completed",
		Params: rawRPCResult(t, appServerItemNotification{
			ThreadID: "thread-primary",
			TurnID:   "turn-primary",
			Item: appServerItem{
				Type:              "collabAgentToolCall",
				ID:                "resume-after-close",
				Tool:              "resumeAgent",
				Status:            "completed",
				ReceiverThreadIDs: []string{"thread-child"},
			},
		}),
	})
	require.Len(t, driver.translatedLines, 1)
	resumeOutput := driver.ParseLine(<-driver.translatedLines)
	require.Len(t, resumeOutput, 1)
	require.Equal(t, consts.ParsedEventKindToolOutput, resumeOutput[0].Kind)

	driver.handleNotification(childUsageNotification(t, 8, 3))
	require.Eventually(t, func() bool { return len(driver.translatedLines) == 1 }, time.Second, time.Millisecond)
	usage := driver.ParseLine(<-driver.translatedLines)
	require.Len(t, usage, 1)
	require.Equal(t, consts.ParsedEventKindUsage, usage[0].Kind)
	require.Equal(t, "gpt-live-child", usage[0].Model)
	require.Equal(t, []string{"thread/read", "thread/resume"}, capturedRPCMethods(writer.requests))
	driver.mu.Lock()
	_, blocked = driver.modelHydrationBlocked["thread-child"]
	require.False(t, blocked)
	driver.mu.Unlock()
}

func TestCodexAppServerInteractedActivityRestoresChildModelHydration(t *testing.T) {
	driver, writer := newTestCodexDriver(t, func(req capturedRPCRequest) appServerEnvelope {
		return appServerEnvelope{Result: rawRPCResult(t, appServerThreadResult{
			Thread: appServerThread{ID: "thread-child"},
			Model:  "gpt-followup-child",
		})}
	})
	driver.translatedLines = make(chan string, 4)
	bindCodexPrimaryActiveThreadForIsolationTest(driver)
	driver.updateEffectiveModel("gpt-parent")

	for index, kind := range []string{"interrupted", "interacted"} {
		driver.handleNotification(appServerEnvelope{
			Method: "item/completed",
			Params: rawRPCResult(t, appServerItemNotification{
				ThreadID: "thread-primary",
				TurnID:   "turn-primary",
				Item: appServerItem{
					Type:          "subAgentActivity",
					ID:            fmt.Sprintf("child-activity-%d", index),
					Kind:          kind,
					AgentThreadID: "thread-child",
					AgentPath:     "/root/child",
				},
			}),
		})
	}
	require.Len(t, driver.translatedLines, 2)
	<-driver.translatedLines
	<-driver.translatedLines

	driver.handleNotification(childUsageNotification(t, 10, 4))
	require.Eventually(t, func() bool { return len(driver.translatedLines) == 1 }, time.Second, time.Millisecond)
	usage := driver.ParseLine(<-driver.translatedLines)
	require.Len(t, usage, 1)
	require.Equal(t, "gpt-followup-child", usage[0].Model)
	require.Equal(t, []string{"thread/read", "thread/resume"}, capturedRPCMethods(writer.requests))
	driver.mu.Lock()
	_, blocked := driver.modelHydrationBlocked["thread-child"]
	require.False(t, blocked)
	driver.mu.Unlock()
}

func TestCodexAppServerFailedResumeDoesNotUnblockClosedChildHydration(t *testing.T) {
	var hydrationRequests atomic.Int32
	driver, _ := newTestCodexDriver(t, func(req capturedRPCRequest) appServerEnvelope {
		hydrationRequests.Add(1)
		return appServerEnvelope{Error: &appServerResponseError{Code: -32000, Message: "unexpected hydration"}}
	})
	driver.translatedLines = make(chan string, 4)
	bindCodexPrimaryActiveThreadForIsolationTest(driver)
	driver.updateEffectiveModel("gpt-parent")
	driver.setThreadModelHydrationBlocked([]string{"thread-child"}, true)

	driver.handleNotification(appServerEnvelope{
		Method: "item/completed",
		Params: rawRPCResult(t, appServerItemNotification{
			ThreadID: "thread-primary",
			TurnID:   "turn-primary",
			Item: appServerItem{
				Type:              "collabAgentToolCall",
				ID:                "failed-resume-child",
				Tool:              "resumeAgent",
				Status:            "failed",
				ReceiverThreadIDs: []string{"thread-child"},
			},
		}),
	})
	<-driver.translatedLines
	driver.handleNotification(childUsageNotification(t, 8, 3))
	time.Sleep(2 * codexThreadModelHydrationBackoff)

	require.Equal(t, int32(0), hydrationRequests.Load())
	driver.mu.Lock()
	_, blocked := driver.modelHydrationBlocked["thread-child"]
	require.True(t, blocked)
	require.Empty(t, driver.pendingTokenUsageByThread["thread-child"])
	require.Empty(t, driver.modelHydrationInFlight)
	driver.mu.Unlock()
	require.Empty(t, driver.translatedLines)
}

func TestCodexAppServerNestedInterruptedActivityBlocksGrandchildHydrationWithoutRootProjection(t *testing.T) {
	var hydrationRequests atomic.Int32
	driver, _ := newTestCodexDriver(t, func(req capturedRPCRequest) appServerEnvelope {
		hydrationRequests.Add(1)
		return appServerEnvelope{Error: &appServerResponseError{Code: -32000, Message: "unexpected hydration"}}
	})
	driver.translatedLines = make(chan string, 4)
	bindCodexPrimaryActiveThreadForIsolationTest(driver)
	driver.updateEffectiveModel("gpt-parent")

	driver.handleNotification(appServerEnvelope{
		Method: "item/completed",
		Params: rawRPCResult(t, appServerItemNotification{
			ThreadID: "thread-child",
			TurnID:   "turn-child",
			Item: appServerItem{
				Type:          "subAgentActivity",
				ID:            "nested-interrupted",
				Kind:          "interrupted",
				AgentThreadID: "thread-grandchild",
				AgentPath:     "/root/child/grandchild",
			},
		}),
	})
	require.Empty(t, driver.translatedLines)
	driver.handleNotification(threadUsageNotification(t, "thread-grandchild", "turn-grandchild", 4, 1))
	time.Sleep(2 * codexThreadModelHydrationBackoff)

	require.Equal(t, int32(0), hydrationRequests.Load())
	driver.mu.Lock()
	_, blocked := driver.modelHydrationBlocked["thread-grandchild"]
	require.True(t, blocked)
	require.Empty(t, driver.pendingTokenUsageByThread["thread-grandchild"])
	require.Empty(t, driver.modelHydrationInFlight)
	_, hasLastUsage := driver.lastTokenUsageByThread["thread-grandchild"]
	require.False(t, hasLastUsage)
	driver.mu.Unlock()
	require.Empty(t, driver.translatedLines)
}

func TestCodexAppServerNestedInteractedActivityRestoresGrandchildHydrationWithoutRootProjection(t *testing.T) {
	driver, writer := newTestCodexDriver(t, func(req capturedRPCRequest) appServerEnvelope {
		return appServerEnvelope{Result: rawRPCResult(t, appServerThreadResult{
			Thread: appServerThread{ID: "thread-grandchild"},
			Model:  "gpt-grandchild",
		})}
	})
	driver.translatedLines = make(chan string, 4)
	bindCodexPrimaryActiveThreadForIsolationTest(driver)
	driver.updateEffectiveModel("gpt-parent")
	driver.setThreadModelHydrationBlocked([]string{"thread-grandchild"}, true)

	driver.handleNotification(appServerEnvelope{
		Method: "item/completed",
		Params: rawRPCResult(t, appServerItemNotification{
			ThreadID: "thread-child",
			TurnID:   "turn-child",
			Item: appServerItem{
				Type:          "subAgentActivity",
				ID:            "nested-interacted",
				Kind:          "interacted",
				AgentThreadID: "thread-grandchild",
				AgentPath:     "/root/child/grandchild",
			},
		}),
	})
	require.Empty(t, driver.translatedLines)
	driver.handleNotification(threadUsageNotification(t, "thread-grandchild", "turn-grandchild", 12, 5))
	require.Eventually(t, func() bool { return len(driver.translatedLines) == 1 }, time.Second, time.Millisecond)
	usage := driver.ParseLine(<-driver.translatedLines)
	require.Len(t, usage, 1)
	require.Equal(t, consts.ParsedEventKindUsage, usage[0].Kind)
	require.Equal(t, "thread-grandchild", usage[0].SessionID)
	require.Equal(t, "gpt-grandchild", usage[0].Model)
	require.Equal(t, []string{"thread/read", "thread/resume"}, capturedRPCMethods(writer.requests))
	driver.mu.Lock()
	_, blocked := driver.modelHydrationBlocked["thread-grandchild"]
	require.False(t, blocked)
	driver.mu.Unlock()
}

func childUsageNotification(t *testing.T, inputTokens int64, outputTokens int64) appServerEnvelope {
	return threadUsageNotification(t, "thread-child", "turn-child", inputTokens, outputTokens)
}

func threadUsageNotification(t *testing.T, threadID string, turnID string, inputTokens int64, outputTokens int64) appServerEnvelope {
	t.Helper()
	return appServerEnvelope{
		Method: "thread/tokenUsage/updated",
		Params: rawRPCResult(t, appServerThreadTokenUsageNotification{
			ThreadID: threadID,
			TurnID:   turnID,
			TokenUsage: appServerThreadTokenUsage{
				Last: appServerTokenUsageBreakdown{
					InputTokens:  inputTokens,
					OutputTokens: outputTokens,
					TotalTokens:  inputTokens + outputTokens,
				},
				Total: appServerTokenUsageBreakdown{
					InputTokens:  inputTokens,
					OutputTokens: outputTokens,
					TotalTokens:  inputTokens + outputTokens,
				},
			},
		}),
	}
}

func TestCodexAppServerForeignThreadServerRequestExecutesWithoutTranslatedToolOutput(t *testing.T) {
	var calls atomic.Int32
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer r.Body.Close()
		calls.Add(1)
		body, err := io.ReadAll(r.Body)
		require.NoError(t, err)
		var request struct {
			Params struct {
				Name string `json:"name"`
			} `json:"params"`
		}
		require.NoError(t, json.Unmarshal(body, &request))
		require.Equal(t, "message_history_read", request.Params.Name)
		w.Header().Set("Content-Type", "application/json")
		_, err = io.WriteString(w, `{"jsonrpc":"2.0","id":1,"result":{"content":[{"type":"text","text":"child tool completed"}]}}`)
		require.NoError(t, err)
	}))
	defer server.Close()

	driver := &CodexAppServerDriver{
		codexRPCState: codexRPCState{
			translatedLines: make(chan string, 4),
		},
		codexLaunchState: codexLaunchState{
			mcpConfigs: map[string]types.NamedMCPConfig{
				"agentteam": {Name: "agentteam", URL: server.URL},
			},
			httpClient: server.Client(),
		},
	}
	bindCodexPrimaryActiveThreadForIsolationTest(driver)

	result, rpcErr := driver.executeServerRequestWithActivity(
		context.Background(),
		"request-child",
		"mcpServer/tool/call",
		json.RawMessage(`{"threadId":"thread-child","turnId":"turn-child","callId":"call-child","server":"agentteam","tool":"message_history_read","arguments":{}}`),
	)
	require.Nil(t, rpcErr)
	require.NotNil(t, result)
	require.Equal(t, int32(1), calls.Load())
	require.Empty(t, driver.translatedLines)

	status, activeTurnID, ok := driver.currentThreadState("thread-primary")
	require.True(t, ok)
	require.Equal(t, appServerThreadStatus{Type: "active"}, status)
	require.Equal(t, "turn-primary", activeTurnID)

	result, rpcErr = driver.executeServerRequestWithActivity(
		context.Background(),
		"request-primary",
		"mcpServer/tool/call",
		json.RawMessage(`{"threadId":"thread-primary","turnId":"turn-primary","callId":"call-primary","server":"agentteam","tool":"message_history_read","arguments":{}}`),
	)
	require.Nil(t, rpcErr)
	require.NotNil(t, result)
	require.Equal(t, int32(2), calls.Load())
	events := driver.ParseLine(<-driver.translatedLines)
	require.Len(t, events, 1)
	require.Equal(t, consts.ParsedEventKindToolOutput, events[0].Kind)
	require.Equal(t, "thread-primary", events[0].SessionID)
}

func TestCodexAppServerThreadEventsDoNotBindPrimaryThreadID(t *testing.T) {
	driver, _ := newTestCodexDriver(t, nil)
	driver.translatedLines = make(chan string, 4)

	driver.handleNotification(appServerEnvelope{
		Method: "turn/started",
		Params: rawRPCResult(t, appServerTurnNotification{
			ThreadID: "thread-before-snapshot",
			Turn: appServerTurn{
				ID:     "turn-before-snapshot",
				Status: appServerTurnStatus{Type: "inProgress"},
			},
		}),
	})
	driver.handleNotification(appServerEnvelope{
		Method: "thread/tokenUsage/updated",
		Params: rawRPCResult(t, appServerThreadTokenUsageNotification{
			ThreadID: "thread-before-snapshot",
			TurnID:   "turn-before-snapshot",
			TokenUsage: appServerThreadTokenUsage{
				Last:  appServerTokenUsageBreakdown{InputTokens: 1, TotalTokens: 1},
				Total: appServerTokenUsageBreakdown{InputTokens: 1, TotalTokens: 1},
			},
		}),
	})

	driver.mu.Lock()
	require.Empty(t, driver.threadID)
	driver.mu.Unlock()

	driver.updateThreadSnapshot(appServerThread{
		ID:     "thread-primary",
		Status: appServerThreadStatus{Type: "idle"},
	})
	driver.mu.Lock()
	require.Equal(t, "thread-primary", driver.threadID)
	driver.mu.Unlock()
}