diff --git a/e2e/agentnetwork/chat_test.go b/e2e/agentnetwork/chat_test.go index 6cd23c34e..5e3a79273 100644 --- a/e2e/agentnetwork/chat_test.go +++ b/e2e/agentnetwork/chat_test.go @@ -212,6 +212,10 @@ func TestProvidersMatrix(t *testing.T) { t.Run(pc.name, func(t *testing.T) { before, _ := srv.ListAccessLogs(ctx) + // Unique per provider so we can find this provider's row by its + // session id and confirm the marker propagated end-to-end. + sessionID := "e2e-session-" + pc.name + // Retry briefly to absorb tunnel/DNS jitter on the first call. var code int var body string @@ -221,9 +225,9 @@ func TestProvidersMatrix(t *testing.T) { var b string var cerr error if pc.kind == harness.WireVertex { - c, b, cerr = cl.Vertex(ctx, settings.Endpoint, proxyIP, pc.project, pc.region, pc.model, "Reply with exactly: pong") + c, b, cerr = cl.Vertex(ctx, settings.Endpoint, proxyIP, pc.project, pc.region, pc.model, "Reply with exactly: pong", sessionID) } else { - c, b, cerr = cl.Chat(ctx, settings.Endpoint, proxyIP, pc.kind, pc.model, "Reply with exactly: pong") + c, b, cerr = cl.Chat(ctx, settings.Endpoint, proxyIP, pc.kind, pc.model, "Reply with exactly: pong", sessionID) } if cerr == nil { code, body = c, b @@ -239,6 +243,21 @@ func TestProvidersMatrix(t *testing.T) { logs, lerr := srv.ListAccessLogs(ctx) return lerr == nil && logs.TotalRecords > before.TotalRecords }, 30*time.Second, 2*time.Second, "an access-log row should be ingested for %s", pc.name) + + // The session id sent as x-session-id must round-trip into the + // access-log row for this provider. + require.Eventually(t, func() bool { + logs, lerr := srv.ListAccessLogs(ctx) + if lerr != nil { + return false + } + for _, r := range logs.Data { + if r.SessionId != nil && *r.SessionId == sessionID { + return true + } + } + return false + }, 30*time.Second, 2*time.Second, "session id %q must be recorded in an access-log row for %s", sessionID, pc.name) }) } diff --git a/e2e/harness/client.go b/e2e/harness/client.go index 1552ad3d4..cf7ef8945 100644 --- a/e2e/harness/client.go +++ b/e2e/harness/client.go @@ -158,8 +158,9 @@ const ( // Chat issues a chat-completion POST to the agent-network endpoint over the // client's tunnel, returning the HTTP status and response body. kind selects -// the wire shape: WireChat (OpenAI) or WireMessages (Anthropic). -func (cl *Client) Chat(ctx context.Context, endpoint, proxyIP, kind, model, prompt string) (int, string, error) { +// the wire shape: WireChat (OpenAI) or WireMessages (Anthropic). A non-empty +// sessionID is sent as the universal x-session-id header the proxy records. +func (cl *Client) Chat(ctx context.Context, endpoint, proxyIP, kind, model, prompt, sessionID string) (int, string, error) { var path, body string var headers []string switch kind { @@ -171,17 +172,26 @@ func (cl *Client) Chat(ctx context.Context, endpoint, proxyIP, kind, model, prom path = "/v1/chat/completions" body = fmt.Sprintf(`{"model":%q,"messages":[{"role":"user","content":%q}]}`, model, prompt) } - return cl.post(ctx, endpoint, proxyIP, path, body, headers) + return cl.post(ctx, endpoint, proxyIP, path, body, withSessionID(headers, sessionID)) } // Vertex issues an Anthropic-on-Vertex rawPredict POST over the tunnel. Unlike // Chat, the model is carried in the request path (project/region/model), so the // proxy routes by path and mints the service-account OAuth token; the body uses -// the Vertex anthropic_version rather than a model field. -func (cl *Client) Vertex(ctx context.Context, endpoint, proxyIP, project, region, model, prompt string) (int, string, error) { +// the Vertex anthropic_version rather than a model field. A non-empty sessionID +// is sent as the universal x-session-id header the proxy records. +func (cl *Client) Vertex(ctx context.Context, endpoint, proxyIP, project, region, model, prompt, sessionID string) (int, string, error) { path := fmt.Sprintf("/v1/projects/%s/locations/%s/publishers/anthropic/models/%s:rawPredict", project, region, model) body := fmt.Sprintf(`{"anthropic_version":"vertex-2023-10-16","max_tokens":64,"messages":[{"role":"user","content":%q}]}`, prompt) - return cl.post(ctx, endpoint, proxyIP, path, body, nil) + return cl.post(ctx, endpoint, proxyIP, path, body, withSessionID(nil, sessionID)) +} + +// withSessionID appends the x-session-id header when sessionID is non-empty. +func withSessionID(headers []string, sessionID string) []string { + if sessionID == "" { + return headers + } + return append(headers, "x-session-id: "+sessionID) } // post runs curl in a throwaway container sharing the client's network