workspace-server/internal/handlers/a2a_proxy_helpers_test.go

@@ -158,6 +158,151 @@ func TestNilIfEmpty_Contract(t *testing.T) {
 	}
 }
 
+// ──────────────────────────────────────────────────────────────────────────────
+// parseUsageFromA2AResponse
+// ──────────────────────────────────────────────────────────────────────────────
+
+func TestParseUsageFromA2AResponse_EmptyAndMalformed(t *testing.T) {
+	cases := []struct {
+		name string
+		body []byte
+	}{
+		{"nil", nil},
+		{"empty", []byte{}},
+		{"non-JSON", []byte("not json")},
+		{"empty object", []byte("{}")},
+		{"null result", []byte(`{"result": null}`)},
+		{"string result", []byte(`{"result": "hello"}`)},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			in, out := parseUsageFromA2AResponse(tc.body)
+			if in != 0 || out != 0 {
+				t.Errorf("parseUsageFromA2AResponse = (%d, %d), want (0, 0)", in, out)
+			}
+		})
+	}
+}
+
+func TestParseUsageFromA2AResponse_ResultUsageShape(t *testing.T) {
+	body := []byte(`{
+		"result": {
+			"usage": {"input_tokens": 1500, "output_tokens": 320}
+		}
+	}`)
+	in, out := parseUsageFromA2AResponse(body)
+	if in != 1500 || out != 320 {
+		t.Errorf("parseUsageFromA2AResponse = (%d, %d), want (1500, 320)", in, out)
+	}
+}
+
+func TestParseUsageFromA2AResponse_TopLevelUsage(t *testing.T) {
+	body := []byte(`{
+		"usage": {"input_tokens": 100, "output_tokens": 50}
+	}`)
+	in, out := parseUsageFromA2AResponse(body)
+	if in != 100 || out != 50 {
+		t.Errorf("parseUsageFromA2AResponse = (%d, %d), want (100, 50)", in, out)
+	}
+}
+
+func TestParseUsageFromA2AResponse_BothPresentPrefersResult(t *testing.T) {
+	// When both result.usage and top-level usage exist, result.usage wins.
+	body := []byte(`{
+		"result": {"usage": {"input_tokens": 500, "output_tokens": 200}},
+		"usage": {"input_tokens": 50, "output_tokens": 20}
+	}`)
+	in, out := parseUsageFromA2AResponse(body)
+	if in != 500 || out != 200 {
+		t.Errorf("parseUsageFromA2AResponse = (%d, %d), want (500, 200) from result.usage", in, out)
+	}
+}
+
+func TestParseUsageFromA2AResponse_ZeroUsage(t *testing.T) {
+	// Zero values are treated as absent (readUsageMap returns ok=false).
+	body := []byte(`{"result": {"usage": {"input_tokens": 0, "output_tokens": 0}}}`)
+	in, out := parseUsageFromA2AResponse(body)
+	if in != 0 || out != 0 {
+		t.Errorf("parseUsageFromA2AResponse = (%d, %d), want (0, 0)", in, out)
+	}
+}
+
+// ──────────────────────────────────────────────────────────────────────────────
+// readUsageMap
+// ──────────────────────────────────────────────────────────────────────────────
+
+func TestReadUsageMap_HappyPath(t *testing.T) {
+	m := map[string]json.RawMessage{
+		"usage": json.RawMessage(`{"input_tokens": 100, "output_tokens": 50}`),
+	}
+	in, out, ok := readUsageMap(m)
+	if !ok {
+		t.Fatal("readUsageMap returned ok=false, want true")
+	}
+	if in != 100 || out != 50 {
+		t.Errorf("readUsageMap = (%d, %d, %v), want (100, 50, true)", in, out, ok)
+	}
+}
+
+func TestReadUsageMap_MissingUsage(t *testing.T) {
+	m := map[string]json.RawMessage{
+		"other": json.RawMessage(`{}`),
+	}
+	in, out, ok := readUsageMap(m)
+	if ok {
+		t.Errorf("readUsageMap returned ok=true for missing usage, want false")
+	}
+}
+
+func TestReadUsageMap_ZeroValues(t *testing.T) {
+	m := map[string]json.RawMessage{
+		"usage": json.RawMessage(`{"input_tokens": 0, "output_tokens": 0}`),
+	}
+	in, out, ok := readUsageMap(m)
+	if ok {
+		t.Errorf("readUsageMap returned ok=true for zero usage, want false")
+	}
+	if in != 0 || out != 0 {
+		t.Errorf("readUsageMap = (%d, %d, %v), want (0, 0, false)", in, out, ok)
+	}
+}
+
+func TestReadUsageMap_OnlyInputTokens(t *testing.T) {
+	m := map[string]json.RawMessage{
+		"usage": json.RawMessage(`{"input_tokens": 200, "output_tokens": 0}`),
+	}
+	in, out, ok := readUsageMap(m)
+	if !ok {
+		t.Fatal("readUsageMap returned ok=false, want true")
+	}
+	if in != 200 || out != 0 {
+		t.Errorf("readUsageMap = (%d, %d, %v), want (200, 0, true)", in, out, ok)
+	}
+}
+
+func TestReadUsageMap_OnlyOutputTokens(t *testing.T) {
+	m := map[string]json.RawMessage{
+		"usage": json.RawMessage(`{"input_tokens": 0, "output_tokens": 150}`),
+	}
+	in, out, ok := readUsageMap(m)
+	if !ok {
+		t.Fatal("readUsageMap returned ok=false, want true")
+	}
+	if in != 0 || out != 150 {
+		t.Errorf("readUsageMap = (%d, %d, %v), want (0, 150, true)", in, out, ok)
+	}
+}
+
+func TestReadUsageMap_MalformedUsageJSON(t *testing.T) {
+	m := map[string]json.RawMessage{
+		"usage": json.RawMessage(`not valid json`),
+	}
+	in, out, ok := readUsageMap(m)
+	if ok {
+		t.Errorf("readUsageMap returned ok=true for malformed usage JSON, want false")
+	}
+}
+
 // Suppress unused import warning — setupTestDB references db.DB but this file
 // only tests pure functions, so db is only needed transitively through helpers.
 var _ = db.DB

workspace-server/internal/handlers/a2a_queue_test.go

@@ -80,6 +80,54 @@ func TestExtractIdempotencyKey_emptyOnMissing(t *testing.T) {
 	}
 }
 
+// ──────────────────────────────────────────────────────────────────────────────
+// extractExpiresInSeconds
+// ──────────────────────────────────────────────────────────────────────────────
+
+func TestExtractExpiresInSeconds_valid(t *testing.T) {
+	cases := []struct {
+		name string
+		body string
+		want int
+	}{
+		{"positive int", `{"params":{"expires_in_seconds":30}}`, 30},
+		{"zero", `{"params":{"expires_in_seconds":0}}`, 0},
+		{"large TTL", `{"params":{"expires_in_seconds":3600}}`, 3600},
+		{"nested message — not affected", `{"params":{"message":{"role":"user"},"expires_in_seconds":60}}`, 60},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			if got := extractExpiresInSeconds([]byte(tc.body)); got != tc.want {
+				t.Errorf("extractExpiresInSeconds = %d, want %d", got, tc.want)
+			}
+		})
+	}
+}
+
+func TestExtractExpiresInSeconds_invalidOrMissing(t *testing.T) {
+	cases := []struct {
+		name string
+		body string
+		want int
+	}{
+		{"negative → 0", `{"params":{"expires_in_seconds":-5}}`, 0},
+		{"missing expires_in_seconds", `{"params":{"message":{"role":"user"}}}`, 0},
+		{"no params at all", `{"method":"message/send"}`, 0},
+		{"malformed JSON", `not json`, 0},
+		{"empty body", ``, 0},
+		{"null value", `{"params":{"expires_in_seconds":null}}`, 0},
+		{"string value", `{"params":{"expires_in_seconds":"30"}}`, 0},
+		{"float value", `{"params":{"expires_in_seconds":30.5}}`, 0},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			if got := extractExpiresInSeconds([]byte(tc.body)); got != tc.want {
+				t.Errorf("extractExpiresInSeconds(%q) = %d, want %d", tc.body, got, tc.want)
+			}
+		})
+	}
+}
+
 func TestExtractDelegationIDFromBody(t *testing.T) {
 	cases := []struct {
 		name string