diff --git a/src/gallium/drivers/asahi/agx_pipe.c b/src/gallium/drivers/asahi/agx_pipe.c
index d34c0599d2b..3c37f8e60a8 100644
--- a/src/gallium/drivers/asahi/agx_pipe.c
+++ b/src/gallium/drivers/asahi/agx_pipe.c
@@ -37,6 +37,7 @@
 #include "frontend/sw_winsys.h"
 #include "gallium/auxiliary/util/u_transfer.h"
 #include "gallium/auxiliary/util/u_surface.h"
+#include "gallium/auxiliary/util/u_framebuffer.h"
 #include "agx_public.h"
 #include "agx_state.h"
 #include "magic.h"
@@ -552,10 +553,14 @@ agx_flush(struct pipe_context *pctx,
 }
 
 static void
-agx_destroy_context(struct pipe_context *ctx)
+agx_destroy_context(struct pipe_context *pctx)
 {
-   if (ctx->stream_uploader)
-      u_upload_destroy(ctx->stream_uploader);
+   struct agx_context *ctx = agx_context(pctx);
+
+   if (pctx->stream_uploader)
+      u_upload_destroy(pctx->stream_uploader);
+
+   util_unreference_framebuffer_state(&ctx->framebuffer);
 
    FREE(ctx);
 }
diff --git a/src/gallium/drivers/asahi/agx_state.c b/src/gallium/drivers/asahi/agx_state.c
index 28fc07c2753..d3bd2dbda0d 100644
--- a/src/gallium/drivers/asahi/agx_state.c
+++ b/src/gallium/drivers/asahi/agx_state.c
@@ -35,6 +35,7 @@
 #include "gallium/auxiliary/util/u_helpers.h"
 #include "gallium/auxiliary/util/u_viewport.h"
 #include "gallium/auxiliary/util/u_blend.h"
+#include "gallium/auxiliary/util/u_framebuffer.h"
 #include "gallium/auxiliary/tgsi/tgsi_from_mesa.h"
 #include "gallium/auxiliary/nir/tgsi_to_nir.h"
 #include "compiler/nir/nir.h"
@@ -176,6 +177,7 @@ agx_create_zsa_state(struct pipe_context *ctx,
    struct agx_zsa *so = CALLOC_STRUCT(agx_zsa);
    assert(!state->depth_bounds_test && "todo");
 
+   so->base = *state;
    so->disable_z_write = !state->depth_writemask;
 
    /* Z func can be used as-is */
@@ -302,7 +304,11 @@ agx_create_sampler_state(struct pipe_context *pctx,
       cfg.compare_func = agx_compare_funcs[state->compare_func];
    }
 
-   return bo;
+   struct agx_sampler_state *so = CALLOC_STRUCT(agx_sampler_state);
+   so->base = *state;
+   so->desc = bo;
+
+   return so;
 }
 
 static void
@@ -320,8 +326,10 @@ agx_bind_sampler_states(struct pipe_context *pctx,
 {
    struct agx_context *ctx = agx_context(pctx);
 
+   ctx->stage[shader].sampler_count = states ? count : 0;
+
    memcpy(&ctx->stage[shader].samplers[start], states,
-          sizeof(struct agx_bo *) * count);
+          sizeof(struct agx_sampler_state *) * count);
 }
 
 /* Channels agree for RGBA but are weird for force 0/1 */
@@ -498,6 +506,8 @@ agx_set_polygon_stipple(struct pipe_context *ctx,
 static void
 agx_set_sample_mask(struct pipe_context *pipe, unsigned sample_mask)
 {
+   struct agx_context *ctx = agx_context(pipe);
+   ctx->sample_mask = sample_mask;
 }
 
 static void
@@ -516,9 +526,11 @@ agx_set_scissor_states(struct pipe_context *pctx,
 }
 
 static void
-agx_set_stencil_ref(struct pipe_context *ctx,
+agx_set_stencil_ref(struct pipe_context *pctx,
                     const struct pipe_stencil_ref state)
 {
+   struct agx_context *ctx = agx_context(pctx);
+   ctx->stencil_ref = state;
 }
 
 static void
@@ -626,6 +638,7 @@ agx_set_framebuffer_state(struct pipe_context *pctx,
    /* XXX: eliminate this flush with batch tracking logic */
    pctx->flush(pctx, NULL, 0);
 
+   util_copy_framebuffer_state(&ctx->framebuffer, state);
    ctx->batch->width = state->width;
    ctx->batch->height = state->height;
    ctx->batch->nr_cbufs = state->nr_cbufs;
@@ -763,6 +776,7 @@ agx_create_shader_state(struct pipe_context *pctx,
                         const struct pipe_shader_state *cso)
 {
    struct agx_uncompiled_shader *so = CALLOC_STRUCT(agx_uncompiled_shader);
+   so->base = *cso;
 
    if (!so)
       return NULL;
@@ -994,11 +1008,12 @@ agx_build_pipeline(struct agx_context *ctx, struct agx_compiled_shader *cs, enum
    }
 
    for (unsigned i = 0; i < PIPE_MAX_SAMPLERS; ++i) {
-      struct agx_bo *bo = ctx->stage[stage].samplers[i];
+      struct agx_sampler_state *sampler = ctx->stage[stage].samplers[i];
 
-      if (!bo)
+      if (!sampler)
          continue;
 
+      struct agx_bo *bo = sampler->desc;
       agx_batch_add_bo(ctx->batch, bo);
 
       agx_pack(record, BIND_SAMPLER, cfg) {
diff --git a/src/gallium/drivers/asahi/agx_state.h b/src/gallium/drivers/asahi/agx_state.h
index 69005961c40..033585a057e 100644
--- a/src/gallium/drivers/asahi/agx_state.h
+++ b/src/gallium/drivers/asahi/agx_state.h
@@ -65,6 +65,7 @@ struct agx_compiled_shader {
 };
 
 struct agx_uncompiled_shader {
+   struct pipe_shader_state base;
    struct nir_shader *nir;
    struct hash_table *variants;
 
@@ -79,13 +80,11 @@ struct agx_stage {
    struct pipe_constant_buffer cb[PIPE_MAX_CONSTANT_BUFFERS];
    uint32_t cb_mask;
 
-   /* BOs for bound samplers. This is all the information we need at
-    * draw time to assemble the pipeline */
-   struct agx_bo *samplers[PIPE_MAX_SAMPLERS];
-
-   /* Sampler views need the full CSO due to Gallium state management */
+   /* Need full CSOs for u_blitter */
+   struct agx_sampler_state *samplers[PIPE_MAX_SAMPLERS];
    struct agx_sampler_view *textures[PIPE_MAX_SHADER_SAMPLER_VIEWS];
-   unsigned texture_count;
+
+   unsigned sampler_count, texture_count;
 };
 
 /* Uploaded scissor descriptors */
@@ -116,6 +115,7 @@ struct agx_batch {
 };
 
 struct agx_zsa {
+   struct pipe_depth_stencil_alpha_state base;
    enum agx_zs_func z_func;
    bool disable_z_write;
 };
@@ -160,7 +160,16 @@ struct agx_context {
    struct pipe_blend_color blend_color;
    struct pipe_viewport_state viewport;
    struct pipe_scissor_state scissor;
+   struct pipe_stencil_ref stencil_ref;
    struct agx_streamout streamout;
+   uint16_t sample_mask;
+   struct pipe_framebuffer_state framebuffer;
+
+   struct pipe_query *cond_query;
+   bool cond_cond;
+   enum pipe_render_cond_flag cond_mode;
+
+   bool is_noop;
 
    uint8_t render_target[8][AGX_RENDER_TARGET_LENGTH];
 };
@@ -181,6 +190,13 @@ struct agx_query {
    unsigned	query;
 };
 
+struct agx_sampler_state {
+   struct pipe_sampler_state base;
+
+   /* Prepared descriptor */
+   struct agx_bo *desc;
+};
+
 struct agx_sampler_view {
    struct pipe_sampler_view base;