X-Git-Url: http://git.giacomolli.uk/onechat/blobdiff_plain/1c1771140063b1b4aa5d508ced1b285e773fce04..ba9faee0cd6c35020997794ef7bf134e8db12e96:/responses_test.go

diff --git a/responses_test.go b/responses_test.go
index c1a7b4d..de928bd 100644
--- a/responses_test.go
+++ b/responses_test.go
@@ -11,6 +11,7 @@ import (
 	"path/filepath"
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/require"
 )
@@ -60,6 +61,42 @@ func TestClientSendsTheFixedRequestFields(t *testing.T) {
 	require.Len(t, body["tools"], 1)
 }
 
+func TestClientSendsTheSessionHeaderToChatGPT(t *testing.T) {
+	login := newFakeAuth(t)
+	auth := newTestAuth(t, login)
+	writeTestTokens(t, auth, login.token(time.Hour), "refresh-0")
+
+	testCases := []struct {
+		name string
+		auth authorizer
+		want string
+	}{
+		{
+			name: "ChatGPT backend",
+			auth: auth,
+			want: "agent-key",
+		},
+		{
+			name: "OpenAI API",
+			auth: apiKey(testKey),
+			want: "",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			server := newFakeResponses(t)
+			client := newClient(server.URL, tc.auth, "agent-key", filepath.Join(t.TempDir(), "usage.jsonl"))
+			_, err := client.Create(t.Context(), testRequest())
+			require.NoError(t, err)
+
+			call := server.Call(0)
+			require.Equal(t, tc.want, call.header.Get("session_id"))
+			require.Equal(t, "agent-key", call.body["prompt_cache_key"])
+		})
+	}
+}
+
 func TestClientReadsTheStream(t *testing.T) {
 	server := newFakeResponses(t)
 	server.reply = func(w http.ResponseWriter, call fakeCall) {
@@ -102,6 +139,7 @@ func TestClientRecordsEachCall(t *testing.T) {
 		w.Header().Set("X-Codex-Primary-Used-Percent", "12")
 		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"}`)