fix(batches): resolve a legacy row's org from the key's organization_id and keep the Batch label on grouped rows

This commit is contained in:
mateo-berri 2026-09-04 12:42:39 -07:00
parent 7e51fbc819
commit 8debf8294c
6 changed files with 22 additions and 14 deletions

View file

@ -158,7 +158,7 @@ class CheckBatchCost:
where={"token": api_key}
)
)
key_org_id = getattr(key_row, "org_id", None) if key_row is not None else None
key_org_id = getattr(key_row, "organization_id", None) if key_row is not None else None
if key_org_id:
return key_org_id
except Exception as e:

View file

@ -2561,7 +2561,7 @@ class TestBatchCostAttribution:
from types import SimpleNamespace
instance = self._instance(
key_row=SimpleNamespace(key_alias="prod-key", org_id="org-moved-to"),
key_row=SimpleNamespace(key_alias="prod-key", organization_id="org-moved-to"),
team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"),
)
@ -2578,7 +2578,7 @@ class TestBatchCostAttribution:
from types import SimpleNamespace
instance = self._instance(
key_row=SimpleNamespace(key_alias="prod-key", org_id="org-42"),
key_row=SimpleNamespace(key_alias="prod-key", organization_id="org-42"),
team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"),
)
@ -2593,7 +2593,7 @@ class TestBatchCostAttribution:
from types import SimpleNamespace
instance = self._instance(
key_row=SimpleNamespace(key_alias="prod-key", org_id=None),
key_row=SimpleNamespace(key_alias="prod-key", organization_id=None),
team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"),
)
@ -2625,7 +2625,7 @@ class TestBatchCostAttribution:
from types import SimpleNamespace
instance = self._instance(
key_row=SimpleNamespace(key_alias="prod-key", org_id=None),
key_row=SimpleNamespace(key_alias="prod-key", organization_id=None),
team_row=SimpleNamespace(team_alias="Team Alpha", organization_id=None),
)

View file

@ -141,6 +141,19 @@ describe("Type column", () => {
expect(screen.getByText("Batch")).toBeInTheDocument();
expect(screen.queryByText("LLM")).not.toBeInTheDocument();
});
it("keeps the Batch label on the grouped create-plus-cost session instead of a row count", () => {
const groupedCostRow: Partial<LogEntry> = {
request_id: "batch_1_batch_cost",
call_type: "aretrieve_batch",
session_id: "batch_1",
session_total_count: 2,
};
renderRows([logEntry(groupedCostRow)]);
expect(screen.getByText("Batch")).toBeInTheDocument();
expect(screen.queryByText("2")).not.toBeInTheDocument();
});
});
describe("batch rows", () => {

View file

@ -65,7 +65,7 @@ export const getRequestLogsTableColumns = ({
const sessionMcpCount = log.mcp_tool_call_count ?? (isMcp ? sessionCount : 0);
if (isBatchCallType(log.call_type)) {
return <BatchBadge count={sessionCount > 1 ? sessionCount : undefined} />;
return <BatchBadge />;
}
if (sessionCount <= 1) {
if (isMcp) return <McpBadge />;

View file

@ -45,14 +45,9 @@ describe("TypeBadges", () => {
});
describe("BatchBadge", () => {
it("should render with default 'Batch' text when no count is provided", () => {
it("should render 'Batch'", () => {
render(<BatchBadge />);
expect(screen.getByText("Batch")).toBeInTheDocument();
});
it("should render the count when provided", () => {
render(<BatchBadge count={4} />);
expect(screen.getByText("4")).toBeInTheDocument();
});
});
});

View file

@ -97,9 +97,9 @@ export const AgentBadge = ({ count }: { count?: number }) => (
</span>
);
export const BatchBadge = ({ count }: { count?: number }) => (
export const BatchBadge = () => (
<span className="inline-flex items-center gap-1 px-2 py-0.5 bg-teal-50 text-teal-700 border border-teal-200 rounded-full text-[11px] font-medium whitespace-nowrap dark:bg-teal-950 dark:text-teal-300 dark:border-teal-800">
<LayersIcon />
{count != null ? count : "Batch"}
Batch
</span>
);