feat: implement error handling and retry mechanism for assistant messages (#11913)

## Pull Request

---

### 📄 Summary
> Why does this change exist?  
> What problem does it solve, and why is this the right approach?

The AI Assistant previously surfaced streaming failures with minimal
structure: a plain error string, no distinction between transient vs
permanent failures, and no way for users to recover without retyping or
refreshing. Backend SSE errors now carry `retryAction` (`auto` /
`manual` / `none`) and structured error codes, but the frontend was not
honoring that contract end-to-end.

This PR wires full error handling and retry into the assistant:

- **Centralized error resolution** — `resolveAssistantErrorMessage` is
replaced by `resolveAssistantError`, which maps backend codes to
user-friendly copy, classifies rate-limit vs non-retryable errors, and
derives the correct `retryAction`.
- **Dual retry budgets in the store** — `streamWithAuthRetry` becomes
`streamWithRetry`, handling auth expiry (one silent re-attempt) and
backend-flagged transient errors (up to 2 auto-retries with 500ms /
1500ms backoff). When auto-retries are exhausted, the error is
downgraded to `manual` so the user can still retry.
- **Manual retry UX** — failed turns commit as styled error bubbles
(`isError`, `errorCode`, `retryAction`). Manual errors show an inline
**Retry** button that replays the originating action (send, approve,
clarify, regenerate) without duplicating the user message. A transient
`retryRegistry` holds the replay thunk for the latest failed turn.
- **Analytics** — `RetryClicked` event fired when the user clicks Retry.

This aligns the UI with the backend error contract and gives users a
clear, actionable path to recover from transient failures.

#### Screenshots / Screen Recordings (if applicable)
> Include screenshots or screen recordings that clearly show the
behavior before the change and the result after the change. This helps
reviewers quickly understand the impact and verify the update.

| Before | After |
|--------|-------|
| Generic/unstructured error text in assistant bubble | Error callout
with warning icon, code-specific copy, and **Retry** button for manual
errors |
| No retry affordance | Retry replays the failed action; auto-retries
happen silently for transient errors |
| Feedback bar shown on errors | Feedback/regenerate hidden on error
bubbles; rate-limit errors still suppress retry |

_Add screenshots of: (1) `thread_busy` manual error with Retry, (2)
rate-limit error with no Retry, (3) successful recovery after Retry._

#### Issues closed by this PR
> Reference issues using `Closes #issue-number` to enable automatic
closure on merge.

Fixes:  https://github.com/SigNoz/nerve-pod/issues/92

---

### ✅ Change Type
_Select all that apply_

- [x] ✨ Feature
- [ ] 🐛 Bug fix
- [ ] ♻️ Refactor
- [ ] 🛠️ Infra / Tooling
- [x] 🧪 Test-only

---

### 🐛 Bug Context
> Required if this PR fixes a bug

N/A — this is primarily a feature/enhancement to error handling UX, not
a targeted bug fix.

---

### 🧪 Testing Strategy
> How was this change validated?

- Tests added/updated:
- `resolveAssistantError.test.ts` — error code copy, rate-limit
classification, non-retryable codes, HTTP/SSE error shapes,
`retryAction` derivation
- `useAIAssistantStore.test.ts` — manual error bubbles, manual retry
replay (send + approve), auto-retry with backoff and downgrade to
manual, silent recovery on auto-retry success, rate-limit errors with no
retry
- `MessageBubble.test.tsx` — error styling, Retry button visibility
(`manual` vs `none`), `onRetry` callback, feedback bar suppressed on
errors
- Removed `resolveAssistantErrorMessage.test.ts` (superseded by
`resolveAssistantError`)

- Manual verification:
- Trigger `thread_busy` during an active execution → error bubble with
Retry → click Retry → successful response without duplicate user message
- Trigger rate-limit error → error bubble, no Retry button, no feedback
bar
- Trigger transient `internal_error` with `retryAction: auto` → silent
retries; if all fail, manual Retry appears

- Edge cases covered:
- Auth expiry mid-stream (`invalid_token`) — one auth retry via
`streamWithRetry`
  - Auto-retry budget exhaustion (2 attempts) → manual Retry affordance
- Retry replays originating action for approve/clarify/regenerate, not
just send
  - Retry no-op when `retryAction: none` or no registry entry
  - Retry disabled while streaming is active

---

### ⚠️ Risk & Impact Assessment
> What could break? How do we recover?

- Blast radius: AI Assistant chat only — `useAIAssistantStore`,
`MessageBubble`, `VirtualizedMessages`, error utils
- Potential regressions:
- Error copy regressions if a new backend code is not in
`ERROR_CODE_COPY` (falls back to backend message)
- Auto-retry backoff may add latency (up to ~2s) before surfacing a
manual error for transient failures
- `retryRegistry` is in-memory only — page reload drops retry capability
(acceptable; user can resend)
- Rollback plan: Revert PR; no schema/migration changes

---

### 📝 Changelog
> Fill only if this affects users, APIs, UI, or documented behavior  
> Use **N/A** for internal or non-user-facing changes

| Field | Value |
|------|-------|
| Deployment Type | Cloud / OSS / Enterprise |
| Change Type | Feature |
| Description | AI Assistant now shows clearer error messages with
inline Retry for recoverable failures, and silently retries transient
backend errors before asking the user to retry. |

---

### 📋 Checklist
- [x] Tests added or explicitly not required
- [ ] Manually tested
- [ ] Breaking changes documented
- [x] Backward compatibility considered

---

## 👀 Notes for Reviewers

<!-- Anything reviewers should keep in mind while reviewing -->

- **`resolveAssistantErrorMessage.ts` → `resolveAssistantError.ts`**:
The new util returns a full `AssistantErrorResolution` object instead of
just a string. All call sites in the store were updated accordingly.
- **`streamWithRetry`**: Auth retry (1×) and auto retry (2× with
backoff) are independent budgets. When auto retries are spent,
`retryAction` is forced to `manual` before the error propagates to
`finalizeStreamingError`.
- **`retryRegistry`**: Transient map keyed by `conversationId`; not
persisted. Pairs with the error bubble's lifetime.
- **Error bubble UI**: Uses `@signozhq/ui` `Button` and
`@signozhq/icons` (`TriangleAlert`, `RotateCw`). Feedback/regenerate bar
is hidden on `isError` messages (same as rate-limit).
- **Files touched (12)**: store, error util, MessageBubble (+ styles),
VirtualizedMessages, types, events, 3 new test files, 1 removed test
file.

---

docs/api/openapi.yml

@@ -1771,12 +1771,15 @@ components:
           additionalProperties: {}
           nullable: true
           type: object
+        syncState:
+          $ref: '#/components/schemas/CloudintegrationtypesSyncState'
         timestampMillis:
           format: int64
           type: integer
       required:
       - timestampMillis
       - data
+      - syncState
       type: object
     CloudintegrationtypesAzureAccountConfig:
       properties:
@@ -2021,6 +2024,8 @@ components:
           format: date-time
           nullable: true
           type: string
+        syncState:
+          $ref: '#/components/schemas/CloudintegrationtypesSyncState'
       required:
       - account_id
       - cloud_account_id
@@ -2030,6 +2035,7 @@ components:
       - providerAccountId
       - integrationConfig
       - removedAt
+      - syncState
       type: object
     CloudintegrationtypesGettableServicesMetadata:
       properties:
@@ -2129,6 +2135,9 @@ components:
           type: object
         providerAccountId:
           type: string
+        syncedVersion:
+          nullable: true
+          type: integer
       required:
       - data
       type: object
@@ -2141,6 +2150,18 @@ components:
         gcp:
           $ref: '#/components/schemas/CloudintegrationtypesGCPIntegrationConfig'
       type: object
+    CloudintegrationtypesRegionState:
+      enum:
+      - enabled
+      - disabled
+      type: string
+    CloudintegrationtypesRegionSyncState:
+      properties:
+        state:
+          $ref: '#/components/schemas/CloudintegrationtypesRegionState'
+      required:
+      - state
+      type: object
     CloudintegrationtypesService:
       properties:
         assets:
@@ -2282,6 +2303,23 @@ components:
         metrics:
           type: boolean
       type: object
+    CloudintegrationtypesSyncState:
+      nullable: true
+      properties:
+        inSync:
+          type: boolean
+        regions:
+          additionalProperties:
+            $ref: '#/components/schemas/CloudintegrationtypesRegionSyncState'
+          type: object
+        version:
+          format: int64
+          type: integer
+      required:
+      - version
+      - inSync
+      - regions
+      type: object
     CloudintegrationtypesUpdatableAccount:
       properties:
         config:

docs/contributing/go/authz.md

@@ -118,7 +118,7 @@ router.Handle("/api/v1/service_accounts", handler.New(
 The pieces:
 
 - **`CheckResources(handlerFn, roles...)`** — the resource-aware authorization wrapper from [pkg/http/middleware/authz.go](/pkg/http/middleware/authz.go). The role list is the community-edition fallback: which managed roles may call this route when per-resource checks are unavailable.
-- **`ResourceDef`** — declares the resource, verb, audit category, how to extract the instance ID, and how to turn that ID into selectors. ID extractors live in [pkg/types/coretypes/extractor.go](/pkg/types/coretypes/extractor.go): `PathParam("id")`, `BodyField(func(req *T) string)` / `BodyFields(func(req *T) []string)` reading the request body the resource middleware decoded into the route's `OpenAPIDef.Request` type `T`, and `ResponseJSONPath("data.id")` for IDs only known after the handler runs (e.g. `create`). A handler on such a route reads the same decoded value with `coretypes.BodyFromContext[T](r.Context())`.
+- **`ResourceDef`** — declares the resource, verb, audit category, how to extract the instance ID, and how to turn that ID into selectors. ID extractors live in [pkg/types/coretypes/extractor.go](/pkg/types/coretypes/extractor.go): `PathParam("id")`, `BodyJSONPath("data.id")`, `BodyJSONArray("ids")`, and `ResponseJSONPath("data.id")` for IDs only known after the handler runs (e.g. `create`).
 - **`SecuritySchemes`** — advertises the required scope (`resource.Scope(verb)`, e.g. `serviceaccount:create`) in the OpenAPI spec.
 
 For routes that link two resources, use `AttachDetachSiblingResourceDef` (both sides are authz-checked, e.g. attaching a role to a service account requires `attach` on **both** the service account and the role). For parent-child routes (e.g. creating an API key under a service account), both sides are checked too, but with different verbs: declare a `BasicResourceDef` checking the child with `create`/`delete`, alongside an `AttachDetachParentChildResourceDef` checking the parent with `attach`/`detach` (within that def the child is only recorded for audit) — see the `/api/v1/service_accounts/{id}/keys` route in [pkg/apiserver/signozapiserver/serviceaccount.go](/pkg/apiserver/signozapiserver/serviceaccount.go).

ee/modules/cloudintegration/implcloudintegration/module.go

@@ -183,32 +183,52 @@ func (module *module) AgentCheckIn(ctx context.Context, orgID valuer.UUID, provi
 		return nil, errors.New(errors.TypeAlreadyExists, cloudintegrationtypes.ErrCodeCloudIntegrationAlreadyConnected, errMessage)
 	}
 
-	account, err := module.store.GetAccountByID(ctx, orgID, req.CloudIntegrationID, provider)
+	storableAccount, err := module.store.GetAccountByID(ctx, orgID, req.CloudIntegrationID, provider)
 	if err != nil {
 		return nil, err
 	}
 
+	account, err := cloudintegrationtypes.NewAccountFromStorable(storableAccount)
+	if err != nil {
+		return nil, err
+	}
+
+	syncState := account.NextSyncState(req.SyncedVersion)
+
 	// If account has been removed (disconnected), return a minimal response with empty integration config.
-	// The agent uses this response to clean up resources
 	if account.RemovedAt != nil {
+		// Heartbeat stays frozen after removal, only the sync state is updated.
+		if account.AgentReport != nil && syncState != nil {
+			account.UpdateSyncState(syncState)
+
+			storableAccount, err = cloudintegrationtypes.NewStorableCloudIntegration(account)
+			if err != nil {
+				return nil, err
+			}
+
+			err = module.store.UpdateAgentReport(ctx, storableAccount)
+			if err != nil {
+				return nil, err
+			}
+		}
+
 		return cloudintegrationtypes.NewAgentCheckInResponse(
 			req.ProviderAccountID,
 			account.ID.StringValue(),
 			new(cloudintegrationtypes.ProviderIntegrationConfig),
 			account.RemovedAt,
+			syncState,
 		), nil
 	}
 
-	// update account with cloud provider account id and agent report (heartbeat)
-	account.Update(&req.ProviderAccountID, cloudintegrationtypes.NewAgentReport(req.Data))
+	account.UpdateAgentReport(&req.ProviderAccountID, cloudintegrationtypes.NewAgentReport(req.Data, syncState))
 
-	err = module.store.UpdateAccount(ctx, account)
+	storableAccount, err = cloudintegrationtypes.NewStorableCloudIntegration(account)
 	if err != nil {
 		return nil, err
 	}
 
-	// Get account as domain object for config access (enabled regions, etc.)
-	domainAccount, err := cloudintegrationtypes.NewAccountFromStorable(account)
+	err = module.store.UpdateAgentReport(ctx, storableAccount)
 	if err != nil {
 		return nil, err
 	}
@@ -223,8 +243,7 @@ func (module *module) AgentCheckIn(ctx context.Context, orgID valuer.UUID, provi
 		return nil, err
 	}
 
-	// Delegate integration config building entirely to the provider module
-	integrationConfig, err := cloudProvider.BuildIntegrationConfig(ctx, domainAccount, storedServices)
+	integrationConfig, err := cloudProvider.BuildIntegrationConfig(ctx, account, storedServices)
 	if err != nil {
 		return nil, err
 	}
@@ -234,6 +253,7 @@ func (module *module) AgentCheckIn(ctx context.Context, orgID valuer.UUID, provi
 		account.ID.StringValue(),
 		integrationConfig,
 		account.RemovedAt,
+		syncState,
 	), nil
 }
 

frontend/src/api/generated/services/sigNoz.schemas.ts

@@ -3396,6 +3396,37 @@ export interface CloudintegrationtypesAWSServiceConfigDTO {
 	metrics?: CloudintegrationtypesAWSServiceMetricsConfigDTO;
 }
 
+export enum CloudintegrationtypesRegionStateDTO {
+	enabled = 'enabled',
+	disabled = 'disabled',
+}
+export interface CloudintegrationtypesRegionSyncStateDTO {
+	state: CloudintegrationtypesRegionStateDTO;
+}
+
+export type CloudintegrationtypesSyncStateDTORegions = {
+	[key: string]: CloudintegrationtypesRegionSyncStateDTO;
+};
+
+/**
+ * @nullable
+ */
+export type CloudintegrationtypesSyncStateDTO = {
+	/**
+	 * @type boolean
+	 */
+	inSync: boolean;
+	/**
+	 * @type object
+	 */
+	regions: CloudintegrationtypesSyncStateDTORegions;
+	/**
+	 * @type integer
+	 * @format int64
+	 */
+	version: number;
+} | null;
+
 export type CloudintegrationtypesAgentReportDTODataAnyOf = {
 	[key: string]: unknown;
 };
@@ -3414,6 +3445,7 @@ export type CloudintegrationtypesAgentReportDTO = {
 	 * @type object,null
 	 */
 	data: CloudintegrationtypesAgentReportDTOData;
+	syncState: CloudintegrationtypesSyncStateDTO | null;
 	/**
 	 * @type integer
 	 * @format int64
@@ -3842,6 +3874,7 @@ export interface CloudintegrationtypesGettableAgentCheckInDTO {
 	 * @format date-time
 	 */
 	removedAt: string | null;
+	syncState: CloudintegrationtypesSyncStateDTO | null;
 }
 
 export interface CloudintegrationtypesServiceMetadataDTO {
@@ -3912,6 +3945,10 @@ export interface CloudintegrationtypesPostableAgentCheckInDTO {
 	 * @type string
 	 */
 	providerAccountId?: string;
+	/**
+	 * @type integer,null
+	 */
+	syncedVersion?: number | null;
 }
 
 export interface CloudintegrationtypesStorableIntegrationDashboardDTO {

frontend/src/container/AIAssistant/components/MessageBubble/MessageBubble.module.scss

@@ -75,6 +75,38 @@
 	word-break: break-word;
 }
 
+// Error bubble: a subtle error-tinted callout replacing the default
+// assistant background, rendered when a turn fails.
+.bubble.error {
+	.assistant & {
+		background: var(--callout-error-background);
+		border: 1px solid var(--callout-error-border);
+	}
+}
+
+.errorContent {
+	display: flex;
+	align-items: flex-start;
+	gap: 8px;
+}
+
+.errorIcon {
+	flex-shrink: 0;
+	margin-top: 2px;
+	color: var(--destructive);
+}
+
+.errorText {
+	color: var(--callout-error-title);
+	white-space: pre-wrap;
+	word-break: break-word;
+}
+
+.retryButton {
+	margin-top: 6px;
+	align-self: flex-start;
+}
+
 // User-bubble row: pencil button sits to the LEFT of the bubble within
 // the right-aligned message line, so it visually "ends" at the bubble's
 // right edge while keeping the bubble in its original position.

frontend/src/container/AIAssistant/components/MessageBubble/MessageBubble.tsx

@@ -2,6 +2,10 @@ import React, { useMemo } from 'react';
 import cx from 'classnames';
 import ReactMarkdown from 'react-markdown';
 import remarkGfm from 'remark-gfm';
+import { Button } from '@signozhq/ui/button';
+import { RotateCw, TriangleAlert } from '@signozhq/icons';
+
+import { RetryActionDTO } from 'api/ai-assistant/sigNozAIAssistantAPI.schemas';
 
 // Side-effect: registers all built-in block types into the BlockRegistry
 import '../blocks';
@@ -104,18 +108,23 @@ function renderGroup(group: RenderGroup): JSX.Element {
 interface MessageBubbleProps {
 	message: Message;
 	onRegenerate?: () => void;
+	onRetry?: () => void;
 	isLastAssistant?: boolean;
 }
 
 export default function MessageBubble({
 	message,
 	onRegenerate,
+	onRetry,
 	isLastAssistant = false,
 }: MessageBubbleProps): JSX.Element {
 	const variant = useVariant();
 	const isCompact = variant === 'panel';
 	const isUser = message.role === 'user';
+	const isError = !isUser && Boolean(message.isError);
 	const hasBlocks = !isUser && message.blocks && message.blocks.length > 0;
+	const showRetry =
+		isError && message.retryAction === RetryActionDTO.manual && Boolean(onRetry);
 
 	// Recompute groups only when the blocks array identity changes — store
 	// updates that don't touch this message's blocks should not re-render the
@@ -138,7 +147,7 @@ export default function MessageBubble({
 		<div className={messageClass} data-testid={`ai-message-${message.id}`}>
 			<div className={bodyClass}>
 				<div className={styles.bubbleRow}>
-					<div className={styles.bubble}>
+					<div className={cx(styles.bubble, { [styles.error]: isError })}>
 						{message.attachments && message.attachments.length > 0 && (
 							<div className={styles.attachments}>
 								{message.attachments.map((att) => {
@@ -161,6 +170,11 @@ export default function MessageBubble({
 
 						{isUser ? (
 							<p className={styles.text}>{message.content}</p>
+						) : isError ? (
+							<div className={styles.errorContent}>
+								<TriangleAlert size={14} className={styles.errorIcon} />
+								<span className={styles.errorText}>{message.content}</span>
+							</div>
 						) : hasBlocks ? (
 							<MessageContext.Provider value={{ messageId: message.id }}>
 								{groups.map((g) => renderGroup(g))}
@@ -183,7 +197,21 @@ export default function MessageBubble({
 					</div>
 				</div>
 
-				{!isUser && !message.isRateLimitError && (
+				{showRetry && (
+					<Button
+						className={styles.retryButton}
+						size="sm"
+						variant="ghost"
+						color="secondary"
+						onClick={onRetry}
+						testId={`ai-message-retry-${message.id}`}
+					>
+						<RotateCw size={12} />
+						Retry
+					</Button>
+				)}
+
+				{!isUser && !isError && !message.isRateLimitError && (
 					<MessageFeedback
 						message={message}
 						onRegenerate={onRegenerate}

frontend/src/container/AIAssistant/components/MessageBubble/__tests__/MessageBubble.test.tsx

@@ -0,0 +1,85 @@
+import React from 'react';
+import { render, screen, userEvent } from 'tests/test-utils';
+import {
+	ErrorCodeDTO,
+	RetryActionDTO,
+} from 'api/ai-assistant/sigNozAIAssistantAPI.schemas';
+
+import { Message } from '../../../types';
+
+// react-markdown + remark-gfm are ESM-only and pull a large untransformed
+// dependency chain into jest. The error-rendering path under test renders
+// plain text (no markdown), so stub them to keep the import graph loadable.
+jest.mock('react-markdown', () => ({
+	__esModule: true,
+	default: ({ children }: { children?: React.ReactNode }): React.ReactNode =>
+		children,
+}));
+jest.mock('remark-gfm', () => ({
+	__esModule: true,
+	default: (): void => undefined,
+}));
+
+// eslint-disable-next-line import/first
+import MessageBubble from '../MessageBubble';
+
+function errorMessage(overrides: Partial<Message> = {}): Message {
+	return {
+		id: 'err-1',
+		role: 'assistant',
+		content: 'This conversation is still finishing a previous response.',
+		isError: true,
+		errorCode: ErrorCodeDTO.thread_busy,
+		retryAction: RetryActionDTO.manual,
+		createdAt: 0,
+		...overrides,
+	};
+}
+
+const retryButton = (): HTMLElement | null =>
+	screen.queryByRole('button', { name: /retry/i });
+
+describe('MessageBubble — error rendering', () => {
+	it('shows a Retry button for a manual error and invokes onRetry on click', async () => {
+		const onRetry = jest.fn();
+		render(<MessageBubble message={errorMessage()} onRetry={onRetry} />);
+
+		// Error copy is rendered, and the feedback bar is suppressed on errors.
+		expect(
+			screen.getByText(/still finishing a previous response/i),
+		).toBeInTheDocument();
+		expect(
+			screen.queryByRole('button', { name: /copy message/i }),
+		).not.toBeInTheDocument();
+
+		const button = retryButton();
+		expect(button).toBeInTheDocument();
+		await userEvent.click(button as HTMLElement);
+		expect(onRetry).toHaveBeenCalledTimes(1);
+	});
+
+	it('hides the Retry button when retryAction is none', () => {
+		render(
+			<MessageBubble
+				message={errorMessage({ retryAction: RetryActionDTO.none })}
+				onRetry={jest.fn()}
+			/>,
+		);
+		expect(retryButton()).not.toBeInTheDocument();
+	});
+
+	it('hides the Retry button when retryAction is auto', () => {
+		render(
+			<MessageBubble
+				message={errorMessage({ retryAction: RetryActionDTO.auto })}
+				onRetry={jest.fn()}
+			/>,
+		);
+		expect(retryButton()).not.toBeInTheDocument();
+	});
+
+	it('hides the Retry button when no onRetry handler is provided', () => {
+		render(<MessageBubble message={errorMessage()} />);
+		expect(retryButton()).not.toBeInTheDocument();
+	});
+});

frontend/src/container/AIAssistant/components/VirtualizedMessages/VirtualizedMessages.tsx

@@ -37,6 +37,9 @@ export default function VirtualizedMessages({
 	const regenerateAssistantMessage = useAIAssistantStore(
 		(s) => s.regenerateAssistantMessage,
 	);
+	const retryAssistantMessage = useAIAssistantStore(
+		(s) => s.retryAssistantMessage,
+	);
 	const { threadId } = useAIAssistantAnalyticsContext(conversationId);
 	const streamingStatus = useAIAssistantStore(
 		(s) => s.streams[conversationId]?.streamingStatus ?? '',
@@ -85,6 +88,14 @@ export default function VirtualizedMessages({
 		[conversationId, isStreaming, regenerateAssistantMessage, threadId],
 	);
 
+	const handleRetry = useCallback((): void => {
+		if (isStreaming) {
+			return;
+		}
+		void logEvent(AIAssistantEvents.RetryClicked, { threadId });
+		void retryAssistantMessage(conversationId);
+	}, [conversationId, isStreaming, retryAssistantMessage, threadId]);
+
 	// Scroll all the way to the actual bottom — including the 64px of bottom
 	// padding on the scroller — so the last bubble has visible breathing room
 	// above the disclaimer / input bar. Virtuoso's `scrollToIndex(LAST,
@@ -206,6 +217,11 @@ export default function VirtualizedMessages({
 									? (): void => handleRegenerate(msg.id)
 									: undefined
 							}
+							onRetry={
+								msg.isError && isLastAssistant && !showStreamingSlot
+									? handleRetry
+									: undefined
+							}
 							isLastAssistant={isLastAssistant}
 						/>
 					);

frontend/src/container/AIAssistant/events.ts

@@ -90,6 +90,7 @@ export enum AIAssistantEvents {
 	SuggestedPromptClicked = 'AI Assistant: Suggested prompt clicked',
 	CancelClicked = 'AI Assistant: Cancel clicked',
 	RegenerateClicked = 'AI Assistant: Regenerate clicked',
+	RetryClicked = 'AI Assistant: Retry clicked',
 	MessageCopied = 'AI Assistant: Message copied',
 	FeedbackSubmitted = 'AI Assistant: Feedback submitted',
 	ResourceOpened = 'AI Assistant: Resource opened',

frontend/src/container/AIAssistant/store/__tests__/useAIAssistantStore.test.ts

@@ -0,0 +1,263 @@
+import {
+	ErrorCodeDTO,
+	RetryActionDTO,
+} from 'api/ai-assistant/sigNozAIAssistantAPI.schemas';
+import type { SSEEvent } from 'api/ai-assistant/chat';
+
+import { useAIAssistantStore } from '../useAIAssistantStore';
+import type { Message } from '../../types';
+
+// The store talks to the chat API only through these named exports. Mock the
+// whole module so we can drive the SSE stream + REST calls deterministically.
+jest.mock('api/ai-assistant/chat', () => ({
+	__esModule: true,
+	createThread: jest.fn(),
+	sendMessage: jest.fn(),
+	streamEvents: jest.fn(),
+	approveExecution: jest.fn(),
+	clarifyExecution: jest.fn(),
+	regenerateMessage: jest.fn(),
+	rejectExecution: jest.fn(),
+	cancelExecution: jest.fn(),
+	listThreads: jest.fn(),
+	getThreadDetail: jest.fn(),
+	updateThread: jest.fn(),
+	submitFeedback: jest.fn(),
+}));
+
+// eslint-disable-next-line @typescript-eslint/no-var-requires, global-require
+const chat = jest.requireMock('api/ai-assistant/chat') as Record<
+	string,
+	jest.Mock
+>;
+
+// Builds a single-use async stream from a fixed list of SSE events.
+async function* eventStream(events: SSEEvent[]): AsyncGenerator<SSEEvent> {
+	for (const event of events) {
+		yield event;
+	}
+}
+
+function errorEvent(
+	executionId: string,
+	code: ErrorCodeDTO,
+	retryAction: RetryActionDTO,
+): SSEEvent {
+	return {
+		type: 'error',
+		executionId,
+		error: { code, message: 'backend message' },
+		retryAction,
+	};
+}
+
+function lastMessage(conversationId: string): Message {
+	const conv = useAIAssistantStore.getState().conversations[conversationId];
+	return conv.messages[conv.messages.length - 1];
+}
+
+describe('useAIAssistantStore — streaming error handling', () => {
+	beforeEach(() => {
+		jest.clearAllMocks();
+		useAIAssistantStore.setState((s) => {
+			s.conversations = {};
+			s.streams = {};
+			s.activeConversationId = null;
+		});
+	});
+
+	it('commits a manually-retryable error bubble with friendly copy and metadata', async () => {
+		chat.createThread.mockResolvedValue('thread-1');
+		chat.sendMessage.mockResolvedValue('exec-1');
+		chat.streamEvents.mockReturnValueOnce(
+			eventStream([
+				errorEvent('exec-1', ErrorCodeDTO.thread_busy, RetryActionDTO.manual),
+			]),
+		);
+
+		useAIAssistantStore.getState().startNewConversation();
+		await useAIAssistantStore.getState().sendMessage('hello');
+
+		const conv = useAIAssistantStore.getState().conversations['thread-1'];
+		expect(conv.messages).toHaveLength(2);
+		expect(conv.messages[0]).toMatchObject({ role: 'user', content: 'hello' });
+		expect(conv.messages[1]).toMatchObject({
+			role: 'assistant',
+			isError: true,
+			errorCode: ErrorCodeDTO.thread_busy,
+			retryAction: RetryActionDTO.manual,
+		});
+		// Code-specific FE copy, not the raw backend message.
+		expect(conv.messages[1].content).toContain(
+			'still finishing a previous response',
+		);
+	});
+
+	it('replays the send on retry without re-pushing the user message', async () => {
+		chat.createThread.mockResolvedValue('thread-1');
+		chat.sendMessage.mockResolvedValue('exec-1');
+		chat.streamEvents.mockReturnValueOnce(
+			eventStream([
+				errorEvent('exec-1', ErrorCodeDTO.thread_busy, RetryActionDTO.manual),
+			]),
+		);
+
+		useAIAssistantStore.getState().startNewConversation();
+		await useAIAssistantStore.getState().sendMessage('hello');
+
+		// The retry succeeds this time.
+		chat.streamEvents.mockReturnValueOnce(
+			eventStream([
+				{
+					type: 'message',
+					executionId: 'exec-1',
+					messageId: 'm1',
+					delta: 'Hi there',
+					done: true,
+				},
+			]),
+		);
+
+		await useAIAssistantStore.getState().retryAssistantMessage('thread-1');
+
+		const conv = useAIAssistantStore.getState().conversations['thread-1'];
+		// Error bubble replaced by the assistant reply; the user message stays.
+		expect(conv.messages).toHaveLength(2);
+		expect(conv.messages[0]).toMatchObject({ role: 'user', content: 'hello' });
+		expect(conv.messages[1]).toMatchObject({
+			role: 'assistant',
+			content: 'Hi there',
+		});
+		expect(conv.messages[1].isError).toBeUndefined();
+		// Thread already existed on retry; the user message was never re-sent as new.
+		expect(chat.createThread).toHaveBeenCalledTimes(1);
+		expect(chat.sendMessage).toHaveBeenCalledTimes(2);
+	});
+
+	it('silently retries auto-flagged errors, then downgrades to manual once spent', async () => {
+		chat.createThread.mockResolvedValue('thread-2');
+		chat.sendMessage.mockResolvedValue('exec');
+		// Always auto-retryable: 1 initial attempt + MAX_AUTO_RETRIES (2) = 3 sends.
+		chat.streamEvents.mockImplementation(() =>
+			eventStream([
+				errorEvent('exec', ErrorCodeDTO.internal_error, RetryActionDTO.auto),
+			]),
+		);
+
+		useAIAssistantStore.getState().startNewConversation();
+		await useAIAssistantStore.getState().sendMessage('hi');
+
+		expect(chat.sendMessage).toHaveBeenCalledTimes(3);
+		expect(lastMessage('thread-2')).toMatchObject({
+			isError: true,
+			errorCode: ErrorCodeDTO.internal_error,
+			// Auto budget exhausted → presented as manual so a Retry button shows.
+			retryAction: RetryActionDTO.manual,
+		});
+	}, 10000);
+
+	it('marks rate-limit errors and offers no retry', async () => {
+		chat.createThread.mockResolvedValue('thread-3');
+		chat.sendMessage.mockResolvedValue('exec');
+		chat.streamEvents.mockReturnValueOnce(
+			eventStream([
+				errorEvent('exec', ErrorCodeDTO.hourly_message_limit, RetryActionDTO.none),
+			]),
+		);
+
+		useAIAssistantStore.getState().startNewConversation();
+		await useAIAssistantStore.getState().sendMessage('hi');
+
+		expect(lastMessage('thread-3')).toMatchObject({
+			isError: true,
+			isRateLimitError: true,
+			retryAction: RetryActionDTO.none,
+		});
+
+		// No retry thunk registered for a non-retryable error — retry is a no-op.
+		const before =
+			useAIAssistantStore.getState().conversations['thread-3'].messages.length;
+		await useAIAssistantStore.getState().retryAssistantMessage('thread-3');
+		expect(
+			useAIAssistantStore.getState().conversations['thread-3'].messages,
+		).toHaveLength(before);
+	});
+
+	it('recovers silently when an auto-flagged error succeeds on retry', async () => {
+		chat.createThread.mockResolvedValue('thread-4');
+		chat.sendMessage.mockResolvedValue('exec');
+		chat.streamEvents
+			.mockReturnValueOnce(
+				eventStream([
+					errorEvent('exec', ErrorCodeDTO.internal_error, RetryActionDTO.auto),
+				]),
+			)
+			.mockReturnValueOnce(
+				eventStream([
+					{
+						type: 'message',
+						executionId: 'exec',
+						messageId: 'm1',
+						delta: 'Recovered',
+						done: true,
+					},
+				]),
+			);
+
+		useAIAssistantStore.getState().startNewConversation();
+		await useAIAssistantStore.getState().sendMessage('hi');
+
+		// 1 initial attempt + 1 silent auto retry, then success — no error bubble.
+		expect(chat.sendMessage).toHaveBeenCalledTimes(2);
+		const conv = useAIAssistantStore.getState().conversations['thread-4'];
+		expect(conv.messages).toHaveLength(2);
+		expect(conv.messages[0]).toMatchObject({ role: 'user', content: 'hi' });
+		expect(conv.messages[1]).toMatchObject({
+			role: 'assistant',
+			content: 'Recovered',
+		});
+		expect(conv.messages.some((m) => m.isError)).toBe(false);
+	}, 10000);
+
+	it('replays the originating action on retry for a non-send error (approve)', async () => {
+		chat.approveExecution.mockResolvedValue('exec-a');
+		chat.streamEvents.mockReturnValueOnce(
+			eventStream([
+				errorEvent('exec-a', ErrorCodeDTO.thread_busy, RetryActionDTO.manual),
+			]),
+		);
+
+		const convId = useAIAssistantStore.getState().startNewConversation();
+		await useAIAssistantStore.getState().approveAction(convId, 'approval-1');
+
+		expect(lastMessage(convId)).toMatchObject({
+			isError: true,
+			retryAction: RetryActionDTO.manual,
+		});
+		expect(chat.approveExecution).toHaveBeenCalledTimes(1);
+
+		// Retry replays the approval (not a send) and succeeds this time.
+		chat.streamEvents.mockReturnValueOnce(
+			eventStream([
+				{
+					type: 'message',
+					executionId: 'exec-a',
+					messageId: 'm1',
+					delta: 'Approved',
+					done: true,
+				},
+			]),
+		);
+		await useAIAssistantStore.getState().retryAssistantMessage(convId);
+
+		const conv = useAIAssistantStore.getState().conversations[convId];
+		expect(conv.messages).toHaveLength(1);
+		expect(conv.messages[0]).toMatchObject({
+			role: 'assistant',
+			content: 'Approved',
+		});
+		expect(conv.messages[0].isError).toBeUndefined();
+		expect(chat.approveExecution).toHaveBeenCalledTimes(2);
+		expect(chat.sendMessage).not.toHaveBeenCalled();
+	});
+});

frontend/src/container/AIAssistant/store/useAIAssistantStore.ts

@@ -8,6 +8,7 @@ import type {
 	MessageActionDTO,
 	MessageSummaryDTOBlocksAnyOfItem,
 } from 'api/ai-assistant/sigNozAIAssistantAPI.schemas';
+import { RetryActionDTO } from 'api/ai-assistant/sigNozAIAssistantAPI.schemas';
 
 import {
 	approveExecution,
@@ -35,7 +36,10 @@ import {
 	MessageBlock,
 	MessageRole,
 } from '../types';
-import { resolveAssistantErrorMessage } from '../utils/resolveAssistantErrorMessage';
+import {
+	resolveAssistantError,
+	type AssistantErrorResolution,
+} from '../utils/resolveAssistantError';
 
 // ---------------------------------------------------------------------------
 // Types used by module-level helpers
@@ -56,6 +60,15 @@ interface SSEStreamCtx {
 
 const streamControllers = new Map<string, AbortController>();
 
+/**
+ * Per-conversation retry thunks for the most recent failed turn. Populated by
+ * `finalizeStreamingError` when the error is manually retryable; consumed by
+ * the `retryAssistantMessage` action when the user clicks Retry. Transient
+ * (not persisted) — it shares the in-memory lifetime of the error bubble it
+ * backs, so a page reload drops both together.
+ */
+const retryRegistry = new Map<string, () => Promise<void>>();
+
 function abortStream(conversationId: string): void {
 	const ctrl = streamControllers.get(conversationId);
 	if (ctrl) {
@@ -197,7 +210,7 @@ function resetStreamingState(
  * Marker thrown by `runStreamingLoop` when an SSE event reports
  * `invalid_token`. Callers that own an originating action (sendMessage /
  * approve / clarify / regenerate) catch this and re-issue that action via
- * `streamWithAuthRetry`; the retry's first REST call will 401, at which point
+ * `streamWithRetry`; the retry's first REST call will 401, at which point
  * the shared axios `interceptorRejected` rotates the access token and replays.
  */
 class AuthExpiredError extends Error {
@@ -207,27 +220,50 @@ class AuthExpiredError extends Error {
 	}
 }
 
+/** Capped silent re-attempts for backend-flagged transient (`auto`) errors. */
+const MAX_AUTO_RETRIES = 2;
+/** Backoff before each auto re-attempt, indexed by prior auto-retry count. */
+const AUTO_RETRY_BACKOFF_MS = [500, 1500];
+
+function delay(ms: number): Promise<void> {
+	return new Promise((resolve) => {
+		setTimeout(resolve, ms);
+	});
+}
+
+/** True when the SSE error carries the backend's `retryAction: 'auto'` flag. */
+function isAutoRetryableError(err: unknown): boolean {
+	return (
+		(err as { retryAction?: unknown } | undefined)?.retryAction ===
+		RetryActionDTO.auto
+	);
+}
+
 /**
  * Runs the originating action (e.g. sendMessage POST) and streams the
- * resulting execution. On `AuthExpiredError`, re-issues `start` once — the
- * retry's REST call hits 401, the shared axios interceptor rotates the
- * access token and replays, and the new SSE picks up the rotated token from
- * localStorage. Backend signals `retryAction: 'manual'` for `invalid_token`,
- * so the dead execution can't be resumed — only a fresh one helps.
+ * resulting execution, with two independent retry budgets:
+ *
+ *   • Auth — on `AuthExpiredError` (SSE `invalid_token`), re-issues `start`
+ *     once. The retry's REST call 401s, the shared axios interceptor rotates
+ *     the access token + replays, and the new SSE picks up the rotated token.
+ *     Backend flags `invalid_token` as `manual`, so only a fresh execution helps.
+ *   • Auto — on an SSE error the backend flagged `retryAction: 'auto'`
+ *     (transient), silently re-issues `start` up to `MAX_AUTO_RETRIES` times
+ *     with backoff. Once exhausted the error propagates so the caller can
+ *     surface a manual Retry affordance.
+ *
+ * Both reset the stream state before re-attempting so a dead execution's
+ * partial output isn't concatenated onto the retry.
  */
-async function streamWithAuthRetry(
+async function streamWithRetry(
 	conversationId: string,
 	start: () => Promise<string>,
 	set: StoreSetter,
 ): Promise<void> {
-	for (let attempt = 0; attempt <= 1; attempt += 1) {
-		if (attempt > 0) {
-			// Drop any partial content/events from the previous attempt so the
-			// retried execution's stream isn't concatenated with the dead one.
-			set((s) => {
-				resetStreamingState(s, conversationId);
-			});
-		}
+	let authRetried = false;
+	let autoRetries = 0;
+
+	for (;;) {
 		// eslint-disable-next-line no-await-in-loop
 		const executionId = await start();
 		const ctrl = newStreamController(conversationId);
@@ -242,10 +278,28 @@ async function streamWithAuthRetry(
 			return;
 		} catch (err) {
 			streamControllers.delete(conversationId);
-			if (err instanceof AuthExpiredError && attempt < 1) {
-				continue;
+
+			if (err instanceof AuthExpiredError && !authRetried) {
+				authRetried = true;
+			} else if (isAutoRetryableError(err) && autoRetries < MAX_AUTO_RETRIES) {
+				// eslint-disable-next-line no-await-in-loop
+				await delay(AUTO_RETRY_BACKOFF_MS[autoRetries] ?? 1500);
+				autoRetries += 1;
+			} else {
+				if (isAutoRetryableError(err)) {
+					// Auto-retry budget spent — present the failure as manually
+					// retryable so the caller surfaces a Retry button rather than
+					// silently giving up.
+					(err as { retryAction?: RetryActionDTO }).retryAction =
+						RetryActionDTO.manual;
+				}
+				throw err;
 			}
-			throw err;
+
+			// Drop partial content/events from the failed attempt before retrying.
+			set((s) => {
+				resetStreamingState(s, conversationId);
+			});
 		}
 	}
 }
@@ -258,7 +312,7 @@ async function streamWithAuthRetry(
  *
  * On an `invalid_token` error event (e.g. MCP auth expired mid-execution),
  * throws `AuthExpiredError` so the caller can re-issue the originating
- * action via `streamWithAuthRetry`. We don't refresh here ourselves — the
+ * action via `streamWithRetry`. We don't refresh here ourselves — the
  * retry's REST call will 401 and the shared axios `interceptorRejected`
  * handles rotation + replay. Throws on any other `error` event — the
  * caller's catch block handles UI feedback.
@@ -484,20 +538,37 @@ function hasPendingInput(conversationId: string, get: StoreGetter): boolean {
 	return Boolean(stream?.pendingApproval || stream?.pendingClarification);
 }
 
+/**
+ * Commits a failed turn as an error message and removes the stream entry.
+ * When the failure is manually retryable and a `retry` thunk is supplied, the
+ * thunk is stashed in `retryRegistry` so the bubble's Retry button can replay
+ * the originating action.
+ */
 function finalizeStreamingError(
 	conversationId: string,
-	errorContent: string,
+	resolution: AssistantErrorResolution,
 	set: StoreSetter,
-	isRateLimit = false,
+	retry?: () => Promise<void>,
 ): void {
+	const { message, code, retryAction, isRateLimit } = resolution;
+
+	if (retryAction === RetryActionDTO.manual && retry) {
+		retryRegistry.set(conversationId, retry);
+	} else {
+		retryRegistry.delete(conversationId);
+	}
+
 	set((s) => {
 		const conv = s.conversations[conversationId];
 		if (conv) {
 			conv.messages.push({
 				id: uuidv4(),
 				role: 'assistant',
-				content: errorContent,
+				content: message,
 				createdAt: Date.now(),
+				isError: true,
+				retryAction,
+				...(code ? { errorCode: code } : {}),
 				...(isRateLimit ? { isRateLimitError: true } : {}),
 			});
 			conv.updatedAt = Date.now();
@@ -506,6 +577,40 @@ function finalizeStreamingError(
 	});
 }
 
+/**
+ * Shared streaming wrapper for actions that have no pre-stream setup beyond
+ * resetting state (approve / clarify / regenerate). Streams the execution,
+ * finalizes the message on success, and on failure resolves the error +
+ * registers `retry` (the caller's own re-invocation) so the bubble can replay
+ * it. `sendMessage` does not use this — it owns thread-creation/re-keying and
+ * runs its own equivalent loop.
+ */
+async function streamAndFinalize(
+	conversationId: string,
+	start: () => Promise<string>,
+	fallback: string,
+	logLabel: string,
+	set: StoreSetter,
+	get: StoreGetter,
+	retry: () => Promise<void>,
+): Promise<void> {
+	try {
+		await streamWithRetry(conversationId, start, set);
+		if (!hasPendingInput(conversationId, get)) {
+			finalizeStreamingMessage(conversationId, set, get);
+		}
+	} catch (err) {
+		// Abort errors are expected when the user cancels — not a failure.
+		if (err instanceof DOMException && err.name === 'AbortError') {
+			return;
+		}
+		// eslint-disable-next-line no-console
+		console.error(logLabel, err);
+		const resolution = resolveAssistantError(err, fallback);
+		finalizeStreamingError(conversationId, resolution, set, retry);
+	}
+}
+
 // ---------------------------------------------------------------------------
 // Store interface
 // ---------------------------------------------------------------------------
@@ -564,6 +669,8 @@ export interface AIAssistantStore {
 		conversationId: string,
 		messageId: string,
 	) => Promise<void>;
+	/** Replays the originating action for a manually-retryable error bubble. */
+	retryAssistantMessage: (conversationId: string) => Promise<void>;
 	submitMessageFeedback: (
 		messageId: string,
 		rating: FeedbackRating,
@@ -877,7 +984,7 @@ export const useAIAssistantStore = create<AIAssistantStore>()(
 					// there's no "originating action" to redo — reopening the
 					// same dead executionId would just re-emit the failure.
 					// Let the error bubble; the user can send a new message,
-					// which will go through `streamWithAuthRetry`.
+					// which will go through `streamWithRetry`.
 					if (
 						detail.activeExecutionId &&
 						!streamControllers.has(threadId) &&
@@ -1060,7 +1167,7 @@ export const useAIAssistantStore = create<AIAssistantStore>()(
 				attachments?: MessageAttachment[],
 				contexts?: MessageContext[],
 			): Promise<void> => {
-				let convId = get().activeConversationId;
+				const convId = get().activeConversationId;
 				if (!convId || !get().conversations[convId]) {
 					return;
 				}
@@ -1093,63 +1200,75 @@ export const useAIAssistantStore = create<AIAssistantStore>()(
 				};
 
 				set((state) => {
-					const conv = state.conversations[convId!];
+					const conv = state.conversations[convId];
 					conv.messages.push(userMessage);
 					conv.updatedAt = Date.now();
 					if (!conv.title && text.trim()) {
 						conv.title = deriveTitle(text);
 					}
-					resetStreamingState(state, convId!);
+					resetStreamingState(state, convId);
 				});
 
-				try {
-					let { threadId } = get().conversations[convId];
-					if (!threadId) {
-						threadId = await createThread();
-						// Re-key the conversation from client UUID to backend threadId
-						// so fetchThreads won't create a duplicate entry later.
-						const oldId = convId;
-						convId = threadId;
-						set((s) => {
-							const conv = s.conversations[oldId];
-							if (conv) {
-								conv.id = convId!;
-								conv.threadId = convId!;
-								s.conversations[convId!] = conv;
-								delete s.conversations[oldId];
-								if (s.activeConversationId === oldId) {
-									s.activeConversationId = convId!;
+				// The full send — ensure a backend thread exists (re-keying the
+				// optimistic client UUID on first send), POST the message, and
+				// stream the reply. Defined as a closure so the error bubble's
+				// Retry button can replay it without re-pushing the user message.
+				const runSend = async (cid: string): Promise<void> => {
+					let targetConvId = cid;
+					try {
+						let { threadId } = get().conversations[targetConvId];
+						if (!threadId) {
+							threadId = await createThread();
+							// Re-key the conversation from client UUID to backend threadId
+							// so fetchThreads won't create a duplicate entry later.
+							const oldId = targetConvId;
+							const newId = threadId;
+							set((s) => {
+								const conv = s.conversations[oldId];
+								if (conv) {
+									conv.id = newId;
+									conv.threadId = newId;
+									s.conversations[newId] = conv;
+									delete s.conversations[oldId];
+									if (s.activeConversationId === oldId) {
+										s.activeConversationId = newId;
+									}
+									const stream = s.streams[oldId];
+									if (stream) {
+										s.streams[newId] = stream;
+										delete s.streams[oldId];
+									}
 								}
-								const stream = s.streams[oldId];
-								if (stream) {
-									s.streams[convId!] = stream;
-									delete s.streams[oldId];
-								}
-							}
-						});
-					}
-					const tid = threadId;
-					await streamWithAuthRetry(
-						convId,
-						() => sendMessageToThread(tid, text, contexts),
-						set,
-					);
+							});
+							targetConvId = newId;
+						}
+						const tid = threadId;
+						await streamWithRetry(
+							targetConvId,
+							() => sendMessageToThread(tid, text, contexts),
+							set,
+						);
 
-					if (!hasPendingInput(convId, get)) {
-						finalizeStreamingMessage(convId, set, get);
+						if (!hasPendingInput(targetConvId, get)) {
+							finalizeStreamingMessage(targetConvId, set, get);
+						}
+					} catch (err) {
+						// Abort errors are expected when the user cancels — not a failure.
+						if (err instanceof DOMException && err.name === 'AbortError') {
+							return;
+						}
+						console.error('[AIAssistant] sendMessage failed:', err);
+						const resolution = resolveAssistantError(
+							err,
+							'Something went wrong while fetching the response. Please try again.',
+						);
+						finalizeStreamingError(targetConvId, resolution, set, () =>
+							runSend(targetConvId),
+						);
 					}
-				} catch (err) {
-					// Abort errors are expected when the user cancels — not a failure
-					if (err instanceof DOMException && err.name === 'AbortError') {
-						return;
-					}
-					console.error('[AIAssistant] sendMessage failed:', err);
-					const { message, isRateLimit } = resolveAssistantErrorMessage(
-						err,
-						'Something went wrong while fetching the response. Please try again.',
-					);
-					finalizeStreamingError(convId, message, set, isRateLimit);
-				}
+				};
+
+				await runSend(convId);
 			},
 
 			approveAction: async (
@@ -1167,26 +1286,17 @@ export const useAIAssistantStore = create<AIAssistantStore>()(
 					}
 				});
 
-				try {
-					await streamWithAuthRetry(
+				const run = (): Promise<void> =>
+					streamAndFinalize(
 						conversationId,
 						() => approveExecution(approvalId),
-						set,
-					);
-					if (!hasPendingInput(conversationId, get)) {
-						finalizeStreamingMessage(conversationId, set, get);
-					}
-				} catch (err) {
-					if (err instanceof DOMException && err.name === 'AbortError') {
-						return;
-					}
-					console.error('[AIAssistant] approveAction failed:', err);
-					const { message, isRateLimit } = resolveAssistantErrorMessage(
-						err,
 						'Something went wrong while processing the approval. Please try again.',
+						'[AIAssistant] approveAction failed:',
+						set,
+						get,
+						run,
 					);
-					finalizeStreamingError(conversationId, message, set, isRateLimit);
-				}
+				await run();
 			},
 
 			rejectAction: async (
@@ -1246,26 +1356,17 @@ export const useAIAssistantStore = create<AIAssistantStore>()(
 					resetStreamingState(s, conversationId);
 				});
 
-				try {
-					await streamWithAuthRetry(
+				const run = (): Promise<void> =>
+					streamAndFinalize(
 						conversationId,
 						() => regenerateMessage(messageId),
-						set,
-					);
-					if (!hasPendingInput(conversationId, get)) {
-						finalizeStreamingMessage(conversationId, set, get);
-					}
-				} catch (err) {
-					if (err instanceof DOMException && err.name === 'AbortError') {
-						return;
-					}
-					console.error('[AIAssistant] regenerateAssistantMessage failed:', err);
-					const { message, isRateLimit } = resolveAssistantErrorMessage(
-						err,
 						'Something went wrong while regenerating the response. Please try again.',
+						'[AIAssistant] regenerateAssistantMessage failed:',
+						set,
+						get,
+						run,
 					);
-					finalizeStreamingError(conversationId, message, set, isRateLimit);
-				}
+				await run();
 			},
 
 			submitMessageFeedback: async (
@@ -1312,26 +1413,42 @@ export const useAIAssistantStore = create<AIAssistantStore>()(
 					}
 				});
 
-				try {
-					await streamWithAuthRetry(
+				const run = (): Promise<void> =>
+					streamAndFinalize(
 						conversationId,
 						() => clarifyExecution(clarificationId, answers),
-						set,
-					);
-					if (!hasPendingInput(conversationId, get)) {
-						finalizeStreamingMessage(conversationId, set, get);
-					}
-				} catch (err) {
-					if (err instanceof DOMException && err.name === 'AbortError') {
-						return;
-					}
-					console.error('[AIAssistant] submitClarification failed:', err);
-					const { message, isRateLimit } = resolveAssistantErrorMessage(
-						err,
 						'Something went wrong while processing your answers. Please try again.',
+						'[AIAssistant] submitClarification failed:',
+						set,
+						get,
+						run,
 					);
-					finalizeStreamingError(conversationId, message, set, isRateLimit);
+				await run();
+			},
+
+			retryAssistantMessage: async (conversationId: string): Promise<void> => {
+				const retry = retryRegistry.get(conversationId);
+				if (!retry) {
+					return;
 				}
+				retryRegistry.delete(conversationId);
+
+				// Drop the trailing error bubble we're retrying from and reset the
+				// stream so the in-progress retry renders immediately. The retry
+				// thunk replays the originating action without re-pushing the
+				// user's message.
+				set((s) => {
+					const conv = s.conversations[conversationId];
+					if (conv) {
+						const last = conv.messages[conv.messages.length - 1];
+						if (last?.isError) {
+							conv.messages.pop();
+						}
+					}
+					resetStreamingState(s, conversationId);
+				});
+
+				await retry();
 			},
 		})),
 		{

frontend/src/container/AIAssistant/types.ts

@@ -15,9 +15,11 @@
 import type {
 	ApprovalEventDTO,
 	ClarificationEventDTO,
+	ErrorCodeDTO,
 	FeedbackRatingDTO,
 	MessageActionDTO,
 	MessageActionKindDTO,
+	RetryActionDTO,
 } from 'api/ai-assistant/sigNozAIAssistantAPI.schemas';
 
 /** Client-only file attachment — no API equivalent (uploads happen via data URLs). */
@@ -91,6 +93,18 @@ export interface Message {
 	 * bar (copy/vote/regenerate) is hidden — retrying would just 429 again.
 	 */
 	isRateLimitError?: boolean;
+	/**
+	 * Marks an assistant message that represents a failed turn. Drives the
+	 * error styling and replaces the feedback bar with a retry affordance.
+	 */
+	isError?: boolean;
+	/** Known backend error code for the failure, when recognised. */
+	errorCode?: ErrorCodeDTO;
+	/**
+	 * Retry semantics for a failed turn — `manual` renders an inline Retry
+	 * button on the error bubble; `none`/`auto` render no button.
+	 */
+	retryAction?: RetryActionDTO;
 	createdAt: number;
 }
 

frontend/src/container/AIAssistant/utils/__tests__/resolveAssistantError.test.ts

@@ -0,0 +1,154 @@
+import { AxiosError } from 'axios';
+import {
+	ErrorCodeDTO,
+	RetryActionDTO,
+} from 'api/ai-assistant/sigNozAIAssistantAPI.schemas';
+
+import { resolveAssistantError } from '../resolveAssistantError';
+
+const FALLBACK = 'Something went wrong. Please try again.';
+
+function restError(status: number, code: string, message: string): AxiosError {
+	const err = new AxiosError('Request failed');
+	err.response = {
+		status,
+		data: { error: { code, message } },
+	} as AxiosError['response'];
+	return err;
+}
+
+describe('resolveAssistantError', () => {
+	describe('message resolution', () => {
+		it('prefers code-specific FE copy over the backend message', () => {
+			const err = restError(409, ErrorCodeDTO.thread_busy, 'raw backend phrasing');
+
+			const { message } = resolveAssistantError(err, FALLBACK);
+			expect(message).toBe(
+				'This conversation is still finishing a previous response. Give it a moment and try again.',
+			);
+		});
+
+		it('falls through to the backend message for a known code without FE copy', () => {
+			const err = restError(
+				400,
+				ErrorCodeDTO.message_not_found,
+				'No such message exists.',
+			);
+
+			expect(resolveAssistantError(err, FALLBACK)).toStrictEqual({
+				message: 'No such message exists.',
+				code: ErrorCodeDTO.message_not_found,
+				retryAction: RetryActionDTO.none,
+				isRateLimit: false,
+			});
+		});
+
+		it('falls back when the error code is not in ErrorCodeDTO', () => {
+			const err = restError(400, 'future_unknown_code', 'Backend-only message');
+
+			expect(resolveAssistantError(err, FALLBACK)).toStrictEqual({
+				message: FALLBACK,
+				code: undefined,
+				retryAction: RetryActionDTO.none,
+				isRateLimit: false,
+			});
+		});
+	});
+
+	describe('rate limiting', () => {
+		it('marks HTTP 429 responses as rate limited and non-retryable', () => {
+			const err = restError(
+				429,
+				ErrorCodeDTO.hourly_message_limit,
+				'Hourly limit reached.',
+			);
+
+			expect(resolveAssistantError(err, FALLBACK)).toStrictEqual({
+				message: "You've reached the hourly message limit. Please try again later.",
+				code: ErrorCodeDTO.hourly_message_limit,
+				retryAction: RetryActionDTO.none,
+				isRateLimit: true,
+			});
+		});
+
+		it('treats known SSE rate-limit codes as rate limited', () => {
+			const err = Object.assign(new Error('Daily token limit exceeded.'), {
+				code: ErrorCodeDTO.daily_token_limit,
+			});
+
+			const res = resolveAssistantError(err, FALLBACK);
+			expect(res.isRateLimit).toBe(true);
+			expect(res.retryAction).toBe(RetryActionDTO.none);
+		});
+
+		it('marks 429 as rate limited even when the code is unknown', () => {
+			const err = restError(429, 'future_unknown_code', 'Too many requests');
+
+			expect(resolveAssistantError(err, FALLBACK)).toStrictEqual({
+				message: FALLBACK,
+				code: undefined,
+				retryAction: RetryActionDTO.none,
+				isRateLimit: true,
+			});
+		});
+	});
+
+	describe('retryAction resolution', () => {
+		it('honours an explicit retryAction from an SSE error event', () => {
+			const err = Object.assign(new Error('Transient hiccup'), {
+				code: ErrorCodeDTO.internal_error,
+				retryAction: RetryActionDTO.auto,
+			});
+
+			expect(resolveAssistantError(err, FALLBACK).retryAction).toBe(
+				RetryActionDTO.auto,
+			);
+		});
+
+		it('forces none for non-retryable permission errors', () => {
+			const err = restError(403, ErrorCodeDTO.permission_denied, 'forbidden');
+
+			expect(resolveAssistantError(err, FALLBACK).retryAction).toBe(
+				RetryActionDTO.none,
+			);
+		});
+
+		it('derives manual for 409 conflicts', () => {
+			const err = restError(409, ErrorCodeDTO.thread_has_active_execution, 'busy');
+
+			expect(resolveAssistantError(err, FALLBACK).retryAction).toBe(
+				RetryActionDTO.manual,
+			);
+		});
+
+		it('derives manual for 5xx responses', () => {
+			const err = restError(503, 'future_unknown_code', 'unavailable');
+
+			expect(resolveAssistantError(err, FALLBACK).retryAction).toBe(
+				RetryActionDTO.manual,
+			);
+		});
+
+		it('derives manual for network failures with no response', () => {
+			const err = new AxiosError('Network Error');
+
+			expect(resolveAssistantError(err, FALLBACK).retryAction).toBe(
+				RetryActionDTO.manual,
+			);
+		});
+
+		it('derives none for other 4xx responses', () => {
+			const err = restError(400, 'future_unknown_code', 'bad request');
+
+			expect(resolveAssistantError(err, FALLBACK).retryAction).toBe(
+				RetryActionDTO.none,
+			);
+		});
+
+		it('defaults to manual for non-Axios errors with no code', () => {
+			expect(resolveAssistantError(new Error('boom'), FALLBACK).retryAction).toBe(
+				RetryActionDTO.manual,
+			);
+		});
+	});
+});

frontend/src/container/AIAssistant/utils/__tests__/resolveAssistantErrorMessage.test.ts

@@ -1,91 +0,0 @@
-import { AxiosError } from 'axios';
-import { ErrorCodeDTO } from 'api/ai-assistant/sigNozAIAssistantAPI.schemas';
-
-import { resolveAssistantErrorMessage } from '../resolveAssistantErrorMessage';
-
-const FALLBACK = 'Something went wrong. Please try again.';
-
-describe('resolveAssistantErrorMessage', () => {
-	it('returns backend message for a known error code', () => {
-		const err = new AxiosError('Request failed');
-		err.response = {
-			status: 400,
-			data: {
-				error: {
-					code: ErrorCodeDTO.thread_busy,
-					message: 'This thread is busy. Try again shortly.',
-				},
-			},
-		} as AxiosError['response'];
-
-		expect(resolveAssistantErrorMessage(err, FALLBACK)).toStrictEqual({
-			message: 'This thread is busy. Try again shortly.',
-			isRateLimit: false,
-		});
-	});
-
-	it('falls back when error code is not in ErrorCodeDTO', () => {
-		const err = new AxiosError('Request failed');
-		err.response = {
-			status: 400,
-			data: {
-				error: {
-					code: 'future_unknown_code',
-					message: 'Backend-only message',
-				},
-			},
-		} as AxiosError['response'];
-
-		expect(resolveAssistantErrorMessage(err, FALLBACK)).toStrictEqual({
-			message: FALLBACK,
-			isRateLimit: false,
-		});
-	});
-
-	it('marks HTTP 429 responses as rate limited', () => {
-		const err = new AxiosError('Too many requests');
-		err.response = {
-			status: 429,
-			data: {
-				error: {
-					code: ErrorCodeDTO.hourly_message_limit,
-					message: 'Hourly limit reached.',
-				},
-			},
-		} as AxiosError['response'];
-
-		expect(resolveAssistantErrorMessage(err, FALLBACK)).toStrictEqual({
-			message: 'Hourly limit reached.',
-			isRateLimit: true,
-		});
-	});
-
-	it('uses backend message for known SSE rate-limit error codes', () => {
-		const err = Object.assign(new Error('Daily token limit exceeded.'), {
-			code: ErrorCodeDTO.daily_token_limit,
-		});
-
-		expect(resolveAssistantErrorMessage(err, FALLBACK)).toStrictEqual({
-			message: 'Daily token limit exceeded.',
-			isRateLimit: true,
-		});
-	});
-
-	it('marks 429 as rate limited even when error code is unknown', () => {
-		const err = new AxiosError('Too many requests');
-		err.response = {
-			status: 429,
-			data: {
-				error: {
-					code: 'future_unknown_code',
-					message: 'Too many requests',
-				},
-			},
-		} as AxiosError['response'];
-
-		expect(resolveAssistantErrorMessage(err, FALLBACK)).toStrictEqual({
-			message: FALLBACK,
-			isRateLimit: true,
-		});
-	});
-});

frontend/src/container/AIAssistant/utils/resolveAssistantError.ts

@@ -0,0 +1,209 @@
+import { isAxiosError } from 'axios';
+import {
+	ErrorCodeDTO,
+	RetryActionDTO,
+	type ErrorBodyDTO,
+	type ErrorResponseDTO,
+} from 'api/ai-assistant/sigNozAIAssistantAPI.schemas';
+
+export interface AssistantErrorResolution {
+	/** User-facing copy: code-specific FE copy → backend message → caller fallback. */
+	message: string;
+	/** Known backend error code, when one we recognise was supplied. */
+	code?: ErrorCodeDTO;
+	/**
+	 * Whether/how the failed action may be retried:
+	 *   • `auto`   — transient; the caller may silently re-attempt (capped).
+	 *   • `manual` — surface a Retry affordance to the user.
+	 *   • `none`   — retrying would re-fail deterministically; offer nothing.
+	 */
+	retryAction: RetryActionDTO;
+	/** Quota/limit error — callers hide the retry + feedback bar (retrying just re-limits). */
+	isRateLimit: boolean;
+}
+
+/** Quota/limit codes — surfaced as rate-limit errors (no retry, feedback bar hidden). */
+const RATE_LIMIT_ERROR_CODES = new Set<ErrorCodeDTO>([
+	ErrorCodeDTO.rate_limit_override_exceeds_ceiling,
+	ErrorCodeDTO.thread_message_limit,
+	ErrorCodeDTO.connection_limit_exceeded,
+	ErrorCodeDTO.hourly_message_limit,
+	ErrorCodeDTO.daily_message_limit,
+	ErrorCodeDTO.daily_token_limit,
+	ErrorCodeDTO.daily_cost_limit,
+	ErrorCodeDTO.budget_exceeded,
+]);
+
+/**
+ * Codes whose retry would re-fail deterministically — permission/config/validation
+ * failures. These force `retryAction: none` regardless of HTTP status.
+ */
+const NON_RETRYABLE_CODES = new Set<ErrorCodeDTO>([
+	ErrorCodeDTO.permission_denied,
+	ErrorCodeDTO.user_disabled,
+	ErrorCodeDTO.org_disabled,
+	ErrorCodeDTO.validation_error,
+	ErrorCodeDTO.invalid_content_length,
+	ErrorCodeDTO.invalid_fork_target,
+	ErrorCodeDTO.missing_signoz_url,
+	ErrorCodeDTO.invalid_signoz_url,
+	ErrorCodeDTO.region_not_configured,
+]);
+
+/**
+ * Code-specific, user-friendly copy. Takes precedence over the backend's raw
+ * `error.message` so the user sees an actionable, consistent sentence rather
+ * than internal phrasing. Codes absent here fall through to the backend message.
+ */
+const ERROR_CODE_COPY: Partial<Record<ErrorCodeDTO, string>> = {
+	[ErrorCodeDTO.permission_denied]:
+		"You don't have permission to do that. Contact your workspace admin if you think this is a mistake.",
+	[ErrorCodeDTO.user_disabled]:
+		'Your access to the AI assistant has been disabled. Contact your workspace admin to re-enable it.',
+	[ErrorCodeDTO.org_disabled]:
+		'The AI assistant is disabled for your organisation. An admin can enable it in settings.',
+	[ErrorCodeDTO.thread_busy]:
+		'This conversation is still finishing a previous response. Give it a moment and try again.',
+	[ErrorCodeDTO.thread_has_active_execution]:
+		'This conversation is still finishing a previous response. Give it a moment and try again.',
+	[ErrorCodeDTO.hourly_message_limit]:
+		"You've reached the hourly message limit. Please try again later.",
+	[ErrorCodeDTO.daily_message_limit]:
+		"You've reached the daily message limit. Please try again tomorrow.",
+	[ErrorCodeDTO.daily_token_limit]:
+		"You've reached today's usage limit. Please try again tomorrow.",
+	[ErrorCodeDTO.daily_cost_limit]:
+		"You've reached today's usage limit. Please try again tomorrow.",
+	[ErrorCodeDTO.budget_exceeded]:
+		"You've reached your usage budget. Contact your workspace admin to raise it.",
+	[ErrorCodeDTO.thread_message_limit]:
+		'This conversation has reached its length limit. Start a new conversation to continue.',
+	[ErrorCodeDTO.connection_limit_exceeded]:
+		'Too many active conversations right now. Close one and try again.',
+	[ErrorCodeDTO.max_turns_exceeded]:
+		'The assistant reached the maximum number of steps for this request. Try rephrasing or breaking it into smaller asks.',
+	[ErrorCodeDTO.region_unreachable]:
+		"Couldn't reach your region's services. Please try again in a moment.",
+	[ErrorCodeDTO.region_not_configured]:
+		'No region is configured for the AI assistant yet. An admin can set this up in settings.',
+	[ErrorCodeDTO.mcp_unavailable]:
+		'A required service is temporarily unavailable. Please try again shortly.',
+	[ErrorCodeDTO.sandbox_unavailable]:
+		'The execution environment is temporarily unavailable. Please try again shortly.',
+	[ErrorCodeDTO.internal_error]:
+		'Something went wrong on our end. Please try again.',
+};
+
+function isErrorCodeDTO(code: string | undefined): code is ErrorCodeDTO {
+	return (
+		code !== undefined && (Object.values(ErrorCodeDTO) as string[]).includes(code)
+	);
+}
+
+function isRetryActionDTO(value: unknown): value is RetryActionDTO {
+	return (
+		typeof value === 'string' &&
+		(Object.values(RetryActionDTO) as string[]).includes(value)
+	);
+}
+
+/**
+ * Pulls the structured error body out of either an Axios REST error or the
+ * SSE error the streaming loop throws (a plain `Error` augmented with `code`).
+ */
+function getErrorBody(err: unknown): ErrorBodyDTO | null {
+	if (isAxiosError(err)) {
+		return (err.response?.data as ErrorResponseDTO | undefined)?.error ?? null;
+	}
+
+	const code = (err as { code?: string } | undefined)?.code;
+	const message = err instanceof Error ? err.message : undefined;
+	if (!code || !message) {
+		return null;
+	}
+
+	return { code: code as ErrorCodeDTO, message };
+}
+
+function isRateLimit(code: ErrorCodeDTO | undefined, err: unknown): boolean {
+	if (isAxiosError(err) && err.response?.status === 429) {
+		return true;
+	}
+	return code !== undefined && RATE_LIMIT_ERROR_CODES.has(code);
+}
+
+/**
+ * Resolves how the failed action may be retried. The backend's explicit signal
+ * (SSE `ErrorEventDTO.retryAction`) is authoritative; otherwise we derive it
+ * from the rate-limit/non-retryable code sets and the HTTP status.
+ */
+function resolveRetryAction(
+	err: unknown,
+	code: ErrorCodeDTO | undefined,
+	rateLimited: boolean,
+): RetryActionDTO {
+	const explicit = (err as { retryAction?: unknown } | undefined)?.retryAction;
+	if (isRetryActionDTO(explicit)) {
+		return explicit;
+	}
+
+	if (rateLimited || (code !== undefined && NON_RETRYABLE_CODES.has(code))) {
+		return RetryActionDTO.none;
+	}
+
+	if (isAxiosError(err)) {
+		const status = err.response?.status;
+		// No response → network/timeout failure; retrying may well succeed.
+		if (status === undefined || status === 408) {
+			return RetryActionDTO.manual;
+		}
+		if (status === 401 || status === 403) {
+			return RetryActionDTO.none;
+		}
+		if (status === 409 || status >= 500) {
+			return RetryActionDTO.manual;
+		}
+		// Other 4xx (validation, bad request) re-fail deterministically.
+		return RetryActionDTO.none;
+	}
+
+	// Non-Axios transport/parse error with no code — let the user retry.
+	return RetryActionDTO.manual;
+}
+
+function resolveMessage(
+	code: ErrorCodeDTO | undefined,
+	body: ErrorBodyDTO | null,
+	fallback: string,
+): string {
+	if (code !== undefined && ERROR_CODE_COPY[code]) {
+		return ERROR_CODE_COPY[code] as string;
+	}
+	// Trust the backend's message only for codes we recognise — never surface
+	// raw text for unknown codes (could be an internal stack trace).
+	if (code !== undefined && body?.message.trim()) {
+		return body.message.trim();
+	}
+	return fallback;
+}
+
+/**
+ * Single resolution point for both SSE and REST assistant errors. Maps the
+ * error onto user-facing copy plus retry semantics, degrading gracefully for
+ * unknown codes (falls back to `fallback` + a `manual` retry where sensible).
+ */
+export function resolveAssistantError(
+	err: unknown,
+	fallback: string,
+): AssistantErrorResolution {
+	const body = getErrorBody(err);
+	const code = isErrorCodeDTO(body?.code) ? body?.code : undefined;
+	const rateLimited = isRateLimit(code, err);
+
+	return {
+		message: resolveMessage(code, body, fallback),
+		code,
+		retryAction: resolveRetryAction(err, code, rateLimited),
+		isRateLimit: rateLimited,
+	};
+}

frontend/src/container/AIAssistant/utils/resolveAssistantErrorMessage.ts

@@ -1,71 +0,0 @@
-import { isAxiosError } from 'axios';
-import {
-	ErrorCodeDTO,
-	type ErrorBodyDTO,
-	type ErrorResponseDTO,
-} from 'api/ai-assistant/sigNozAIAssistantAPI.schemas';
-
-export interface AssistantErrorResolution {
-	message: string;
-	isRateLimit: boolean;
-}
-
-function isErrorCodeDTO(code: string | undefined): code is ErrorCodeDTO {
-	return (
-		code !== undefined && (Object.values(ErrorCodeDTO) as string[]).includes(code)
-	);
-}
-
-const RATE_LIMIT_ERROR_CODES = new Set<ErrorCodeDTO>([
-	ErrorCodeDTO.rate_limit_override_exceeds_ceiling,
-	ErrorCodeDTO.thread_message_limit,
-	ErrorCodeDTO.connection_limit_exceeded,
-	ErrorCodeDTO.hourly_message_limit,
-	ErrorCodeDTO.daily_message_limit,
-	ErrorCodeDTO.daily_token_limit,
-	ErrorCodeDTO.daily_cost_limit,
-	ErrorCodeDTO.budget_exceeded,
-]);
-
-function isRateLimitError(code: string | undefined, err: unknown): boolean {
-	if (isAxiosError(err) && err.response?.status === 429) {
-		return true;
-	}
-
-	return isErrorCodeDTO(code) && RATE_LIMIT_ERROR_CODES.has(code);
-}
-
-function getErrorBody(err: unknown): ErrorBodyDTO | null {
-	if (isAxiosError(err)) {
-		return (err.response?.data as ErrorResponseDTO | undefined)?.error ?? null;
-	}
-
-	const code = (err as { code?: string } | undefined)?.code;
-	const message = err instanceof Error ? err.message : undefined;
-	if (!code || !message) {
-		return null;
-	}
-
-	return { code: code as ErrorCodeDTO, message };
-}
-
-/**
- * Uses `error.message` when `error.code` is a known `ErrorCodeDTO`;
- * otherwise returns `fallback`.
- */
-export function resolveAssistantErrorMessage(
-	err: unknown,
-	fallback: string,
-): AssistantErrorResolution {
-	const body = getErrorBody(err);
-	const isRateLimit = isRateLimitError(body?.code, err);
-
-	if (body && isErrorCodeDTO(body.code) && body.message.trim()) {
-		return {
-			message: body.message.trim(),
-			isRateLimit,
-		};
-	}
-
-	return { message: fallback, isRateLimit: Boolean(isRateLimit) };
-}

frontend/src/container/Integrations/CloudIntegration/AmazonWebServices/__tests__/mockData.ts

@@ -24,6 +24,7 @@ const accountsResponse: ListAccounts200 = {
 				agentReport: {
 					timestampMillis: 1747114366214,
 					data: null,
+					syncState: null,
 				},
 				providerAccountId: PROVIDER_ACCOUNT_ID,
 				removedAt: null,

frontend/src/pages/IntegrationsDetailsPage/CloudIntegration/stories/__story_mockdata__/cloudIntegration.ts

@@ -295,7 +295,11 @@ const account = (
 	provider,
 	providerAccountId: ACCOUNTS[provider][index],
 	config: accountConfig(provider),
-	agentReport: { timestampMillis: Date.now() - 45 * 1000, data: null },
+	agentReport: {
+		timestampMillis: Date.now() - 45 * 1000,
+		data: null,
+		syncState: null,
+	},
 	createdAt: new Date(Date.now() - 21 * 24 * 60 * 60 * 1000).toISOString(),
 	updatedAt: new Date(Date.now() - 60 * 60 * 1000).toISOString(),
 	removedAt: null,

pkg/apiserver/signozapiserver/authdomain.go

@@ -1,15 +1,18 @@
 package signozapiserver
 
 import (
+	"encoding/json"
 	"net/http"
 	"slices"
 
+	"github.com/SigNoz/signoz/pkg/errors"
 	"github.com/SigNoz/signoz/pkg/http/handler"
 	"github.com/SigNoz/signoz/pkg/types"
 	"github.com/SigNoz/signoz/pkg/types/authtypes"
 	"github.com/SigNoz/signoz/pkg/types/coretypes"
 	"github.com/SigNoz/signoz/pkg/valuer"
 	"github.com/gorilla/mux"
+	"github.com/tidwall/gjson"
 )
 
 func (provider *provider) addAuthDomainRoutes(router *mux.Router) error {
@@ -74,7 +77,7 @@ func (provider *provider) addAuthDomainRoutes(router *mux.Router) error {
 				SourceIDs:      coretypes.OneID(coretypes.ResponseJSONPath("data.id")),
 				SourceSelector: coretypes.WildcardSelector,
 				TargetResource: coretypes.ResourceRole,
-				TargetIDs:      authDomainPostableRoleNamesExtractor(),
+				TargetIDs:      authDomainRoleNamesExtractor(),
 				TargetSelector: coretypes.IDSelector,
 			},
 		),
@@ -146,7 +149,7 @@ func (provider *provider) addAuthDomainRoutes(router *mux.Router) error {
 				SourceIDs:      coretypes.OneID(coretypes.PathParam("id")),
 				SourceSelector: coretypes.IDSelector,
 				TargetResource: coretypes.ResourceRole,
-				TargetIDs:      authDomainUpdatableRoleNamesExtractor(),
+				TargetIDs:      authDomainRoleNamesExtractor(),
 				TargetSelector: coretypes.IDSelector,
 			},
 			handler.AttachDetachSiblingResourceDef{
@@ -196,16 +199,20 @@ func (provider *provider) addAuthDomainRoutes(router *mux.Router) error {
 
 // The extracted names are the roles the request body's mapping grants at SSO
 // login — see authDomainEffectiveRoleNames.
-func authDomainPostableRoleNamesExtractor() coretypes.ResourceIDsExtractor {
-	return coretypes.BodyFields(func(req *authtypes.PostableAuthDomain) []string {
-		return authDomainEffectiveRoleNames(req.RoleMapping)
-	})
-}
+func authDomainRoleNamesExtractor() coretypes.ResourceIDsExtractor {
+	return coretypes.ResourceIDsExtractor{Phase: coretypes.PhaseRequest, Fn: func(ec coretypes.ExtractorContext) ([]string, error) {
+		roleMappingJSON := gjson.GetBytes(ec.RequestBody, "roleMapping")
+		if !roleMappingJSON.Exists() || roleMappingJSON.Type == gjson.Null {
+			return authDomainEffectiveRoleNames(nil), nil
+		}
 
-func authDomainUpdatableRoleNamesExtractor() coretypes.ResourceIDsExtractor {
-	return coretypes.BodyFields(func(req *authtypes.UpdatableAuthDomain) []string {
-		return authDomainEffectiveRoleNames(req.RoleMapping)
-	})
+		roleMapping := new(authtypes.RoleMapping)
+		if err := json.Unmarshal([]byte(roleMappingJSON.Raw), roleMapping); err != nil {
+			return nil, errors.NewInvalidInputf(errors.CodeInvalidInput, "invalid role mapping: %v", err)
+		}
+
+		return authDomainEffectiveRoleNames(roleMapping), nil
+	}}
 }
 
 // The extracted names are the roles the stored domain's mapping grants at SSO

pkg/apiserver/signozapiserver/cloudintegration.go

@@ -350,7 +350,7 @@ func (provider *provider) addCloudIntegrationRoutes(router *mux.Router) error {
 			Resource: coretypes.ResourceMetaResourceCloudIntegration,
 			Verb:     coretypes.VerbRead,
 			Category: coretypes.ActionCategoryDataAccess,
-			ID:       coretypes.BodyField(func(req *citypes.PostableAgentCheckIn) string { return req.ID }),
+			ID:       coretypes.BodyJSONPath("account_id"),
 			Selector: coretypes.IDSelector,
 		}),
 	)).Methods(http.MethodPost).GetError(); err != nil {
@@ -377,12 +377,7 @@ func (provider *provider) addCloudIntegrationRoutes(router *mux.Router) error {
 			Resource: coretypes.ResourceMetaResourceCloudIntegration,
 			Verb:     coretypes.VerbRead,
 			Category: coretypes.ActionCategoryDataAccess,
-			ID: coretypes.BodyField(func(req *citypes.PostableAgentCheckIn) string {
-				if req.CloudIntegrationID.IsZero() {
-					return ""
-				}
-				return req.CloudIntegrationID.StringValue()
-			}),
+			ID:       coretypes.BodyJSONPath("cloudIntegrationId"),
 			Selector: coretypes.IDSelector,
 		}),
 	)).Methods(http.MethodPost).GetError(); err != nil {

pkg/apiserver/signozapiserver/gateway.go

@@ -332,7 +332,7 @@ func (provider *provider) addGatewayRoutes(router *mux.Router) error {
 				Verb:           coretypes.VerbAttach,
 				Category:       coretypes.ActionCategoryConfigurationChange,
 				ParentResource: coretypes.ResourceMetaResourceIngestionKey,
-				ParentID:       coretypes.BodyField(func(req *gatewaytypes.PostableIngestionKeyLimit) string { return req.KeyID }),
+				ParentID:       coretypes.BodyJSONPath("keyId"),
 				ParentSelector: coretypes.IDSelector,
 				ChildResource:  coretypes.ResourceMetaResourceIngestionLimit,
 				ChildIDs:       coretypes.OneID(coretypes.ResponseJSONPath("data.id")),

pkg/apiserver/signozapiserver/prometheus.go

@@ -5,7 +5,6 @@ import (
 	"strings"
 
 	"github.com/SigNoz/signoz/pkg/errors"
-	"github.com/SigNoz/signoz/pkg/http/binding"
 	"github.com/SigNoz/signoz/pkg/http/handler"
 	"github.com/SigNoz/signoz/pkg/http/render"
 	"github.com/SigNoz/signoz/pkg/prometheus"
@@ -80,14 +79,6 @@ func (h *prometheusOpenAPIHandler) ResourceDefs() []handler.ResourceDef {
 	}}
 }
 
-func (h *prometheusOpenAPIHandler) Request() any {
-	return nil
-}
-
-func (h *prometheusOpenAPIHandler) BindBodyOptions() []binding.BindBodyOption {
-	return nil
-}
-
 func (provider *provider) addPrometheusRoutes(router *mux.Router) error {
 	if err := router.Handle("/prometheus/api/v1/query", &prometheusOpenAPIHandler{
 		handlerFunc: provider.authzMiddleware.CheckResources(provider.prometheusHandler.Query, authtypes.SigNozAdminRoleName, authtypes.SigNozEditorRoleName, authtypes.SigNozViewerRoleName),

pkg/apiserver/signozapiserver/querier.go

@@ -461,11 +461,10 @@ func (provider *provider) addQuerierRoutes(router *mux.Router) error {
 		ErrorStatusCodes:    []int{http.StatusBadRequest},
 		SecuritySchemes:     newScopedSecuritySchemes(telemetryReadScopes()),
 	}, handler.WithResourceDefs(handler.TelemetryResourceDef{
-		Verb:         coretypes.VerbRead,
-		Category:     coretypes.ActionCategoryDataAccess,
-		Selector:     querybuilder.TelemetrySelector,
-		Resources:    querybuilder.QueryRangeResources,
-		RequiresBody: true,
+		Verb:      coretypes.VerbRead,
+		Category:  coretypes.ActionCategoryDataAccess,
+		Selector:  querybuilder.TelemetrySelector,
+		Resources: querybuilder.QueryRangeResources,
 	}))).Methods(http.MethodPost).GetError(); err != nil {
 		return err
 	}
@@ -484,11 +483,10 @@ func (provider *provider) addQuerierRoutes(router *mux.Router) error {
 		ErrorStatusCodes:    []int{http.StatusBadRequest},
 		SecuritySchemes:     newScopedSecuritySchemes(telemetryReadScopes()),
 	}, handler.WithResourceDefs(handler.TelemetryResourceDef{
-		Verb:         coretypes.VerbRead,
-		Category:     coretypes.ActionCategoryDataAccess,
-		Selector:     querybuilder.TelemetrySelector,
-		Resources:    querybuilder.QueryRangeResources,
-		RequiresBody: true,
+		Verb:      coretypes.VerbRead,
+		Category:  coretypes.ActionCategoryDataAccess,
+		Selector:  querybuilder.TelemetrySelector,
+		Resources: querybuilder.QueryRangeResources,
 	}))).Methods(http.MethodPost).GetError(); err != nil {
 		return err
 	}

pkg/apiserver/signozapiserver/registry.go

@@ -4,7 +4,6 @@ import (
 	"net/http"
 
 	"github.com/SigNoz/signoz/pkg/factory"
-	"github.com/SigNoz/signoz/pkg/http/binding"
 	pkghandler "github.com/SigNoz/signoz/pkg/http/handler"
 	"github.com/SigNoz/signoz/pkg/http/render"
 	"github.com/gorilla/mux"
@@ -56,14 +55,6 @@ func (handler *healthOpenAPIHandler) ResourceDefs() []pkghandler.ResourceDef {
 	return nil
 }
 
-func (handler *healthOpenAPIHandler) Request() any {
-	return nil
-}
-
-func (handler *healthOpenAPIHandler) BindBodyOptions() []binding.BindBodyOption {
-	return nil
-}
-
 func (provider *provider) addRegistryRoutes(router *mux.Router) error {
 	if err := router.Handle("/api/v2/healthz", newHealthOpenAPIHandler(
 		provider.authzMiddleware.OpenAccess(provider.factoryHandler.Healthz),

pkg/apiserver/signozapiserver/serviceaccount.go

@@ -358,20 +358,10 @@ func (provider *provider) addServiceAccountRoutes(router *mux.Router) error {
 			Verb:           coretypes.VerbAttach,
 			Category:       coretypes.ActionCategoryAccessControl,
 			SourceResource: coretypes.ResourceServiceAccount,
-			SourceIDs: coretypes.OneID(coretypes.BodyField(func(req *serviceaccounttypes.PostableServiceAccountRole) string {
-				if req.ServiceAccountID.IsZero() {
-					return ""
-				}
-				return req.ServiceAccountID.StringValue()
-			})),
+			SourceIDs:      coretypes.OneID(coretypes.BodyJSONPath("serviceAccountId")),
 			SourceSelector: coretypes.IDSelector,
 			TargetResource: coretypes.ResourceRole,
-			TargetIDs: coretypes.OneID(coretypes.BodyField(func(req *serviceaccounttypes.PostableServiceAccountRole) string {
-				if req.RoleID.IsZero() {
-					return ""
-				}
-				return req.RoleID.StringValue()
-			})),
+			TargetIDs:      coretypes.OneID(coretypes.BodyJSONPath("roleId")),
 			TargetSelector: provider.roleSelector,
 		}),
 	)).Methods(http.MethodPost).GetError(); err != nil {

pkg/apiserver/signozapiserver/zeus.go

@@ -68,7 +68,7 @@ func (provider *provider) addZeusRoutes(router *mux.Router) error {
 		Resource: coretypes.ResourceMetaResourceDeploymentHost,
 		Verb:     coretypes.VerbUpdate,
 		Category: coretypes.ActionCategoryConfigurationChange,
-		ID:       coretypes.BodyField(func(req *zeustypes.PostableHost) string { return req.Name }),
+		ID:       coretypes.BodyJSONPath("name"),
 		Selector: coretypes.WildcardSelector,
 	}))).Methods(http.MethodPut).GetError(); err != nil {
 		return err

pkg/gateway/handler.go

@@ -8,7 +8,6 @@ import (
 	"github.com/SigNoz/signoz/pkg/http/binding"
 	"github.com/SigNoz/signoz/pkg/http/render"
 	"github.com/SigNoz/signoz/pkg/types/authtypes"
-	"github.com/SigNoz/signoz/pkg/types/coretypes"
 	"github.com/SigNoz/signoz/pkg/types/gatewaytypes"
 	"github.com/SigNoz/signoz/pkg/valuer"
 	"github.com/gorilla/mux"
@@ -285,8 +284,8 @@ func (handler *handler) CreateIngestionKeyLimit(rw http.ResponseWriter, r *http.
 
 	orgID := valuer.MustNewUUID(claims.OrgID)
 
-	req, err := coretypes.BodyFromContext[gatewaytypes.PostableIngestionKeyLimit](r.Context())
-	if err != nil {
+	var req gatewaytypes.PostableIngestionKeyLimit
+	if err := binding.JSON.BindBody(r.Body, &req); err != nil {
 		render.Error(rw, err)
 		return
 	}

pkg/http/handler/handler.go

@@ -1,13 +1,10 @@
 package handler
 
 import (
-	"fmt"
 	"net/http"
-	"reflect"
 	"slices"
 
 	"github.com/SigNoz/signoz/pkg/errors"
-	"github.com/SigNoz/signoz/pkg/http/binding"
 	"github.com/SigNoz/signoz/pkg/http/render"
 	"github.com/swaggest/openapi-go"
 	"github.com/swaggest/openapi-go/openapi3"
@@ -19,15 +16,12 @@ type Handler interface {
 	http.Handler
 	ServeOpenAPI(openapi.OperationContext)
 	ResourceDefs() []ResourceDef
-	Request() any
-	BindBodyOptions() []binding.BindBodyOption
 }
 
 type handler struct {
-	handlerFunc     http.HandlerFunc
-	openAPIDef      OpenAPIDef
-	resourceDefs    []ResourceDef
-	bindBodyOptions []binding.BindBodyOption
+	handlerFunc  http.HandlerFunc
+	openAPIDef   OpenAPIDef
+	resourceDefs []ResourceDef
 }
 
 func New(handlerFunc http.HandlerFunc, openAPIDef OpenAPIDef, opts ...Option) Handler {
@@ -53,10 +47,6 @@ func New(handlerFunc http.HandlerFunc, openAPIDef OpenAPIDef, opts ...Option) Ha
 		opt(handler)
 	}
 
-	if RequiresBody(handler.resourceDefs) && (openAPIDef.Request == nil || reflect.TypeOf(openAPIDef.Request).Kind() != reflect.Pointer) {
-		panic(fmt.Sprintf("handler %s: a body extractor needs OpenAPIDef.Request to be a pointer, got %T", openAPIDef.ID, openAPIDef.Request))
-	}
-
 	return handler
 }
 
@@ -145,11 +135,3 @@ func (handler *handler) ServeOpenAPI(opCtx openapi.OperationContext) {
 func (handler *handler) ResourceDefs() []ResourceDef {
 	return handler.resourceDefs
 }
-
-func (handler *handler) Request() any {
-	return handler.openAPIDef.Request
-}
-
-func (handler *handler) BindBodyOptions() []binding.BindBodyOption {
-	return handler.bindBodyOptions
-}

pkg/http/handler/handler_test.go

@@ -4,8 +4,6 @@ import (
 	"net/http"
 	"testing"
 
-	"github.com/SigNoz/signoz/pkg/http/binding"
-	"github.com/SigNoz/signoz/pkg/types/coretypes"
 	"github.com/gorilla/mux"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -24,41 +22,6 @@ func (bespokeOpenAPIHandler) ServeOpenAPI(opCtx openapi.OperationContext) {
 
 func (bespokeOpenAPIHandler) ResourceDefs() []ResourceDef { return nil }
 
-func (bespokeOpenAPIHandler) Request() any { return nil }
-
-func (bespokeOpenAPIHandler) BindBodyOptions() []binding.BindBodyOption { return nil }
-
-func TestNewPanicsWhenBodyExtractorHasNoPointerRequest(t *testing.T) {
-	type body struct{ ID string }
-	bodyDef := BasicResourceDef{Resource: coretypes.ResourceRole, Verb: coretypes.VerbRead, ID: coretypes.BodyField(func(req *body) string { return req.ID }), Selector: coretypes.IDSelector}
-	pathDef := BasicResourceDef{Resource: coretypes.ResourceRole, Verb: coretypes.VerbRead, ID: coretypes.PathParam("id"), Selector: coretypes.IDSelector}
-
-	testCases := []struct {
-		name    string
-		request any
-		def     ResourceDef
-		panics  bool
-	}{
-		{name: "BodyExtractor_ValueRequest_Panics", request: body{}, def: bodyDef, panics: true},
-		{name: "BodyExtractor_NilRequest_Panics", request: nil, def: bodyDef, panics: true},
-		{name: "BodyExtractor_PointerRequest_Registers", request: new(body), def: bodyDef, panics: false},
-		{name: "PathExtractor_ValueRequest_Registers", request: body{}, def: pathDef, panics: false},
-	}
-
-	for _, testCase := range testCases {
-		t.Run(testCase.name, func(t *testing.T) {
-			register := func() {
-				New(func(http.ResponseWriter, *http.Request) {}, OpenAPIDef{ID: testCase.name, Request: testCase.request}, WithResourceDefs(testCase.def))
-			}
-			if testCase.panics {
-				assert.Panics(t, register)
-			} else {
-				assert.NotPanics(t, register)
-			}
-		})
-	}
-}
-
 func TestAttachStabilities(t *testing.T) {
 	router := mux.NewRouter()
 	router.Handle("/development", New(func(http.ResponseWriter, *http.Request) {}, OpenAPIDef{ID: "Development", SuccessStatusCode: http.StatusOK, Stability: StabilityDevelopment})).Methods(http.MethodGet)

pkg/http/handler/option.go

@@ -1,7 +1,5 @@
 package handler
 
-import "github.com/SigNoz/signoz/pkg/http/binding"
-
 type Option func(*handler)
 
 func WithResourceDefs(defs ...ResourceDef) Option {
@@ -9,9 +7,3 @@ func WithResourceDefs(defs ...ResourceDef) Option {
 		h.resourceDefs = append(h.resourceDefs, defs...)
 	}
 }
-
-func WithBindBodyOptions(opts ...binding.BindBodyOption) Option {
-	return func(h *handler) {
-		h.bindBodyOptions = append(h.bindBodyOptions, opts...)
-	}
-}

pkg/http/handler/resourcedef.go

@@ -9,7 +9,6 @@ type ResourceDef interface {
 	// resolveRequest is unexported to seal the interface. It returns a slice so a
 	// single def can fan out (e.g. a telemetry query touching multiple signals).
 	resolveRequest(ec coretypes.ExtractorContext) []coretypes.ResolvedResource
-	requiresBody() bool
 }
 
 func ResolveRequest(defs []ResourceDef, ec coretypes.ExtractorContext) []coretypes.ResolvedResource {
@@ -21,17 +20,6 @@ func ResolveRequest(defs []ResourceDef, ec coretypes.ExtractorContext) []coretyp
 	return resolved
 }
 
-// RequiresBody reports whether any def needs the decoded request body.
-func RequiresBody(defs []ResourceDef) bool {
-	for _, def := range defs {
-		if def.requiresBody() {
-			return true
-		}
-	}
-
-	return false
-}
-
 // BasicResourceDef checks a single resource for one verb.
 type BasicResourceDef struct {
 	Resource coretypes.Resource
@@ -54,10 +42,6 @@ func (def BasicResourceDef) resolveRequest(ec coretypes.ExtractorContext) []core
 	}
 }
 
-func (def BasicResourceDef) requiresBody() bool {
-	return def.ID.RequiresBody
-}
-
 // AttachDetachSiblingResourceDef checks an attach/detach between peer resources;
 // both source and target are authz-checked.
 type AttachDetachSiblingResourceDef struct {
@@ -88,10 +72,6 @@ func (def AttachDetachSiblingResourceDef) resolveRequest(ec coretypes.ExtractorC
 	}
 }
 
-func (def AttachDetachSiblingResourceDef) requiresBody() bool {
-	return def.SourceIDs.RequiresBody || def.TargetIDs.RequiresBody
-}
-
 // AttachDetachParentChildResourceDef authz-checks only the parent; the child
 // rides along for audit context.
 type AttachDetachParentChildResourceDef struct {
@@ -121,20 +101,11 @@ func (def AttachDetachParentChildResourceDef) resolveRequest(ec coretypes.Extrac
 	}
 }
 
-func (def AttachDetachParentChildResourceDef) requiresBody() bool {
-	return def.ParentID.RequiresBody || def.ChildIDs.RequiresBody
-}
-
 type TelemetryResourceDef struct {
-	Verb         coretypes.Verb
-	Category     coretypes.ActionCategory
-	Selector     coretypes.SelectorFunc
-	Resources    coretypes.ResourceExtractor
-	RequiresBody bool
-}
-
-func (def TelemetryResourceDef) requiresBody() bool {
-	return def.RequiresBody
+	Verb      coretypes.Verb
+	Category  coretypes.ActionCategory
+	Selector  coretypes.SelectorFunc
+	Resources coretypes.ResourceExtractor
 }
 
 func (def TelemetryResourceDef) resolveRequest(ec coretypes.ExtractorContext) []coretypes.ResolvedResource {

pkg/http/middleware/resource.go

@@ -5,9 +5,7 @@ import (
 	"io"
 	"log/slog"
 	"net/http"
-	"reflect"
 
-	"github.com/SigNoz/signoz/pkg/http/binding"
 	"github.com/SigNoz/signoz/pkg/http/handler"
 	"github.com/SigNoz/signoz/pkg/types/coretypes"
 	"github.com/gorilla/mux"
@@ -25,8 +23,8 @@ func NewResource(logger *slog.Logger) *Resource {
 
 func (middleware *Resource) Wrap(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
-		provider := handlerFromRequest(req)
-		if provider == nil || len(provider.ResourceDefs()) == 0 {
+		defs := resourceDefsFromRequest(req)
+		if len(defs) == 0 {
 			next.ServeHTTP(rw, req)
 			return
 		}
@@ -38,40 +36,18 @@ func (middleware *Resource) Wrap(next http.Handler) http.Handler {
 			req.Body = io.NopCloser(bytes.NewReader(body))
 		}
 
-		defs := provider.ResourceDefs()
-
-		var decoded any
-		var decodeErr error
-		if handler.RequiresBody(defs) {
-			decoded, decodeErr = decodeBody(provider.Request(), body, provider.BindBodyOptions()...)
+		extractorCtx := coretypes.ExtractorContext{
+			Request:     req,
+			RequestBody: body,
 		}
+		resolved := handler.ResolveRequest(defs, extractorCtx)
 
-		extractorCtx := coretypes.ExtractorContext{Request: req, RequestBody: decoded}
-
-		var resolved []coretypes.ResolvedResource
-		if decodeErr != nil {
-			// authz renders the error inside the audit middleware, so the request is still logged
-			resolved = []coretypes.ResolvedResource{coretypes.NewResolvedResourceWithError(coretypes.Verb{}, coretypes.ActionCategory{}, decodeErr)}
-		} else {
-			resolved = handler.ResolveRequest(defs, extractorCtx)
-		}
-
-		ctx := coretypes.NewContextWithExtractorContext(req.Context(), extractorCtx)
-		ctx = coretypes.NewContextWithResolvedResources(ctx, resolved)
+		ctx := coretypes.NewContextWithResolvedResources(req.Context(), resolved)
 		next.ServeHTTP(rw, req.WithContext(ctx))
 	})
 }
 
-func decodeBody(prototype any, body []byte, opts ...binding.BindBodyOption) (any, error) {
-	decoded := reflect.New(reflect.TypeOf(prototype).Elem()).Interface()
-	if err := binding.JSON.BindBody(bytes.NewReader(body), decoded, opts...); err != nil {
-		return nil, err
-	}
-
-	return decoded, nil
-}
-
-func handlerFromRequest(req *http.Request) handler.Handler {
+func resourceDefsFromRequest(req *http.Request) []handler.ResourceDef {
 	route := mux.CurrentRoute(req)
 	if route == nil {
 		return nil
@@ -87,5 +63,5 @@ func handlerFromRequest(req *http.Request) handler.Handler {
 		return nil
 	}
 
-	return provider
+	return provider.ResourceDefs()
 }

pkg/modules/authdomain/implauthdomain/handler.go

@@ -5,11 +5,11 @@ import (
 	"net/http"
 	"time"
 
+	"github.com/SigNoz/signoz/pkg/http/binding"
 	"github.com/SigNoz/signoz/pkg/http/render"
 	"github.com/SigNoz/signoz/pkg/modules/authdomain"
 	"github.com/SigNoz/signoz/pkg/types"
 	"github.com/SigNoz/signoz/pkg/types/authtypes"
-	"github.com/SigNoz/signoz/pkg/types/coretypes"
 	"github.com/SigNoz/signoz/pkg/valuer"
 	"github.com/gorilla/mux"
 )
@@ -32,8 +32,8 @@ func (handler *handler) Create(rw http.ResponseWriter, req *http.Request) {
 		return
 	}
 
-	body, err := coretypes.BodyFromContext[authtypes.PostableAuthDomain](req.Context())
-	if err != nil {
+	body := new(authtypes.PostableAuthDomain)
+	if err := binding.JSON.BindBody(req.Body, body); err != nil {
 		render.Error(rw, err)
 		return
 	}
@@ -142,8 +142,8 @@ func (handler *handler) Update(rw http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	body, err := coretypes.BodyFromContext[authtypes.UpdatableAuthDomain](r.Context())
-	if err != nil {
+	body := new(authtypes.UpdatableAuthDomain)
+	if err := binding.JSON.BindBody(r.Body, body); err != nil {
 		render.Error(rw, err)
 		return
 	}

pkg/modules/cloudintegration/config.go

@@ -22,7 +22,7 @@ func newConfig() factory.Config {
 		Agent: AgentConfig{
 			// we will maintain the latest version of cloud integration agent from here,
 			// till we automate it externally or figure out a way to validate it.
-			Version: "v0.0.14",
+			Version: "v0.0.15",
 		},
 	}
 }

pkg/modules/cloudintegration/implcloudintegration/handler.go

@@ -10,7 +10,6 @@ import (
 	"github.com/SigNoz/signoz/pkg/modules/cloudintegration"
 	"github.com/SigNoz/signoz/pkg/types/authtypes"
 	"github.com/SigNoz/signoz/pkg/types/cloudintegrationtypes"
-	"github.com/SigNoz/signoz/pkg/types/coretypes"
 	"github.com/SigNoz/signoz/pkg/valuer"
 	"github.com/gorilla/mux"
 )
@@ -468,8 +467,8 @@ func (handler *handler) AgentCheckIn(rw http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	req, err := coretypes.BodyFromContext[cloudintegrationtypes.PostableAgentCheckIn](r.Context())
-	if err != nil {
+	req := new(cloudintegrationtypes.PostableAgentCheckIn)
+	if err := binding.JSON.BindBody(r.Body, req); err != nil {
 		render.Error(rw, err)
 		return
 	}

pkg/modules/cloudintegration/implcloudintegration/store.go

@@ -134,6 +134,24 @@ func (store *store) UpdateAccount(ctx context.Context, account *cloudintegration
 		BunDBCtx(ctx).
 		NewUpdate().
 		Model(account).
+		Column("config").
+		Column("updated_at").
+		WherePK().
+		Where("org_id = ?", account.OrgID).
+		Where("provider = ?", account.Provider).
+		Exec(ctx)
+
+	return err
+}
+
+func (store *store) UpdateAgentReport(ctx context.Context, account *cloudintegrationtypes.StorableCloudIntegration) error {
+	_, err := store.
+		store.
+		BunDBCtx(ctx).
+		NewUpdate().
+		Model(account).
+		Column("account_id").
+		Column("last_agent_report").
 		WherePK().
 		Where("org_id = ?", account.OrgID).
 		Where("provider = ?", account.Provider).

pkg/modules/serviceaccount/implserviceaccount/handler.go

@@ -8,7 +8,6 @@ import (
 	"github.com/SigNoz/signoz/pkg/modules/serviceaccount"
 	"github.com/SigNoz/signoz/pkg/types"
 	"github.com/SigNoz/signoz/pkg/types/authtypes"
-	"github.com/SigNoz/signoz/pkg/types/coretypes"
 	"github.com/SigNoz/signoz/pkg/types/serviceaccounttypes"
 	"github.com/SigNoz/signoz/pkg/valuer"
 	"github.com/gorilla/mux"
@@ -223,8 +222,8 @@ func (handler *handler) CreateServiceAccountRole(rw http.ResponseWriter, r *http
 		return
 	}
 
-	req, err := coretypes.BodyFromContext[serviceaccounttypes.PostableServiceAccountRole](r.Context())
-	if err != nil {
+	req := new(serviceaccounttypes.PostableServiceAccountRole)
+	if err := binding.JSON.BindBody(r.Body, req); err != nil {
 		render.Error(rw, err)
 		return
 	}

pkg/querier/api.go

@@ -14,7 +14,6 @@ import (
 	"github.com/SigNoz/signoz/pkg/http/binding"
 	"github.com/SigNoz/signoz/pkg/http/render"
 	"github.com/SigNoz/signoz/pkg/types/authtypes"
-	"github.com/SigNoz/signoz/pkg/types/coretypes"
 	"github.com/SigNoz/signoz/pkg/types/ctxtypes"
 	"github.com/SigNoz/signoz/pkg/types/instrumentationtypes"
 	qbtypes "github.com/SigNoz/signoz/pkg/types/querybuildertypes/querybuildertypesv5"
@@ -53,8 +52,8 @@ func (handler *handler) QueryRange(rw http.ResponseWriter, req *http.Request) {
 		return
 	}
 
-	queryRangeRequest, err := coretypes.BodyFromContext[qbtypes.QueryRangeRequest](req.Context())
-	if err != nil {
+	var queryRangeRequest qbtypes.QueryRangeRequest
+	if err := binding.JSON.BindBody(req.Body, &queryRangeRequest); err != nil {
 		render.Error(rw, err)
 		return
 	}
@@ -71,7 +70,7 @@ func (handler *handler) QueryRange(rw http.ResponseWriter, req *http.Request) {
 		return
 	}
 
-	queryRangeResponse, err := handler.querier.QueryRange(ctx, orgID, queryRangeRequest)
+	queryRangeResponse, err := handler.querier.QueryRange(ctx, orgID, &queryRangeRequest)
 	if err != nil {
 		render.Error(rw, err)
 		return
@@ -97,8 +96,8 @@ func (handler *handler) QueryRangePreview(rw http.ResponseWriter, req *http.Requ
 		return
 	}
 
-	queryRangeRequest, err := coretypes.BodyFromContext[qbtypes.QueryRangeRequest](req.Context())
-	if err != nil {
+	var queryRangeRequest qbtypes.QueryRangeRequest
+	if err := json.NewDecoder(req.Body).Decode(&queryRangeRequest); err != nil {
 		render.Error(rw, err)
 		return
 	}
@@ -119,7 +118,7 @@ func (handler *handler) QueryRangePreview(rw http.ResponseWriter, req *http.Requ
 		return
 	}
 
-	preview, err := handler.querier.QueryRangePreview(ctx, orgID, queryRangeRequest, previewOpts)
+	preview, err := handler.querier.QueryRangePreview(ctx, orgID, &queryRangeRequest, previewOpts)
 	if err != nil {
 		render.Error(rw, err)
 		return

pkg/querybuilder/query_range_resources.go

@@ -2,6 +2,7 @@ package querybuilder
 
 import (
 	"context"
+	"encoding/json"
 	"strings"
 
 	"github.com/SigNoz/signoz/pkg/errors"
@@ -9,6 +10,7 @@ import (
 	qbtypes "github.com/SigNoz/signoz/pkg/types/querybuildertypes/querybuildertypesv5"
 	"github.com/SigNoz/signoz/pkg/types/telemetrytypes"
 	"github.com/SigNoz/signoz/pkg/valuer"
+	"github.com/tidwall/gjson"
 )
 
 func TelemetrySelector(_ context.Context, resource coretypes.Resource, id string, _ valuer.UUID) ([]coretypes.Selector, error) {
@@ -27,19 +29,20 @@ func TelemetrySelector(_ context.Context, resource coretypes.Resource, id string
 }
 
 func QueryRangeResources(ec coretypes.ExtractorContext) ([]coretypes.ResourceWithID, error) {
-	req, err := coretypes.BodyAs[qbtypes.QueryRangeRequest](ec)
+	queries := gjson.GetBytes(ec.RequestBody, "compositeQuery.queries")
+	if !queries.IsArray() || len(queries.Array()) == 0 {
+		return nil, errors.NewInvalidInputf(errors.CodeInvalidInput, "atleast one query is required")
+	}
+
+	variables, err := queryRangeVariables(ec.RequestBody)
 	if err != nil {
 		return nil, err
 	}
 
-	if len(req.CompositeQuery.Queries) == 0 {
-		return nil, errors.NewInvalidInputf(errors.CodeInvalidInput, "atleast one query is required")
-	}
-
-	refs := make([]coretypes.ResourceWithID, 0, len(req.CompositeQuery.Queries))
+	refs := make([]coretypes.ResourceWithID, 0, len(queries.Array()))
 	seen := make(map[string]struct{})
-	for _, query := range req.CompositeQuery.Queries {
-		queryRefs, err := resourcesForQuery(query, req.Variables)
+	for _, query := range queries.Array() {
+		queryRefs, err := resourcesForQuery(query, variables)
 		if err != nil {
 			return nil, err
 		}
@@ -57,6 +60,21 @@ func QueryRangeResources(ec coretypes.ExtractorContext) ([]coretypes.ResourceWit
 	return refs, nil
 }
 
+func queryRangeVariables(body []byte) (map[string]qbtypes.VariableItem, error) {
+	variables := make(map[string]qbtypes.VariableItem)
+
+	raw := gjson.GetBytes(body, "variables")
+	if !raw.Exists() {
+		return variables, nil
+	}
+
+	if err := json.Unmarshal([]byte(raw.Raw), &variables); err != nil {
+		return nil, errors.NewInvalidInputf(errors.CodeInvalidInput, "invalid variables in query range request")
+	}
+
+	return variables, nil
+}
+
 // PromQLResources is the resource set of a bare PromQL query: metrics on
 // the promql wildcard, the same ID resourcesForQuery assigns to a PromQL
 // query inside a composite — one grant covers both entry points.
@@ -67,53 +85,42 @@ func PromQLResources(coretypes.ExtractorContext) ([]coretypes.ResourceWithID, er
 	}}, nil
 }
 
-func resourcesForQuery(query qbtypes.QueryEnvelope, variables map[string]qbtypes.VariableItem) ([]coretypes.ResourceWithID, error) {
-	queryType := query.Type.StringValue()
+func resourcesForQuery(query gjson.Result, variables map[string]qbtypes.VariableItem) ([]coretypes.ResourceWithID, error) {
+	queryType := query.Get("type").String()
 	typeWildcard := queryType + "/" + coretypes.WildCardSelectorString
 
-	switch query.Type {
-	case qbtypes.QueryTypeBuilder, qbtypes.QueryTypeSubQuery:
-		return resourcesForBuilderQuery(queryType, query.Spec, variables)
-	case qbtypes.QueryTypeBuilderAI:
+	switch queryType {
+	case qbtypes.QueryTypeBuilder.StringValue(), qbtypes.QueryTypeSubQuery.StringValue():
+		return resourcesForBuilderQuery(queryType, query.Get("spec"), variables)
+	case qbtypes.QueryTypeBuilderAI.StringValue():
 		// always a traces query; the signal may be absent from the payload
-		_, _, expression, err := builderQuerySpec(query.Spec)
-		if err != nil {
-			return nil, err
-		}
-
-		return builderQueryResourceRefs(queryType, coretypes.ResourceTelemetryResourceTraces, expression, variables)
-	case qbtypes.QueryTypePromQL:
+		return builderQueryResourceRefs(queryType, coretypes.ResourceTelemetryResourceTraces, query.Get("spec"), variables)
+	case qbtypes.QueryTypePromQL.StringValue():
 		return []coretypes.ResourceWithID{{Resource: coretypes.ResourceTelemetryResourceMetrics, ID: typeWildcard}}, nil
-	case qbtypes.QueryTypeClickHouseSQL:
+	case qbtypes.QueryTypeClickHouseSQL.StringValue():
 		return []coretypes.ResourceWithID{
 			{Resource: coretypes.ResourceTelemetryResourceLogs, ID: typeWildcard},
 			{Resource: coretypes.ResourceTelemetryResourceTraces, ID: typeWildcard},
 			{Resource: coretypes.ResourceTelemetryResourceMetrics, ID: typeWildcard},
 			{Resource: coretypes.ResourceTelemetryResourceMeterMetrics, ID: typeWildcard},
 		}, nil
-	case qbtypes.QueryTypeFormula, qbtypes.QueryTypeJoin, qbtypes.QueryTypeTraceOperator:
+	case qbtypes.QueryTypeFormula.StringValue(), qbtypes.QueryTypeJoin.StringValue(), qbtypes.QueryTypeTraceOperator.StringValue():
 		return nil, nil
 	default:
 		return nil, errors.NewInvalidInputf(errors.CodeInvalidInput, "unsupported query type %q", queryType)
 	}
 }
 
-func resourcesForBuilderQuery(queryType string, spec any, variables map[string]qbtypes.VariableItem) ([]coretypes.ResourceWithID, error) {
-	signal, source, expression, err := builderQuerySpec(spec)
+func resourcesForBuilderQuery(queryType string, spec gjson.Result, variables map[string]qbtypes.VariableItem) ([]coretypes.ResourceWithID, error) {
+	resource, err := builderQueryResource(spec)
 	if err != nil {
 		return nil, err
 	}
-
-	resource, err := builderQueryResource(signal, source)
-	if err != nil {
-		return nil, err
-	}
-
-	return builderQueryResourceRefs(queryType, resource, expression, variables)
+	return builderQueryResourceRefs(queryType, resource, spec, variables)
 }
 
-func builderQueryResourceRefs(queryType string, resource coretypes.Resource, expression string, variables map[string]qbtypes.VariableItem) ([]coretypes.ResourceWithID, error) {
-	ids, err := builderQuerySelectors(queryType, expression, variables)
+func builderQueryResourceRefs(queryType string, resource coretypes.Resource, spec gjson.Result, variables map[string]qbtypes.VariableItem) ([]coretypes.ResourceWithID, error) {
+	ids, err := builderQuerySelectors(queryType, spec.Get("filter.expression").String(), variables)
 	if err != nil {
 		return nil, err
 	}
@@ -126,46 +133,27 @@ func builderQueryResourceRefs(queryType string, resource coretypes.Resource, exp
 	return refs, nil
 }
 
-func builderQueryResource(signal telemetrytypes.Signal, source telemetrytypes.Source) (coretypes.Resource, error) {
-	switch signal {
-	case telemetrytypes.SignalTraces:
+func builderQueryResource(spec gjson.Result) (coretypes.Resource, error) {
+	source := spec.Get("source").String()
+
+	switch spec.Get("signal").String() {
+	case telemetrytypes.SignalTraces.StringValue():
 		return coretypes.ResourceTelemetryResourceTraces, nil
-	case telemetrytypes.SignalLogs:
-		if source == telemetrytypes.SourceAudit {
+	case telemetrytypes.SignalLogs.StringValue():
+		if source == telemetrytypes.SourceAudit.StringValue() {
 			return coretypes.ResourceTelemetryResourceAuditLogs, nil
 		}
 		return coretypes.ResourceTelemetryResourceLogs, nil
-	case telemetrytypes.SignalMetrics:
-		if source == telemetrytypes.SourceMeter {
+	case telemetrytypes.SignalMetrics.StringValue():
+		if source == telemetrytypes.SourceMeter.StringValue() {
 			return coretypes.ResourceTelemetryResourceMeterMetrics, nil
 		}
 		return coretypes.ResourceTelemetryResourceMetrics, nil
 	default:
-		return nil, errors.NewInvalidInputf(errors.CodeInvalidInput, "unsupported signal %q", signal.StringValue())
+		return nil, errors.NewInvalidInputf(errors.CodeInvalidInput, "unsupported signal %q", spec.Get("signal").String())
 	}
 }
 
-func builderQuerySpec(spec any) (telemetrytypes.Signal, telemetrytypes.Source, string, error) {
-	switch typed := spec.(type) {
-	case qbtypes.QueryBuilderQuery[qbtypes.TraceAggregation]:
-		return typed.Signal, typed.Source, filterExpression(typed.Filter), nil
-	case qbtypes.QueryBuilderQuery[qbtypes.LogAggregation]:
-		return typed.Signal, typed.Source, filterExpression(typed.Filter), nil
-	case qbtypes.QueryBuilderQuery[qbtypes.MetricAggregation]:
-		return typed.Signal, typed.Source, filterExpression(typed.Filter), nil
-	default:
-		return telemetrytypes.Signal{}, telemetrytypes.Source{}, "", errors.Newf(errors.TypeInternal, errors.CodeInternal, "unexpected builder query spec %T", spec)
-	}
-}
-
-func filterExpression(filter *qbtypes.Filter) string {
-	if filter == nil {
-		return ""
-	}
-
-	return filter.Expression
-}
-
 func builderQuerySelectors(queryType, expression string, variables map[string]qbtypes.VariableItem) ([]string, error) {
 	typeWildcard := queryType + "/" + coretypes.WildCardSelectorString
 

pkg/querybuilder/query_range_resources_test.go

@@ -2,24 +2,14 @@ package querybuilder
 
 import (
 	"context"
-	"strings"
 	"testing"
 
-	"github.com/SigNoz/signoz/pkg/http/binding"
 	"github.com/SigNoz/signoz/pkg/types/coretypes"
-	qbtypes "github.com/SigNoz/signoz/pkg/types/querybuildertypes/querybuildertypesv5"
 	"github.com/SigNoz/signoz/pkg/valuer"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
-func queryRangeExtractorContext(t *testing.T, body string) coretypes.ExtractorContext {
-	t.Helper()
-	req := new(qbtypes.QueryRangeRequest)
-	require.NoError(t, binding.JSON.BindBody(strings.NewReader(body), req))
-	return coretypes.ExtractorContext{RequestBody: req}
-}
-
 func builderQueryBody(signal, filterExpression string) string {
 	return `{"compositeQuery":{"queries":[{"type":"builder_query","spec":{"signal":"` + signal + `","filter":{"expression":"` + filterExpression + `"}}}]}}`
 }
@@ -152,13 +142,6 @@ func TestQueryRangeResources(t *testing.T) {
 				{Resource: coretypes.ResourceTelemetryResourceLogs, ID: "builder_query/signoz.workspace.key.id/checkout"},
 			},
 		},
-		{
-			name: "DuplicateSignalKey_LastValueWins",
-			body: `{"compositeQuery":{"queries":[{"type":"builder_query","spec":{"signal":"logs","signal":"traces","filter":{"expression":"signoz.workspace.key.id = 'a'"}}}]}}`,
-			expected: []coretypes.ResourceWithID{
-				{Resource: coretypes.ResourceTelemetryResourceTraces, ID: "builder_query/signoz.workspace.key.id/a"},
-			},
-		},
 		{
 			name: "duplicate queries dedupe",
 			body: `{"compositeQuery":{"queries":[{"type":"builder_query","spec":{"signal":"logs","filter":{"expression":"signoz.workspace.key.id = 'a'"}}},{"type":"builder_query","spec":{"signal":"logs","filter":{"expression":"signoz.workspace.key.id='a'"}}}]}}`,
@@ -170,7 +153,7 @@ func TestQueryRangeResources(t *testing.T) {
 
 	for _, testCase := range testCases {
 		t.Run(testCase.name, func(t *testing.T) {
-			refs, err := QueryRangeResources(queryRangeExtractorContext(t, testCase.body))
+			refs, err := QueryRangeResources(coretypes.ExtractorContext{RequestBody: []byte(testCase.body)})
 			require.NoError(t, err)
 			assert.Equal(t, testCase.expected, refs)
 		})
@@ -182,20 +165,14 @@ func TestQueryRangeResourcesErrors(t *testing.T) {
 		`{"compositeQuery":{"queries":[]}}`,
 		`{}`,
 		builderQueryBody("logs", "signoz.workspace.key.id = "),
+		`{"compositeQuery":{"queries":[{"type":"builder_query","spec":{"signal":"unknown"}}]}}`,
+		`{"compositeQuery":{"queries":[{"type":"unknown_type"}]}}`,
 	}
 
 	for _, body := range bodies {
-		_, err := QueryRangeResources(queryRangeExtractorContext(t, body))
+		_, err := QueryRangeResources(coretypes.ExtractorContext{RequestBody: []byte(body)})
 		assert.Error(t, err, "body %s", body)
 	}
-
-	// rejected by the decode the middleware runs, before any extractor
-	for _, body := range []string{
-		`{"compositeQuery":{"queries":[{"type":"builder_query","spec":{"signal":"unknown"}}]}}`,
-		`{"compositeQuery":{"queries":[{"type":"unknown_type"}]}}`,
-	} {
-		assert.Error(t, binding.JSON.BindBody(strings.NewReader(body), new(qbtypes.QueryRangeRequest)), "body %s", body)
-	}
 }
 
 func TestTelemetrySelector(t *testing.T) {

pkg/types/cloudintegrationtypes/account.go

@@ -3,6 +3,7 @@ package cloudintegrationtypes
 import (
 	"encoding/json"
 	"fmt"
+	"maps"
 	"time"
 
 	"github.com/SigNoz/signoz/pkg/errors"
@@ -26,6 +27,17 @@ type Account struct {
 type AgentReport struct {
 	TimestampMillis int64          `json:"timestampMillis" required:"true"`
 	Data            map[string]any `json:"data" required:"true" nullable:"true"`
+	SyncState       *SyncState     `json:"syncState" required:"true" nullable:"true"`
+}
+
+type SyncState struct {
+	Version int64                       `json:"version" required:"true"`
+	InSync  bool                        `json:"inSync" required:"true"`
+	Regions map[string]*RegionSyncState `json:"regions" required:"true" nullable:"false"`
+}
+
+type RegionSyncState struct {
+	State RegionState `json:"state" required:"true"`
 }
 
 type AccountConfig struct {
@@ -150,6 +162,7 @@ func NewAccountFromStorable(storableAccount *StorableCloudIntegration) (*Account
 		account.AgentReport = &AgentReport{
 			TimestampMillis: storableAccount.LastAgentReport.TimestampMillis,
 			Data:            storableAccount.LastAgentReport.Data,
+			SyncState:       NewSyncStateFromStorable(storableAccount.LastAgentReport.SyncState),
 		}
 	}
 
@@ -308,10 +321,28 @@ func NewAccountConfigFromUpdatable(provider CloudProviderType, config *Updatable
 	}
 }
 
-func NewAgentReport(data map[string]any) *AgentReport {
+func NewAgentReport(data map[string]any, syncState *SyncState) *AgentReport {
 	return &AgentReport{
 		TimestampMillis: time.Now().UnixMilli(),
 		Data:            data,
+		SyncState:       syncState,
+	}
+}
+
+func NewSyncStateFromStorable(storableSyncState *StorableSyncState) *SyncState {
+	if storableSyncState == nil {
+		return nil
+	}
+
+	regions := make(map[string]*RegionSyncState, len(storableSyncState.Regions))
+	for region, regionSyncState := range storableSyncState.Regions {
+		regions[region] = &RegionSyncState{State: regionSyncState.State}
+	}
+
+	return &SyncState{
+		Version: storableSyncState.Version,
+		InSync:  storableSyncState.InSync,
+		Regions: regions,
 	}
 }
 
@@ -335,6 +366,40 @@ func (account *Account) Update(provider CloudProviderType, config *AccountConfig
 	return nil
 }
 
+func (account *Account) UpdateAgentReport(providerAccountID *string, agentReport *AgentReport) {
+	account.ProviderAccountID = providerAccountID
+	account.AgentReport = agentReport
+}
+
+// UpdateSyncState keeps the rest of the agent report, and is a no-op when the agent has never checked in.
+func (account *Account) UpdateSyncState(syncState *SyncState) {
+	if account.AgentReport == nil {
+		return
+	}
+
+	account.AgentReport.SyncState = syncState
+}
+
+// NextSyncState returns the sync state for this check-in, or nil for providers without one.
+func (account *Account) NextSyncState(syncedVersion *int64) *SyncState {
+	if account.Provider != CloudProviderTypeAWS {
+		return nil
+	}
+
+	var previous *SyncState
+	if account.AgentReport != nil {
+		previous = account.AgentReport.SyncState
+	}
+
+	regions := account.Config.AWS.Regions
+	// Removed before the agent ever checked in: no region was sent to it, so there is nothing to clean up.
+	if account.AgentReport == nil && account.RemovedAt != nil {
+		regions = nil
+	}
+
+	return newSyncState(previous, regions, account.RemovedAt != nil, syncedVersion)
+}
+
 func (postableAccount *PostableAccount) UnmarshalJSON(data []byte) error {
 	type Alias PostableAccount
 
@@ -406,3 +471,79 @@ func (config *AccountConfig) ToJSON() ([]byte, error) {
 func NewIngestionKeyName(provider CloudProviderType) string {
 	return fmt.Sprintf("%s-integration", provider.StringValue())
 }
+
+// newSyncState returns the sync state after a check-in without mutating previous.
+func newSyncState(previous *SyncState, regions []string, removed bool, syncedVersion *int64) *SyncState {
+	if previous == nil {
+		previous = newSyncStateFromRegions(regions)
+	}
+
+	next := previous.copy()
+
+	// The agent synced this version, so its disabled regions are cleaned up and can be dropped.
+	if syncedVersion != nil && *syncedVersion == next.Version {
+		next.InSync = true
+		maps.DeleteFunc(next.Regions, func(_ string, regionSyncState *RegionSyncState) bool {
+			return regionSyncState.State == RegionStateDisabled
+		})
+	}
+
+	// Once the integration is removed, every region is disabled.
+	if removed {
+		regions = nil
+	}
+
+	changed := false
+	desiredRegionsMap := make(map[string]struct{}, len(regions))
+
+	for _, region := range regions {
+		desiredRegionsMap[region] = struct{}{}
+
+		if regionSyncState, ok := next.Regions[region]; ok && regionSyncState.State == RegionStateEnabled {
+			continue
+		}
+
+		next.Regions[region] = &RegionSyncState{State: RegionStateEnabled}
+		changed = true
+	}
+
+	for region, regionSyncState := range next.Regions {
+		_, ok := desiredRegionsMap[region]
+		if ok && regionSyncState.State == RegionStateEnabled {
+			continue
+		}
+
+		if !ok && regionSyncState.State == RegionStateDisabled {
+			continue
+		}
+
+		regionSyncState.State = RegionStateDisabled
+		changed = true
+	}
+
+	if changed {
+		next.Version++
+		next.InSync = false
+	}
+
+	return next
+}
+
+// newSyncStateFromRegions is used on the first check-in, when the agent has already deployed regions, so it starts in sync.
+func newSyncStateFromRegions(regions []string) *SyncState {
+	syncState := &SyncState{Version: 1, InSync: true, Regions: make(map[string]*RegionSyncState, len(regions))}
+	for _, region := range regions {
+		syncState.Regions[region] = &RegionSyncState{State: RegionStateEnabled}
+	}
+
+	return syncState
+}
+
+func (syncState *SyncState) copy() *SyncState {
+	regions := make(map[string]*RegionSyncState, len(syncState.Regions))
+	for region, regionSyncState := range syncState.Regions {
+		regions[region] = &RegionSyncState{State: regionSyncState.State}
+	}
+
+	return &SyncState{Version: syncState.Version, InSync: syncState.InSync, Regions: regions}
+}

pkg/types/cloudintegrationtypes/checkin.go

@@ -12,7 +12,8 @@ type AgentCheckInRequest struct {
 	ProviderAccountID  string      `json:"providerAccountId" required:"false"`
 	CloudIntegrationID valuer.UUID `json:"cloudIntegrationId" required:"false"`
 
-	Data map[string]any `json:"data" required:"true" nullable:"true"`
+	Data          map[string]any `json:"data" required:"true" nullable:"true"`
+	SyncedVersion *int64         `json:"syncedVersion" required:"false" nullable:"true"`
 }
 
 type PostableAgentCheckIn struct {
@@ -28,6 +29,7 @@ type AgentCheckInResponse struct {
 	ProviderAccountID  string                     `json:"providerAccountId" required:"true"`
 	IntegrationConfig  *ProviderIntegrationConfig `json:"integrationConfig" required:"true"`
 	RemovedAt          *time.Time                 `json:"removedAt" required:"true" nullable:"true"`
+	SyncState          *SyncState                 `json:"syncState" required:"true" nullable:"true"`
 }
 
 type GettableAgentCheckIn struct {
@@ -73,12 +75,13 @@ func NewGettableAgentCheckIn(provider CloudProviderType, resp *AgentCheckInRespo
 	return gettable
 }
 
-func NewAgentCheckInResponse(providerAccountID, cloudIntegrationID string, integrationConfig *ProviderIntegrationConfig, removedAt *time.Time) *AgentCheckInResponse {
+func NewAgentCheckInResponse(providerAccountID, cloudIntegrationID string, integrationConfig *ProviderIntegrationConfig, removedAt *time.Time, syncState *SyncState) *AgentCheckInResponse {
 	return &AgentCheckInResponse{
 		CloudIntegrationID: cloudIntegrationID,
 		ProviderAccountID:  providerAccountID,
 		IntegrationConfig:  integrationConfig,
 		RemovedAt:          removedAt,
+		SyncState:          syncState,
 	}
 }
 

pkg/types/cloudintegrationtypes/cloudintegration.go

@@ -25,6 +25,17 @@ var (
 	ErrCodeServiceDefinitionNotFound            = errors.MustNewCode("service_definition_not_found")
 )
 
+var (
+	RegionStateEnabled  = RegionState{valuer.NewString("enabled")}
+	RegionStateDisabled = RegionState{valuer.NewString("disabled")}
+)
+
+type RegionState struct{ valuer.String }
+
+func (RegionState) Enum() []any {
+	return []any{RegionStateEnabled, RegionStateDisabled}
+}
+
 // StorableCloudIntegration represents a cloud integration stored in the database.
 // This is also referred as "Account" in the context of cloud integrations.
 type StorableCloudIntegration struct {
@@ -43,8 +54,16 @@ type StorableCloudIntegration struct {
 // StorableAgentReport represents the last heartbeat and arbitrary data sent by the agent
 // as of now there is no use case for Data field, but keeping it for backwards compatibility with older structure.
 type StorableAgentReport struct {
-	TimestampMillis int64          `json:"timestamp_millis"` // backward compatibility
-	Data            map[string]any `json:"data"`
+	TimestampMillis int64              `json:"timestamp_millis"` // backward compatibility
+	Data            map[string]any     `json:"data"`
+	SyncState       *StorableSyncState `json:"sync_state,omitempty"`
+}
+
+// StorableSyncState holds every region sent to the agent. A disabled region is dropped only after the agent acks Version.
+type StorableSyncState struct {
+	Version int64                       `json:"version"`
+	InSync  bool                        `json:"in_sync"`
+	Regions map[string]*RegionSyncState `json:"regions"`
 }
 
 // StorableCloudIntegrationService is to store service config for a cloud integration, which is a cloud provider specific configuration.
@@ -148,12 +167,30 @@ func NewStorableCloudIntegration(account *Account) (*StorableCloudIntegration, e
 		storableAccount.LastAgentReport = &StorableAgentReport{
 			TimestampMillis: account.AgentReport.TimestampMillis,
 			Data:            account.AgentReport.Data,
+			SyncState:       NewStorableSyncState(account.AgentReport.SyncState),
 		}
 	}
 
 	return storableAccount, nil
 }
 
+func NewStorableSyncState(syncState *SyncState) *StorableSyncState {
+	if syncState == nil {
+		return nil
+	}
+
+	regions := make(map[string]*RegionSyncState, len(syncState.Regions))
+	for region, regionSyncState := range syncState.Regions {
+		regions[region] = &RegionSyncState{State: regionSyncState.State}
+	}
+
+	return &StorableSyncState{
+		Version: syncState.Version,
+		InSync:  syncState.InSync,
+		Regions: regions,
+	}
+}
+
 // NewStorableCloudIntegrationService creates a new StorableCloudIntegrationService with
 // generated ID and timestamps from a CloudIntegrationService and its serialized config JSON.
 func NewStorableCloudIntegrationService(svc *CloudIntegrationService, configJSON string) *StorableCloudIntegrationService {
@@ -172,6 +209,7 @@ func (account *StorableCloudIntegration) Update(providerAccountID *string, agent
 		account.LastAgentReport = &StorableAgentReport{
 			TimestampMillis: agentReport.TimestampMillis,
 			Data:            agentReport.Data,
+			SyncState:       NewStorableSyncState(agentReport.SyncState),
 		}
 	}
 }

pkg/types/cloudintegrationtypes/store.go

@@ -25,9 +25,12 @@ type Store interface {
 	// CreateAccount creates a new cloud integration account
 	CreateAccount(ctx context.Context, account *StorableCloudIntegration) error
 
-	// UpdateAccount updates an existing cloud integration account
+	// UpdateAccount updates the user updatable fields (config) of an existing cloud integration account
 	UpdateAccount(ctx context.Context, account *StorableCloudIntegration) error
 
+	// UpdateAgentReport updates the provider account id and last agent report of an existing cloud integration account
+	UpdateAgentReport(ctx context.Context, account *StorableCloudIntegration) error
+
 	// RemoveAccount marks a cloud integration account as removed by setting the RemovedAt field
 	RemoveAccount(ctx context.Context, orgID, id valuer.UUID, provider CloudProviderType) error
 

pkg/types/coretypes/extractor.go

@@ -1,10 +1,8 @@
 package coretypes
 
 import (
-	"context"
 	"net/http"
 
-	"github.com/SigNoz/signoz/pkg/errors"
 	"github.com/gorilla/mux"
 	"github.com/tidwall/gjson"
 )
@@ -14,70 +12,24 @@ const (
 	PhaseResponse
 )
 
-var (
-	errCodeExtractorContextNotFound = errors.MustNewCode("extractor_context_not_found")
-	errCodeRequestTypeUndeclared    = errors.MustNewCode("request_type_undeclared")
-	errCodeRequestTypeMismatch      = errors.MustNewCode("request_type_mismatch")
-)
-
 type ExtractPhase int
 
-type extractorContextKey struct{}
-
 // ExtractorContext carries everything an extractor may read: Request + RequestBody
-// are filled pre-handler, ResponseBody post-handler. RequestBody is the body
-// decoded by the resource middleware into the route's declared request type.
+// are filled pre-handler, ResponseBody post-handler.
 type ExtractorContext struct {
 	Request      *http.Request
-	RequestBody  any
+	RequestBody  []byte
 	ResponseBody []byte
 }
 
-func NewContextWithExtractorContext(ctx context.Context, ec ExtractorContext) context.Context {
-	return context.WithValue(ctx, extractorContextKey{}, ec)
-}
-
-func ExtractorContextFromContext(ctx context.Context) (ExtractorContext, error) {
-	ec, ok := ctx.Value(extractorContextKey{}).(ExtractorContext)
-	if !ok {
-		return ExtractorContext{}, errors.New(errors.TypeInternal, errCodeExtractorContextNotFound, "extractor context not found in context")
-	}
-
-	return ec, nil
-}
-
-func BodyAs[T any](ec ExtractorContext) (*T, error) {
-	if ec.RequestBody == nil {
-		return nil, errors.New(errors.TypeInternal, errCodeRequestTypeUndeclared, "route does not declare a request type")
-	}
-
-	typed, ok := ec.RequestBody.(*T)
-	if !ok {
-		return nil, errors.Newf(errors.TypeInternal, errCodeRequestTypeMismatch, "route declares request type %T, expected %T", ec.RequestBody, (*T)(nil))
-	}
-
-	return typed, nil
-}
-
-func BodyFromContext[T any](ctx context.Context) (*T, error) {
-	ec, err := ExtractorContextFromContext(ctx)
-	if err != nil {
-		return nil, err
-	}
-
-	return BodyAs[T](ec)
-}
-
 type ResourceIDExtractor struct {
-	Phase        ExtractPhase
-	RequiresBody bool
-	Fn           func(ExtractorContext) (string, error)
+	Phase ExtractPhase
+	Fn    func(ExtractorContext) (string, error)
 }
 
 type ResourceIDsExtractor struct {
-	Phase        ExtractPhase
-	RequiresBody bool
-	Fn           func(ExtractorContext) ([]string, error)
+	Phase ExtractPhase
+	Fn    func(ExtractorContext) ([]string, error)
 }
 
 func NewResourceIDExtractor(phase ExtractPhase, fn func(ExtractorContext) (string, error)) ResourceIDExtractor {
@@ -98,7 +50,7 @@ func OneID(extractor ResourceIDExtractor) ResourceIDsExtractor {
 		return ResourceIDsExtractor{}
 	}
 
-	return ResourceIDsExtractor{Phase: extractor.Phase, RequiresBody: extractor.RequiresBody, Fn: func(ec ExtractorContext) ([]string, error) {
+	return ResourceIDsExtractor{Phase: extractor.Phase, Fn: func(ec ExtractorContext) ([]string, error) {
 		id, err := extractor.Fn(ec)
 		if err != nil || id == "" {
 			return nil, err
@@ -123,25 +75,26 @@ func PathParam(name string) ResourceIDExtractor {
 	}}
 }
 
-func BodyField[T any](pick func(*T) string) ResourceIDExtractor {
-	return ResourceIDExtractor{Phase: PhaseRequest, RequiresBody: true, Fn: func(ec ExtractorContext) (string, error) {
-		req, err := BodyAs[T](ec)
-		if err != nil {
-			return "", err
-		}
-
-		return pick(req), nil
+func BodyJSONPath(path string) ResourceIDExtractor {
+	return ResourceIDExtractor{Phase: PhaseRequest, Fn: func(ec ExtractorContext) (string, error) {
+		return gjson.GetBytes(ec.RequestBody, path).String(), nil
 	}}
 }
 
-func BodyFields[T any](pick func(*T) []string) ResourceIDsExtractor {
-	return ResourceIDsExtractor{Phase: PhaseRequest, RequiresBody: true, Fn: func(ec ExtractorContext) ([]string, error) {
-		req, err := BodyAs[T](ec)
-		if err != nil {
-			return nil, err
+func BodyJSONArray(path string) ResourceIDsExtractor {
+	return ResourceIDsExtractor{Phase: PhaseRequest, Fn: func(ec ExtractorContext) ([]string, error) {
+		result := gjson.GetBytes(ec.RequestBody, path)
+		if !result.Exists() {
+			return nil, nil
 		}
 
-		return pick(req), nil
+		array := result.Array()
+		ids := make([]string, 0, len(array))
+		for _, r := range array {
+			ids = append(ids, r.String())
+		}
+
+		return ids, nil
 	}}
 }
 

pkg/zeus/handler.go

@@ -8,7 +8,6 @@ import (
 	"github.com/SigNoz/signoz/pkg/http/render"
 	"github.com/SigNoz/signoz/pkg/licensing"
 	"github.com/SigNoz/signoz/pkg/types/authtypes"
-	"github.com/SigNoz/signoz/pkg/types/coretypes"
 	"github.com/SigNoz/signoz/pkg/types/zeustypes"
 	"github.com/SigNoz/signoz/pkg/valuer"
 )
@@ -95,8 +94,8 @@ func (h *handler) PutHost(rw http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	req, err := coretypes.BodyFromContext[zeustypes.PostableHost](r.Context())
-	if err != nil {
+	req := new(zeustypes.PostableHost)
+	if err := binding.JSON.BindBody(r.Body, req); err != nil {
 		render.Error(rw, err)
 		return
 	}

tests/fixtures/cloudintegrations.py

@@ -34,6 +34,8 @@ class ProviderAccountSpec:
     expected_config: Callable[[dict], dict]
     # only the suites that exercise updates need to supply it.
     updated_params: dict = field(default_factory=dict)
+    # params -> the agentReport.syncState the API is expected to return after the first check-in.
+    expected_sync_state: Callable[[dict], dict | None] = lambda p: None
     # id shown in parametrized test names; defaults to the provider slug.
     id: str = field(default="")
 
@@ -315,6 +317,7 @@ def simulate_agent_checkin(
     account_id: str,
     cloud_account_id: str,
     data: dict | None = None,
+    synced_version: int | None = None,
 ) -> requests.Response:
     endpoint = f"/api/v1/cloud_integrations/{cloud_provider}/accounts/check_in"
 
@@ -323,6 +326,8 @@ def simulate_agent_checkin(
         "providerAccountId": cloud_account_id,
         "data": data or {},
     }
+    if synced_version is not None:
+        checkin_payload["syncedVersion"] = synced_version
 
     response = requests.post(
         signoz.self.host_configs["8080"].get(endpoint),

tests/integration/tests/cloudintegrations/03_agent_check_in.py

@@ -3,6 +3,7 @@ from collections.abc import Callable
 from http import HTTPStatus
 
 import pytest
+import requests
 
 from fixtures import types
 from fixtures.auth import USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD, add_license
@@ -152,3 +153,230 @@ def test_duplicate_cloud_account_checkins(
     # Second check-in: account2 tries to claim the same provider account ID → 409
     response = simulate_agent_checkin(signoz, admin_token, spec.provider, account2["id"], same_provider_account_id)
     assert response.status_code == HTTPStatus.CONFLICT, f"Expected 409 for duplicate providerAccountId, got {response.status_code}: {response.text}"
+
+
+def test_sync_state_drops_removed_region_after_ack(
+    signoz: types.SigNoz,
+    create_user_admin: types.Operation,  # pylint: disable=unused-argument
+    get_token: Callable[[str, str], str],
+    create_cloud_integration_account: Callable,
+) -> None:
+    admin_token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
+    account_id = create_cloud_integration_account(admin_token, "aws", regions=["us-east-1", "us-west-2"])["id"]
+    provider_account_id = str(uuid.uuid4())
+
+    response = simulate_agent_checkin(signoz, admin_token, "aws", account_id, provider_account_id)
+    assert response.status_code == HTTPStatus.OK, response.text
+
+    response = requests.put(
+        signoz.self.host_configs["8080"].get(f"/api/v1/cloud_integrations/aws/accounts/{account_id}"),
+        headers={"Authorization": f"Bearer {admin_token}"},
+        json={"config": {"aws": {"regions": ["us-east-1"]}}},
+        timeout=10,
+    )
+    assert response.status_code == HTTPStatus.NO_CONTENT, response.text
+
+    response = simulate_agent_checkin(signoz, admin_token, "aws", account_id, provider_account_id)
+    assert response.status_code == HTTPStatus.OK, response.text
+    assert response.json()["data"]["syncState"] == {
+        "version": 2,
+        "inSync": False,
+        "regions": {"us-east-1": {"state": "enabled"}, "us-west-2": {"state": "disabled"}},
+    }, "removed region should be marked disabled and the version bumped"
+
+    response = simulate_agent_checkin(signoz, admin_token, "aws", account_id, provider_account_id, synced_version=2)
+    assert response.status_code == HTTPStatus.OK, response.text
+    assert response.json()["data"]["syncState"] == {
+        "version": 2,
+        "inSync": True,
+        "regions": {"us-east-1": {"state": "enabled"}},
+    }, "acked removed region should be dropped"
+
+
+def test_sync_state_keeps_removed_region_without_ack(
+    signoz: types.SigNoz,
+    create_user_admin: types.Operation,  # pylint: disable=unused-argument
+    get_token: Callable[[str, str], str],
+    create_cloud_integration_account: Callable,
+) -> None:
+    """The agent failed to clean up or crashed, so it never acks: the removed region stays and the version stays put."""
+    admin_token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
+    account_id = create_cloud_integration_account(admin_token, "aws", regions=["us-east-1", "us-west-2"])["id"]
+    provider_account_id = str(uuid.uuid4())
+
+    response = simulate_agent_checkin(signoz, admin_token, "aws", account_id, provider_account_id)
+    assert response.status_code == HTTPStatus.OK, response.text
+
+    response = requests.put(
+        signoz.self.host_configs["8080"].get(f"/api/v1/cloud_integrations/aws/accounts/{account_id}"),
+        headers={"Authorization": f"Bearer {admin_token}"},
+        json={"config": {"aws": {"regions": ["us-east-1"]}}},
+        timeout=10,
+    )
+    assert response.status_code == HTTPStatus.NO_CONTENT, response.text
+
+    for _ in range(3):
+        response = simulate_agent_checkin(signoz, admin_token, "aws", account_id, provider_account_id)
+        assert response.status_code == HTTPStatus.OK, response.text
+        assert response.json()["data"]["syncState"] == {
+            "version": 2,
+            "inSync": False,
+            "regions": {"us-east-1": {"state": "enabled"}, "us-west-2": {"state": "disabled"}},
+        }, "unacked removed region should stay without bumping the version"
+
+
+@pytest.mark.parametrize("synced_version", [2, 9], ids=["stale", "ahead"])
+def test_sync_state_ignores_mismatched_ack(
+    signoz: types.SigNoz,
+    create_user_admin: types.Operation,  # pylint: disable=unused-argument
+    get_token: Callable[[str, str], str],
+    create_cloud_integration_account: Callable,
+    synced_version: int,
+) -> None:
+    """An ack for any version other than the current one (v3) is ignored,
+    so us-west-2, removed at v2 and still unacked, is not dropped.
+    """
+    admin_token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
+    account_id = create_cloud_integration_account(admin_token, "aws", regions=["us-east-1", "us-west-2"])["id"]
+    provider_account_id = str(uuid.uuid4())
+
+    response = simulate_agent_checkin(signoz, admin_token, "aws", account_id, provider_account_id)
+    assert response.status_code == HTTPStatus.OK, response.text
+
+    for regions in (["us-east-1"], ["us-east-1", "eu-west-1"]):
+        response = requests.put(
+            signoz.self.host_configs["8080"].get(f"/api/v1/cloud_integrations/aws/accounts/{account_id}"),
+            headers={"Authorization": f"Bearer {admin_token}"},
+            json={"config": {"aws": {"regions": regions}}},
+            timeout=10,
+        )
+        assert response.status_code == HTTPStatus.NO_CONTENT, response.text
+
+        response = simulate_agent_checkin(signoz, admin_token, "aws", account_id, provider_account_id)
+        assert response.status_code == HTTPStatus.OK, response.text
+
+    expected_sync_state = {
+        "version": 3,
+        "inSync": False,
+        "regions": {"us-east-1": {"state": "enabled"}, "us-west-2": {"state": "disabled"}, "eu-west-1": {"state": "enabled"}},
+    }
+    assert response.json()["data"]["syncState"] == expected_sync_state
+
+    response = simulate_agent_checkin(signoz, admin_token, "aws", account_id, provider_account_id, synced_version=synced_version)
+    assert response.status_code == HTTPStatus.OK, response.text
+    assert response.json()["data"]["syncState"] == expected_sync_state, "an ack for another version should be ignored"
+
+
+def test_sync_state_applies_ack_before_config_change(
+    signoz: types.SigNoz,
+    create_user_admin: types.Operation,  # pylint: disable=unused-argument
+    get_token: Callable[[str, str], str],
+    create_cloud_integration_account: Callable,
+) -> None:
+    """The user changes regions while the agent syncs: the ack for the version it synced still lands."""
+    admin_token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
+    account_id = create_cloud_integration_account(admin_token, "aws", regions=["us-east-1", "us-west-2"])["id"]
+    provider_account_id = str(uuid.uuid4())
+
+    response = simulate_agent_checkin(signoz, admin_token, "aws", account_id, provider_account_id)
+    assert response.status_code == HTTPStatus.OK, response.text
+
+    for regions, synced_version in ((["us-east-1"], None), (["us-east-1", "eu-west-1"], 2)):
+        response = requests.put(
+            signoz.self.host_configs["8080"].get(f"/api/v1/cloud_integrations/aws/accounts/{account_id}"),
+            headers={"Authorization": f"Bearer {admin_token}"},
+            json={"config": {"aws": {"regions": regions}}},
+            timeout=10,
+        )
+        assert response.status_code == HTTPStatus.NO_CONTENT, response.text
+
+        response = simulate_agent_checkin(signoz, admin_token, "aws", account_id, provider_account_id, synced_version=synced_version)
+        assert response.status_code == HTTPStatus.OK, response.text
+
+    assert response.json()["data"]["syncState"] == {
+        "version": 3,
+        "inSync": False,
+        "regions": {"us-east-1": {"state": "enabled"}, "eu-west-1": {"state": "enabled"}},
+    }, "ack should drop the removed region before the new region bumps the version"
+
+
+@pytest.mark.parametrize("synced_version", [1, None], ids=["agent_acks_synced_version", "agent_crashed"])
+def test_sync_state_region_removed_during_sync(
+    signoz: types.SigNoz,
+    create_user_admin: types.Operation,  # pylint: disable=unused-argument
+    get_token: Callable[[str, str], str],
+    create_cloud_integration_account: Callable,
+    synced_version: int | None,
+) -> None:
+    """The user removes a region while the agent syncs v1; whether the agent acks v1 or crashed, the region must not be lost."""
+    admin_token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
+    account_id = create_cloud_integration_account(admin_token, "aws", regions=["us-east-1", "us-west-2"])["id"]
+    provider_account_id = str(uuid.uuid4())
+
+    response = simulate_agent_checkin(signoz, admin_token, "aws", account_id, provider_account_id)
+    assert response.status_code == HTTPStatus.OK, response.text
+    assert response.json()["data"]["syncState"] == {
+        "version": 1,
+        "inSync": True,
+        "regions": {"us-east-1": {"state": "enabled"}, "us-west-2": {"state": "enabled"}},
+    }
+
+    response = requests.put(
+        signoz.self.host_configs["8080"].get(f"/api/v1/cloud_integrations/aws/accounts/{account_id}"),
+        headers={"Authorization": f"Bearer {admin_token}"},
+        json={"config": {"aws": {"regions": ["us-east-1"]}}},
+        timeout=10,
+    )
+    assert response.status_code == HTTPStatus.NO_CONTENT, response.text
+
+    response = simulate_agent_checkin(signoz, admin_token, "aws", account_id, provider_account_id, synced_version=synced_version)
+    assert response.status_code == HTTPStatus.OK, response.text
+    assert response.json()["data"]["syncState"] == {
+        "version": 2,
+        "inSync": False,
+        "regions": {"us-east-1": {"state": "enabled"}, "us-west-2": {"state": "disabled"}},
+    }, "region removed mid-sync should be marked disabled"
+
+    response = simulate_agent_checkin(signoz, admin_token, "aws", account_id, provider_account_id, synced_version=2)
+    assert response.status_code == HTTPStatus.OK, response.text
+    assert response.json()["data"]["syncState"] == {
+        "version": 2,
+        "inSync": True,
+        "regions": {"us-east-1": {"state": "enabled"}},
+    }
+
+
+def test_sync_state_after_disconnect(
+    signoz: types.SigNoz,
+    create_user_admin: types.Operation,  # pylint: disable=unused-argument
+    get_token: Callable[[str, str], str],
+    create_cloud_integration_account: Callable,
+) -> None:
+    admin_token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
+    account_id = create_cloud_integration_account(admin_token, "aws", regions=["us-east-1", "us-west-2"])["id"]
+    provider_account_id = str(uuid.uuid4())
+
+    response = simulate_agent_checkin(signoz, admin_token, "aws", account_id, provider_account_id)
+    assert response.status_code == HTTPStatus.OK, response.text
+
+    response = requests.delete(
+        signoz.self.host_configs["8080"].get(f"/api/v1/cloud_integrations/aws/accounts/{account_id}"),
+        headers={"Authorization": f"Bearer {admin_token}"},
+        timeout=10,
+    )
+    assert response.status_code == HTTPStatus.NO_CONTENT, response.text
+
+    for _ in range(2):
+        response = simulate_agent_checkin(signoz, admin_token, "aws", account_id, provider_account_id)
+        assert response.status_code == HTTPStatus.OK, response.text
+        assert response.json()["data"]["removedAt"] is not None, "removedAt should be set after disconnect"
+        assert response.json()["data"]["syncState"] == {
+            "version": 2,
+            "inSync": False,
+            "regions": {"us-east-1": {"state": "disabled"}, "us-west-2": {"state": "disabled"}},
+        }, "every region should be disabled once, without bumping the version on later check-ins"
+
+    for _ in range(2):
+        response = simulate_agent_checkin(signoz, admin_token, "aws", account_id, provider_account_id, synced_version=2)
+        assert response.status_code == HTTPStatus.OK, response.text
+        assert response.json()["data"]["syncState"] == {"version": 2, "inSync": True, "regions": {}}, "acked removal should leave no regions"

tests/integration/tests/cloudintegrations/04_accounts.py

@@ -21,6 +21,11 @@ AWS_ACCOUNT_SPEC = ProviderAccountSpec(
     updated_params={"deployment_region": "us-east-1", "regions": ["us-east-1", "us-west-2", "eu-west-1"]},
     build_config=lambda p: {"aws": {"deploymentRegion": p["deployment_region"], "regions": p["regions"]}},
     expected_config=lambda p: {"regions": p["regions"]},
+    expected_sync_state=lambda p: {
+        "version": 1,
+        "inSync": True,
+        "regions": {region: {"state": "enabled"} for region in p["regions"]},
+    },
 )
 
 GCP_ACCOUNT_SPEC = ProviderAccountSpec(
@@ -128,6 +133,7 @@ def test_list_accounts_after_checkin(
     assert found["providerAccountId"] == provider_account_id, "providerAccountId should match"
     assert found["config"][spec.provider] == spec.expected_config(spec.initial_params), "config should match account config"
     assert found["agentReport"] is not None, "agentReport should be present after check-in"
+    assert found["agentReport"]["syncState"] == spec.expected_sync_state(spec.initial_params), "syncState should be seeded from the account regions on first check-in"
     assert found["removedAt"] is None, "removedAt should be null for a live account"
 
 
@@ -282,6 +288,7 @@ def test_update_account_after_checkin_preserves_connected_status(
     assert found_after is not None, "Account must still be listed after config update (account_id should not be reset)"
     assert found_after["providerAccountId"] == provider_account_id, "providerAccountId should be preserved after update"
     assert found_after["agentReport"] is not None, "agentReport should be preserved after update"
+    assert found_after["agentReport"]["syncState"] == found_before["agentReport"]["syncState"], "config update must not change syncState"
     assert found_after["config"][spec.provider] == spec.expected_config(spec.updated_params), "Config should reflect the update"
     assert found_after["removedAt"] is None, "removedAt should still be null"
 

tests/integration/tests/querierauthz/03_query_types.py

@@ -2,8 +2,6 @@ from collections.abc import Callable
 from datetime import UTC, datetime, timedelta
 from http import HTTPStatus
 
-import requests
-
 from fixtures import querier, types
 from fixtures.auth import USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD, change_user_role, create_active_user
 from fixtures.querier import make_query_request
@@ -152,27 +150,3 @@ def test_managed_viewer_meter_and_clickhouse_allowed_audit_denied(
     # audit-logs builder queries remain admin-only.
     audit = make_query_request(signoz, token, start, end, audit_query, request_type=querier.RequestType.RAW)
     assert audit.status_code == HTTPStatus.FORBIDDEN, audit.text
-
-
-def test_duplicate_signal_key_is_checked_on_the_bound_value(
-    signoz: types.SigNoz,
-    get_token: Callable[[str, str], str],
-) -> None:
-    now = datetime.now(tz=UTC)
-    start, end = int((now - timedelta(hours=1)).timestamp() * 1000), int(now.timestamp() * 1000)
-
-    # raw string: json= would collapse the duplicate "signal" key
-    body = (
-        f'{{"schemaVersion":"v1","start":{start},"end":{end},"requestType":"scalar",'
-        '"compositeQuery":{"queries":[{"type":"builder_query","spec":{"name":"A","signal":"traces","signal":"logs",'
-        '"disabled":false,"filter":{"expression":"signoz.workspace.key.id = \'key-a\'"},'
-        '"aggregations":[{"expression":"count()"}]}}]},"noCache":true}'
-    )
-
-    response = requests.post(
-        signoz.self.host_configs["8080"].get("/api/v5/query_range"),
-        timeout=querier.QUERY_TIMEOUT,
-        headers={"authorization": f"Bearer {get_token(key_a_email, user_password)}", "content-type": "application/json"},
-        data=body,
-    )
-    assert response.status_code == HTTPStatus.FORBIDDEN, response.text

tests/integration/tests/serviceaccount/06_fga.py

@@ -246,15 +246,6 @@ def test_attach_detach_dual_scoped(
     )
     assert resp.status_code == HTTPStatus.FORBIDDEN, f"assign viewer to target: expected 403, got {resp.status_code}: {resp.text}"
 
-    # duplicate roleId: the server binds the last one (viewer) -> forbidden. Raw string, json= would collapse the key.
-    resp = requests.post(
-        signoz.self.host_configs["8080"].get("/api/v1/service_account_roles"),
-        data=f'{{"serviceAccountId": "{target_id}", "roleId": "{editor_role_id}", "roleId": "{viewer_role_id}"}}',
-        headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"},
-        timeout=5,
-    )
-    assert resp.status_code == HTTPStatus.FORBIDDEN, f"assign duplicate roleId to target: expected 403, got {resp.status_code}: {resp.text}"
-
     # Both SA-detach (target id) and role-detach (editor) present -> remove allowed.
     resp = requests.delete(
         signoz.self.host_configs["8080"].get(f"/api/v1/service_account_roles/{editor_entry_id}"),