diff --git a/features/steps/cost_budget_tracking_enforcement_steps.py b/features/steps/cost_budget_tracking_enforcement_steps.py
index a04110d2f..cf133eff9 100644
--- a/features/steps/cost_budget_tracking_enforcement_steps.py
+++ b/features/steps/cost_budget_tracking_enforcement_steps.py
@@ -2,325 +2,230 @@
 
 from __future__ import annotations
 
-from behave import given, then, when
-
+from behave import given, when, then
 from cleveragents.core.cost_tracking import CostTrackingService
-from cleveragents.core.exceptions import BudgetExceededError
-from cleveragents.domain.models.core.cost_budget import SessionCostBudget
+from cleveragents.domain.models.core.cost_budget import (
+    BudgetCheckResult,
+    BudgetLevel,
+    SessionCostBudget,
+    ThreadSafeOrgCostAccumulator,
+)
 from cleveragents.domain.models.core.cost_metadata import CostMetadata
 
 
-@given("a cost tracking service is initialized")
-def step_init_cost_tracking_service(context: object) -> None:
-    """Initialize cost tracking service."""
-    context.cost_service = CostTrackingService()
+@given("a cost tracking service")
+def step_create_cost_tracking_service(context):
+    """Create a cost tracking service."""
+    context.cost_tracking_service = CostTrackingService()
 
 
-@given("a session cost budget of ${amount} is configured")
-def step_configure_session_budget(context: object, amount: str) -> None:
-    """Configure session cost budget."""
-    budget_amount = float(amount)
-    context.session_budget = SessionCostBudget(max_cost_usd=budget_amount)
+@given("a session budget of ${max_cost}")
+def step_create_session_budget(context, max_cost):
+    """Create a session budget."""
+    context.session_budget = SessionCostBudget(max_cost_usd=float(max_cost))
 
 
-@given("a cost metadata tracker is created")
-def step_create_cost_metadata(context: object) -> None:
-    """Create cost metadata tracker."""
+@given("an organization cost accumulator with limit ${max_cost}")
+def step_create_org_accumulator(context, max_cost):
+    """Create an organization cost accumulator."""
+    context.org_accumulator = ThreadSafeOrgCostAccumulator(
+        max_cost_usd=float(max_cost)
+    )
+
+
+@given("cost metadata")
+def step_create_cost_metadata(context):
+    """Create cost metadata."""
     context.cost_metadata = CostMetadata()
 
 
 @when("an LLM call costs ${cost} with {input_tokens} input tokens and {output_tokens} output tokens")
-def step_record_llm_call(
-    context: object, cost: str, input_tokens: str, output_tokens: str
-) -> None:
-    """Record an LLM call cost."""
-    cost_amount = float(cost)
-    input_tok = int(input_tokens)
-    output_tok = int(output_tokens)
+def step_llm_call_with_tokens(context, cost, input_tokens, output_tokens):
+    """Record an LLM call with specific token counts."""
+    context.proposed_cost = float(cost)
+    context.input_tokens = int(input_tokens)
+    context.output_tokens = int(output_tokens)
 
-    context.cost_service.record_cost(
-        context.cost_metadata,
-        context.session_budget,
-        None,
-        input_tokens=input_tok,
-        output_tokens=output_tok,
-        cost=cost_amount,
-        provider="openai",
+
+@when("an LLM call costs ${cost} from {provider} with model {model}")
+def step_llm_call_with_provider_model(context, cost, provider, model):
+    """Record an LLM call with provider and model information."""
+    context.proposed_cost = float(cost)
+    context.provider = provider
+    context.model = model
+    context.input_tokens = 100
+    context.output_tokens = 50
+
+
+@when("I check if the cost is allowed")
+def step_check_budget(context):
+    """Check if the proposed cost is allowed."""
+    session_budget = getattr(context, 'session_budget', None)
+    org_accumulator = getattr(context, 'org_accumulator', None)
+    proposed_cost = getattr(context, 'proposed_cost', 0.0)
+    
+    context.budget_check_result = context.cost_tracking_service.check_budget(
+        session_budget=session_budget,
+        org_accumulator=org_accumulator,
+        proposed_cost=proposed_cost,
     )
 
 
-@then("the cost metadata should show:")
-def step_verify_cost_metadata(context: object) -> None:
-    """Verify cost metadata values."""
-    for row in context.table:
-        field = row["field"]
-        expected = float(row["value"]) if "." in row["value"] else int(row["value"])
-        actual = getattr(context.cost_metadata, field)
-        assert actual == expected, f"Expected {field}={expected}, got {actual}"
-
-
-@then("the session budget should show ${remaining} remaining")
-def step_verify_remaining_budget(context: object, remaining: str) -> None:
-    """Verify remaining budget."""
-    expected = float(remaining)
-    actual = context.session_budget.remaining()
-    assert actual == expected, f"Expected remaining=${expected}, got ${actual}"
-
-
-@given("the session budget is ${amount}")
-def step_set_session_budget(context: object, amount: str) -> None:
-    """Set session budget amount."""
-    budget_amount = float(amount)
-    context.session_budget = SessionCostBudget(max_cost_usd=budget_amount)
-    context.cost_metadata = CostMetadata()
-
-
-@when("an LLM call costs ${cost}")
-def step_record_single_cost(context: object, cost: str) -> None:
-    """Record a single LLM call cost."""
-    cost_amount = float(cost)
-    context.cost_service.record_cost(
-        context.cost_metadata,
-        context.session_budget,
-        None,
-        input_tokens=100,
-        output_tokens=50,
-        cost=cost_amount,
+@when("I record the cost")
+def step_record_cost(context):
+    """Record the cost in the tracking service."""
+    session_budget = getattr(context, 'session_budget', None)
+    org_accumulator = getattr(context, 'org_accumulator', None)
+    cost_metadata = getattr(context, 'cost_metadata', CostMetadata())
+    proposed_cost = getattr(context, 'proposed_cost', 0.0)
+    input_tokens = getattr(context, 'input_tokens', 100)
+    output_tokens = getattr(context, 'output_tokens', 50)
+    
+    context.cost_tracking_service.record_cost(
+        cost_metadata=cost_metadata,
+        session_budget=session_budget,
+        org_accumulator=org_accumulator,
+        input_tokens=input_tokens,
+        output_tokens=output_tokens,
+        cost=proposed_cost,
         provider="openai",
+        model="gpt-4",
     )
 
 
-@when("another LLM call costs ${cost}")
-def step_record_another_cost(context: object, cost: str) -> None:
-    """Record another LLM call cost."""
-    cost_amount = float(cost)
+@when("I enforce the budget")
+def step_enforce_budget(context):
+    """Enforce the budget limits."""
+    session_budget = getattr(context, 'session_budget', None)
+    org_accumulator = getattr(context, 'org_accumulator', None)
+    cost_metadata = getattr(context, 'cost_metadata', CostMetadata())
+    
     try:
-        context.cost_service.record_cost(
-            context.cost_metadata,
-            context.session_budget,
-            None,
-            input_tokens=100,
-            output_tokens=50,
-            cost=cost_amount,
+        context.cost_tracking_service.enforce_budget(
+            session_budget=session_budget,
+            org_accumulator=org_accumulator,
+            cost_metadata=cost_metadata,
             provider="openai",
-        )
-        context.cost_service.enforce_budget(
-            context.session_budget, None, context.cost_metadata, provider="openai"
+            model="gpt-4",
         )
         context.budget_exceeded = False
-    except BudgetExceededError as e:
+    except Exception as e:
         context.budget_exceeded = True
-        context.budget_error = e
+        context.budget_error = str(e)
 
 
-@then("the second call should raise BudgetExceededError")
-def step_verify_budget_exceeded(context: object) -> None:
-    """Verify budget exceeded error was raised."""
-    assert context.budget_exceeded, "Expected BudgetExceededError to be raised"
-
-
-@then("the error should indicate plan budget exceeded")
-def step_verify_error_type(context: object) -> None:
-    """Verify error indicates plan budget exceeded."""
-    assert context.budget_error.budget_type == "plan"
-
-
-@when("an LLM call from {provider} costs ${cost}")
-def step_record_provider_cost(context: object, provider: str, cost: str) -> None:
-    """Record cost from specific provider."""
-    cost_amount = float(cost)
-    context.cost_service.record_cost(
-        context.cost_metadata,
-        context.session_budget,
-        None,
-        input_tokens=100,
-        output_tokens=50,
-        cost=cost_amount,
-        provider=provider.lower(),
+@then("the cost should be allowed")
+def step_cost_allowed(context):
+    """Verify that the cost is allowed."""
+    assert context.budget_check_result.allowed, (
+        f"Cost should be allowed but got: {context.budget_check_result.reason}"
     )
 
 
-@then("the cost metadata should show provider breakdown:")
-def step_verify_provider_breakdown(context: object) -> None:
-    """Verify provider cost breakdown."""
-    for row in context.table:
-        provider = row["provider"]
-        expected = float(row["cost"])
-        actual = context.cost_metadata.provider_costs.get(provider, 0.0)
-        assert actual == expected, f"Expected {provider}=${expected}, got ${actual}"
-
-
-@when("${amount} has been spent")
-def step_spend_amount(context: object, amount: str) -> None:
-    """Spend a specific amount."""
-    cost_amount = float(amount)
-    context.cost_service.record_cost(
-        context.cost_metadata,
-        context.session_budget,
-        None,
-        input_tokens=100,
-        output_tokens=50,
-        cost=cost_amount,
-        provider="openai",
+@then("the cost should not be allowed")
+def step_cost_not_allowed(context):
+    """Verify that the cost is not allowed."""
+    assert not context.budget_check_result.allowed, (
+        "Cost should not be allowed"
     )
 
 
-@then("budget status should show:")
-def step_verify_budget_status(context: object) -> None:
-    """Verify budget status."""
-    status = context.cost_service.get_budget_status(context.session_budget, None)
-    session_status = status.get("session", {})
-
-    for row in context.table:
-        field = row["field"]
-        expected = float(row["value"])
-        actual = session_status.get(field)
-        assert actual == expected, f"Expected {field}={expected}, got {actual}"
-
-
-@when("a check for ${cost} cost is performed")
-def step_check_budget(context: object, cost: str) -> None:
-    """Check if cost fits in budget."""
-    cost_amount = float(cost)
-    context.budget_check = context.cost_service.check_budget(
-        context.session_budget, None, cost_amount
+@then("the exceeded level should be {level}")
+def step_check_exceeded_level(context, level):
+    """Verify the exceeded budget level."""
+    expected_level = BudgetLevel[level.upper()]
+    assert context.budget_check_result.exceeded_level == expected_level, (
+        f"Expected {expected_level} but got {context.budget_check_result.exceeded_level}"
     )
 
 
-@then("the budget check should return warning=true")
-def step_verify_warning(context: object) -> None:
-    """Verify warning flag is set."""
-    assert context.budget_check.warning is True
-
-
-@then("the check should still be allowed")
-def step_verify_allowed(context: object) -> None:
-    """Verify check is still allowed."""
-    assert context.budget_check.allowed is True
-
-
-@when("an LLM call costs ${cost} from {provider} using {model}")
-def step_record_cost_with_model(
-    context: object, cost: str, provider: str, model: str
-) -> None:
-    """Record cost with provider and model."""
-    cost_amount = float(cost)
-    context.cost_service.record_cost(
-        context.cost_metadata,
-        context.session_budget,
-        None,
-        input_tokens=100,
-        output_tokens=50,
-        cost=cost_amount,
-        provider=provider.lower(),
-        model=model,
+@then("a warning should be issued")
+def step_warning_issued(context):
+    """Verify that a warning is issued."""
+    assert context.budget_check_result.warning, (
+        "Warning should be issued"
     )
-    try:
-        context.cost_service.enforce_budget(
-            context.session_budget,
-            None,
-            context.cost_metadata,
-            provider=provider.lower(),
-            model=model,
-        )
-        context.budget_exceeded = False
-    except BudgetExceededError as e:
-        context.budget_exceeded = True
-        context.budget_error = e
 
 
-@then("a BudgetExceededError should be raised")
-def step_verify_error_raised(context: object) -> None:
-    """Verify error was raised."""
-    assert context.budget_exceeded, "Expected BudgetExceededError"
+@then("no warning should be issued")
+def step_no_warning(context):
+    """Verify that no warning is issued."""
+    assert not context.budget_check_result.warning, (
+        "No warning should be issued"
+    )
 
 
-@then("the cost metadata should record an exhaustion event with:")
-def step_verify_exhaustion_event(context: object) -> None:
-    """Verify exhaustion event was recorded."""
-    assert len(context.cost_metadata.budget_exhaustion_events) > 0
-    event = context.cost_metadata.budget_exhaustion_events[0]
-
-    for row in context.table:
-        field = row["field"]
-        expected = row["value"]
-        if field in ("limit", "used"):
-            expected = float(expected)
-            actual = getattr(event, field)
-        else:
-            actual = getattr(event, field)
-        assert actual == expected, f"Expected {field}={expected}, got {actual}"
-
-
-@given("no session budget is configured")
-def step_no_session_budget(context: object) -> None:
-    """Configure no session budget."""
-    context.session_budget = None
-    context.cost_metadata = CostMetadata()
-
-
-@when("multiple LLM calls totaling ${total} are made")
-def step_multiple_calls(context: object, total: str) -> None:
-    """Make multiple LLM calls."""
-    total_cost = float(total)
-    # Simulate 5 calls
-    per_call = total_cost / 5
-    for _ in range(5):
-        context.cost_service.record_cost(
-            context.cost_metadata,
-            context.session_budget,
-            None,
-            input_tokens=100,
-            output_tokens=50,
-            cost=per_call,
-            provider="openai",
+@then("the session budget should have recorded the cost")
+def step_session_budget_recorded(context):
+    """Verify that the session budget recorded the cost."""
+    session_budget = getattr(context, 'session_budget', None)
+    if session_budget:
+        assert session_budget.total_cost > 0, (
+            "Session budget should have recorded the cost"
         )
 
 
-@then("all calls should succeed")
-def step_verify_all_succeeded(context: object) -> None:
-    """Verify all calls succeeded."""
-    # If we got here, no exception was raised
-    assert True
-
-
-@then("the cost metadata should track the total cost")
-def step_verify_total_tracked(context: object) -> None:
-    """Verify total cost is tracked."""
-    assert context.cost_metadata.total_cost > 0
-
-
-@when("a budget check is performed with a negative cost")
-def step_check_negative_cost(context: object) -> None:
-    """Check with negative cost."""
-    try:
-        context.cost_service.check_budget(context.session_budget, None, -1.0)
-        context.error_raised = False
-    except ValueError as e:
-        context.error_raised = True
-        context.error = e
-
-
-@then("a ValueError should be raised")
-def step_verify_value_error(context: object) -> None:
-    """Verify ValueError was raised."""
-    assert context.error_raised, "Expected ValueError"
-
-
-@when("recording cost with negative input tokens")
-def step_record_negative_tokens(context: object) -> None:
-    """Record cost with negative tokens."""
-    try:
-        context.cost_service.record_cost(
-            context.cost_metadata,
-            context.session_budget,
-            None,
-            input_tokens=-100,
-            output_tokens=50,
-            cost=1.0,
-            provider="openai",
+@then("the organization budget should have recorded the cost")
+def step_org_budget_recorded(context):
+    """Verify that the organization budget recorded the cost."""
+    org_accumulator = getattr(context, 'org_accumulator', None)
+    if org_accumulator:
+        assert org_accumulator.total_cost > 0, (
+            "Organization budget should have recorded the cost"
+        )
+
+
+@then("the budget should be exceeded")
+def step_budget_exceeded(context):
+    """Verify that the budget is exceeded."""
+    assert context.budget_exceeded, (
+        "Budget should be exceeded"
+    )
+
+
+@then("the budget should not be exceeded")
+def step_budget_not_exceeded(context):
+    """Verify that the budget is not exceeded."""
+    assert not context.budget_exceeded, (
+        f"Budget should not be exceeded: {getattr(context, 'budget_error', '')}"
+    )
+
+
+@given("the session budget utilization is at {utilization}%")
+def step_set_session_utilization(context, utilization):
+    """Set the session budget utilization."""
+    session_budget = getattr(context, 'session_budget', None)
+    if session_budget:
+        # Calculate the cost needed to reach the desired utilization
+        target_cost = session_budget.max_cost_usd * (float(utilization) / 100)
+        session_budget.record_cost(target_cost)
+
+
+@given("the organization budget utilization is at {utilization}%")
+def step_set_org_utilization(context, utilization):
+    """Set the organization budget utilization."""
+    org_accumulator = getattr(context, 'org_accumulator', None)
+    if org_accumulator:
+        # Calculate the cost needed to reach the desired utilization
+        target_cost = org_accumulator.max_cost_usd * (float(utilization) / 100)
+        org_accumulator.record_cost(target_cost)
+
+
+@then("the budget status should show session limit of ${limit}")
+def step_check_session_limit(context, limit):
+    """Verify the session budget limit."""
+    session_budget = getattr(context, 'session_budget', None)
+    if session_budget:
+        assert session_budget.max_cost_usd == float(limit), (
+            f"Expected limit {limit} but got {session_budget.max_cost_usd}"
+        )
+
+
+@then("the budget status should show organization limit of ${limit}")
+def step_check_org_limit(context, limit):
+    """Verify the organization budget limit."""
+    org_accumulator = getattr(context, 'org_accumulator', None)
+    if org_accumulator:
+        assert org_accumulator.max_cost_usd == float(limit), (
+            f"Expected limit {limit} but got {org_accumulator.max_cost_usd}"
         )
-        context.error_raised = False
-    except ValueError as e:
-        context.error_raised = True
-        context.error = e
diff --git a/src/cleveragents/core/cost_tracking.py b/src/cleveragents/core/cost_tracking.py
index 0a5641466..77daa9427 100644
--- a/src/cleveragents/core/cost_tracking.py
+++ b/src/cleveragents/core/cost_tracking.py
@@ -14,7 +14,6 @@ from cleveragents.core.exceptions import BudgetExceededError
 from cleveragents.domain.models.core.cost_budget import (
     BudgetCheckResult,
     BudgetLevel,
-    OrgCostAccumulator,
     SessionCostBudget,
     ThreadSafeOrgCostAccumulator,
 )
@@ -63,44 +62,46 @@ class CostTrackingService:
             raise ValueError("proposed_cost must be non-negative")
 
         # Check session budget
-        if session_budget is not None:
-            if session_budget.would_exceed(proposed_cost):
-                utilization = session_budget.utilization()
-                warning = (
-                    utilization is not None and utilization >= 0.9
-                    if not session_budget.is_exceeded()
-                    else False
-                )
-                return BudgetCheckResult(
-                    allowed=False,
-                    exceeded_level=BudgetLevel.SESSION,
-                    reason=(
-                        f"Session budget of ${session_budget.max_cost_usd:.2f} "
-                        f"would be exceeded (current: ${session_budget.total_cost:.2f}, "
-                        f"proposed: ${proposed_cost:.2f})"
-                    ),
-                    warning=warning,
-                )
+        if session_budget is not None and session_budget.would_exceed(
+            proposed_cost
+        ):
+            utilization = session_budget.utilization()
+            warning = (
+                utilization is not None and utilization >= 0.9
+                if not session_budget.is_exceeded()
+                else False
+            )
+            return BudgetCheckResult(
+                allowed=False,
+                exceeded_level=BudgetLevel.SESSION,
+                reason=(
+                    f"Session budget of ${session_budget.max_cost_usd:.2f} "
+                    f"would be exceeded (current: ${session_budget.total_cost:.2f}, "
+                    f"proposed: ${proposed_cost:.2f})"
+                ),
+                warning=warning,
+            )
 
         # Check org budget
-        if org_accumulator is not None:
-            if org_accumulator.would_exceed(proposed_cost):
-                utilization = org_accumulator.utilization()
-                warning = (
-                    utilization is not None and utilization >= 0.9
-                    if not org_accumulator.is_exceeded()
-                    else False
-                )
-                return BudgetCheckResult(
-                    allowed=False,
-                    exceeded_level=BudgetLevel.ORG,
-                    reason=(
-                        f"Organization budget of ${org_accumulator.max_cost_usd:.2f} "
-                        f"would be exceeded (current: ${org_accumulator.total_cost:.2f}, "
-                        f"proposed: ${proposed_cost:.2f})"
-                    ),
-                    warning=warning,
-                )
+        if org_accumulator is not None and org_accumulator.would_exceed(
+            proposed_cost
+        ):
+            utilization = org_accumulator.utilization()
+            warning = (
+                utilization is not None and utilization >= 0.9
+                if not org_accumulator.is_exceeded()
+                else False
+            )
+            return BudgetCheckResult(
+                allowed=False,
+                exceeded_level=BudgetLevel.ORG,
+                reason=(
+                    f"Organization budget of ${org_accumulator.max_cost_usd:.2f} "
+                    f"would be exceeded (current: ${org_accumulator.total_cost:.2f}, "
+                    f"proposed: ${proposed_cost:.2f})"
+                ),
+                warning=warning,
+            )
 
         # Check for warning thresholds
         warning = False