X-Git-Url: https://git.giacomolli.uk/onechat/blobdiff_plain/6b21fcf78372995c6d8130c63f2f8aa2609bf7fe..59e1537e1969f7c109084497287d729edbedba04:/model_test.go?ds=sidebyside

diff --git a/model_test.go b/model_test.go
index 84b767e..90d0e93 100644
--- a/model_test.go
+++ b/model_test.go
@@ -85,13 +85,26 @@ func TestClientSendsTheSessionHeaderToChatGPT(t *testing.T) {
 			call := server.Call(0)
 			require.Equal(t, tc.want, call.header.Get("session_id"))
 			require.Equal(t, "agent-key", call.body["prompt_cache_key"])
+
+			// A subagent has its own key.
+			request := modelTestRequest()
+			request.Agent = "GreenOtter"
+			_, err = client.Create(t.Context(), request)
+			require.NoError(t, err)
+
+			call = server.Call(1)
+			if tc.want != "" {
+				require.Equal(t, "agent-key-GreenOtter", call.header.Get("session_id"))
+			}
+
+			require.Equal(t, "agent-key-GreenOtter", call.body["prompt_cache_key"])
 		})
 	}
 }
 
 func TestClientReadsTheStream(t *testing.T) {
 	server := newModelFake(t)
-	server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
+	server.reply = func(w http.ResponseWriter, _ ModelFakeCall) {
 		modelWriteEvents(w,
 			map[string]any{"type": "response.created", "response": map[string]any{}},
 			map[string]any{"type": "response.output_item.done", "item": map[string]any{"type": "reasoning", "encrypted_content": "opaque"}},
@@ -116,7 +129,7 @@ func TestClientReadsTheStream(t *testing.T) {
 
 func TestClientTakesTheOutputOfTheCompletedResponse(t *testing.T) {
 	server := newModelFake(t)
-	server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
+	server.reply = func(w http.ResponseWriter, _ ModelFakeCall) {
 		modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("only here")}, 10, 0, 2))
 	}
 
@@ -132,8 +145,10 @@ func TestClientRecordsEachCall(t *testing.T) {
 		w.Header().Set("X-Codex-Plan-Type", "plus")
 		w.Header().Set("X-Other", "not quota")
 		w.Header().Set("X-Codex-Turn-State", "opaque state")
+
 		if call.number == 1 {
 			w.WriteHeader(http.StatusTooManyRequests)
+
 			_, _ = io.WriteString(w, `{"detail": "slow down"}`)
 
 			return
@@ -145,10 +160,12 @@ func TestClientRecordsEachCall(t *testing.T) {
 	client := newTestModelClient(t, server)
 	request := modelTestRequest()
 	request.Node = &NodeKey{2, 5}
+	request.Agent = "GreenOtter"
 	_, err := client.Create(t.Context(), request)
 	require.NoError(t, err)
 	require.NotContains(t, server.Call(0).body, "node")
 	require.NotContains(t, string(server.Call(0).raw), "[2,5]")
+	require.NotContains(t, server.Call(0).body, "agent")
 
 	_, err = client.Create(t.Context(), modelTestRequest())
 	require.ErrorContains(t, err, "HTTP 429")
@@ -161,6 +178,7 @@ func TestClientRecordsEachCall(t *testing.T) {
 		"x-codex-primary-used-percent": "12",
 		"x-codex-plan-type":            "plus",
 	}
+
 	require.Equal(t, "gpt-test", records[0].Model)
 	require.Equal(t, http.StatusOK, records[0].Status)
 	require.Equal(t, 500, records[0].InputTokens)
@@ -168,12 +186,14 @@ func TestClientRecordsEachCall(t *testing.T) {
 	require.Equal(t, 7, records[0].OutputTokens)
 	require.Equal(t, quota, records[0].Quota)
 	require.Equal(t, &[2]int{2, 5}, records[0].Node)
+	require.Equal(t, "GreenOtter", records[0].Agent)
 	require.Equal(t, len("ok"), records[0].TextBytes)
 	require.Greater(t, records[0].Seconds, 0.0)
 	require.Empty(t, records[0].Error)
 	require.NotEmpty(t, records[0].Date)
 
 	require.Nil(t, records[1].Node)
+	require.Empty(t, records[1].Agent)
 	require.Zero(t, records[1].TextBytes)
 
 	require.Equal(t, http.StatusTooManyRequests, records[1].Status)
@@ -196,11 +216,13 @@ func TestClientRefusesCallsOverTheLimit(t *testing.T) {
 
 	for range 2 {
 		require.False(t, client.LimitReached())
+
 		_, err := client.Create(t.Context(), modelTestRequest())
 		require.NoError(t, err)
 	}
 
 	require.True(t, client.LimitReached())
+
 	_, err := client.Create(t.Context(), modelTestRequest())
 	require.ErrorIs(t, err, errModelCallLimit)
 	require.Equal(t, 2, server.Calls())
@@ -238,7 +260,7 @@ func TestClientFailsWhenTheResponseFails(t *testing.T) {
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
 			server := newModelFake(t)
-			server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
+			server.reply = func(w http.ResponseWriter, _ ModelFakeCall) {
 				modelWriteEvents(w, tc.events...)
 			}